diff --git a/src/beaver_gateway/backends/transcript.py b/src/beaver_gateway/backends/transcript.py index f955692..aebca1e 100644 --- a/src/beaver_gateway/backends/transcript.py +++ b/src/beaver_gateway/backends/transcript.py @@ -20,13 +20,14 @@ except ImportError: # pragma: no cover _cli_version = "2.1.248" if TYPE_CHECKING: - from collections.abc import Iterable + from collections.abc import Iterable, Sequence __all__ = [ "CLI_VERSION", "build_entries", "close_open_tool_uses", "fingerprint", + "last_message_before", "messages_from_entries", "open_tool_uses", "prompt_count", @@ -35,6 +36,7 @@ __all__ = [ "spoken_text", "strip_tool_entries", "text_of", + "turn_starts", "window_entries", ] @@ -438,6 +440,33 @@ def prompt_count(entries: Iterable[Mapping[str, Any]]) -> int: return sum(1 for e in entries if _is_prompt_entry(e)) +def turn_starts(entries: Iterable[Mapping[str, Any]]) -> list[int]: + """Where each turn opens: the index of every prompt the SDK was handed. + + The CLI stamps only those with ``promptSource``; what it writes as a + user on its own (an interruption note, a skill body) stays inside the + turn it belongs to. + """ + return [ + index + for index, entry in enumerate(entries) + if entry.get("promptSource") and _is_prompt_entry(entry) + ] + + +def last_message_before(entries: Sequence[Mapping[str, Any]], index: int) -> str | None: + """Uuid of the last ``user``/``assistant`` entry ahead of ``index``.""" + for entry in reversed(entries[:index]): + uuid = entry.get("uuid") + if ( + entry.get("type") in _PROMPT_TYPES + and isinstance(uuid, str) + and not entry.get("isSidechain") + ): + return uuid + return None + + def spoken_text(block: Any, *, thinking_is_progress: bool) -> str: """What a content block (dict or model) says to the human. diff --git a/src/beaver_gateway/conversations/questions.py b/src/beaver_gateway/conversations/questions.py index cf60e81..51547ef 100644 --- a/src/beaver_gateway/conversations/questions.py +++ b/src/beaver_gateway/conversations/questions.py @@ -7,7 +7,7 @@ import contextlib from typing import TYPE_CHECKING, Any, cast from uuid import uuid4 -from beaver_gateway.conversations.spawning import Spawning +from beaver_gateway.conversations.rewinding import Rewinding from beaver_gateway.conversations.state import Question if TYPE_CHECKING: @@ -16,7 +16,7 @@ if TYPE_CHECKING: __all__ = ["Questions"] -class Questions(Spawning): +class Questions(Rewinding): async def ask(self, key: str, payload: dict[str, Any]) -> str | None: """Show the question and wait; ``None`` when nobody answered in time.""" conv = await self.get(key) diff --git a/src/beaver_gateway/conversations/rewinding.py b/src/beaver_gateway/conversations/rewinding.py new file mode 100644 index 0000000..6f826cd --- /dev/null +++ b/src/beaver_gateway/conversations/rewinding.py @@ -0,0 +1,114 @@ +"""Taking turns back: the thread moves to a fork that ends before them.""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, cast + +from claude_agent_sdk import fork_session_via_store + +from beaver_gateway.backends.transcript import ( + last_message_before, + messages_from_entries, + reply_text, + text_of, + turn_starts, +) +from beaver_gateway.conversations.spawning import Spawning + +if TYPE_CHECKING: + from collections.abc import Sequence + + from beaver_gateway.storage.models import Conversation + +__all__ = ["Rewinding", "Rewound"] + +_log = logging.getLogger(__name__) + +PREVIEW = 280 + + +@dataclass(frozen=True, slots=True) +class Rewound: + conversation: Conversation + turns: int + previous_session: str + + +class Rewinding(Spawning): + async def turns( + self, conv: Conversation, *, limit: int = 10 + ) -> list[dict[str, Any]]: + """The latest turns that can be taken back, oldest first; the first stays.""" + entries = await self.entries(conv) + starts = turn_starts(entries) + ends = [*starts[1:], len(entries)] + spans = list(zip(starts, ends, strict=True))[1:] + agent = self._agents.get(conv.agent_name) + progress = agent is not None and agent.thinking_is_progress + return [ + self._turn(entries[start:end], thinking_is_progress=progress) + for start, end in spans[-limit:] + ] + + def _turn( + self, entries: Sequence[Any], *, thinking_is_progress: bool + ) -> dict[str, Any]: + prompt, *rest = messages_from_entries(entries) + replies = ( + reply_text(m["content"], thinking_is_progress=thinking_is_progress) + for m in reversed(rest) + if m["role"] == "assistant" + ) + return { + "id": entries[0].get("uuid"), + "ts": entries[0].get("timestamp"), + "text": text_of(self._unwrapped(prompt)["content"])[:PREVIEW], + "reply": next((text for text in replies if text), "")[:PREVIEW], + "failed": any(e.get("isApiErrorMessage") for e in entries), + } + + async def rewind(self, conv: Conversation, turn: str) -> Rewound: + """Take ``turn`` and every turn after it out of what the model sees. + + Nothing is deleted: the conversation moves to a fork of its session + that ends on the message before ``turn``, the cut session stays in + the store under its old id. + """ + row_id = cast("int", conv.id) + runner = self._runner(row_id) + if runner.lock.locked(): + msg = f"conversation {conv.external_id} is mid-turn; interrupt it first" + raise RuntimeError(msg) + async with runner.lock: + row = await self.get_row(row_id) or conv + entries = await self.entries(row) + starts = turn_starts(entries) + cut = next((i for i in starts[1:] if entries[i].get("uuid") == turn), None) + keep = last_message_before(entries, cut) if cut is not None else None + if cut is None or keep is None: + msg = f"conversation {row.external_id} has no turn {turn} to take back" + raise LookupError(msg) + previous = cast("str", row.session_id) + forked = await fork_session_via_store( + self._store, + previous, + directory=str(self._claude_agent(row.agent_name).cwd), + up_to_message_id=keep, + ) + await self._backend(row.agent_name).close(row.external_id) + + async def apply(moved: Conversation) -> None: + moved.session_id = forked.session_id + + row = await self._update(row, apply) + turns = sum(1 for start in starts if start >= cut) + _log.info( + "rewound %s by %d turns: session %s -> %s", + row.external_id, + turns, + previous, + forked.session_id, + ) + return Rewound(conversation=row, turns=turns, previous_session=previous) diff --git a/src/beaver_gateway/conversations/service.py b/src/beaver_gateway/conversations/service.py index e6c400b..6639cca 100644 --- a/src/beaver_gateway/conversations/service.py +++ b/src/beaver_gateway/conversations/service.py @@ -1,7 +1,7 @@ """``Conversations`` - the service every frontend, job and gateway tool talks to. Built as layers, one file each: rows → seeds → turns → messaging → -spawning → questions → closing; this file adds start, stop and restart +spawning → rewinding → questions → closing; this file adds start, stop and restart recovery. A turn started by a user message streams back to whoever asked; a turn started by an inject streams nowhere and can only speak via ``say``. """ diff --git a/src/beaver_gateway/conversations/turns.py b/src/beaver_gateway/conversations/turns.py index 89dc55b..5cc2b16 100644 --- a/src/beaver_gateway/conversations/turns.py +++ b/src/beaver_gateway/conversations/turns.py @@ -68,14 +68,14 @@ class Turns(Seeds): backend = self._backend(conv.agent_name) turn_id = turn_id or f"turn_{uuid4().hex[:12]}" capture = capture or TurnCapture() - resume = session_id if session_id is not None else conv.session_id async with runner.lock: runner.turn_id = turn_id runner.origin = origin runner.text = shown if shown is not None else _prompt_preview(messages) runner.started_at = datetime.now(UTC) runner.tools = {} - await self._mark_running(conv, turn_id) + row = await self._mark_running(conv, turn_id) + resume = session_id if session_id is not None else row.session_id before = await self.before_turn(conv) self._bus.publish( "turn.start", @@ -463,12 +463,12 @@ class Turns(Seeds): except Exception: # noqa: BLE001 _log.exception("rate limit write failed") - async def _mark_running(self, conv: Conversation, turn_id: str) -> None: + async def _mark_running(self, conv: Conversation, turn_id: str) -> Conversation: async def apply(row: Conversation) -> None: row.running_turn = turn_id row.last_activity_at = datetime.now(UTC) - await self._update(conv, apply) + return await self._update(conv, apply) async def _mark_done( self, conv: Conversation, capture: TurnCapture, *, cut: bool = False diff --git a/src/beaver_gateway/frontends/api/frontend.py b/src/beaver_gateway/frontends/api/frontend.py index 0e4d655..1db1d4d 100644 --- a/src/beaver_gateway/frontends/api/frontend.py +++ b/src/beaver_gateway/frontends/api/frontend.py @@ -79,6 +79,8 @@ MEMORY_MAX_ENTRIES = 5000 SEARCH_MIN = 2 GRAPH_LIMIT = 120 SEARCH_LIMIT = 20 +TURNS_LIMIT = 10 +TURNS_MAX = 50 SNIPPET_MAX = 160 @@ -578,6 +580,42 @@ def build_app( # noqa: PLR0915 "text": result.text, } + @app.get("/conversations/{public_id}/turns") + async def get_turns(public_id: str, request: Request) -> dict[str, Any]: + await require_token(request, runtime, scope=SCOPE) + conv = await conv_of(public_id) + limit = min(query_int(request, "limit", TURNS_LIMIT), TURNS_MAX) + return { + "id": conv.external_id, + "turns": await conversations.turns(conv, limit=limit), + } + + @app.post("/conversations/{public_id}/rewind") + async def post_rewind(public_id: str, request: Request) -> dict[str, Any]: + token = await require_token(request, runtime, scope=ADMIN_SCOPE) + conv = await conv_of(public_id) + turn = text_of(await body_of(request), "turn") + try: + result = await conversations.rewind(conv, turn) + except LookupError as exc: + raise HTTPException(status.HTTP_404_NOT_FOUND, str(exc)) from exc + except RuntimeError as exc: + raise HTTPException(status.HTTP_409_CONFLICT, str(exc)) from exc + out = { + "turns": result.turns, + "session_id": result.conversation.session_id, + "previous_session": result.previous_session, + } + await audit.log( + runtime, + actor=f"token:{token}", + kind="api_rewind", + agent_name=conv.agent_name, + conversation=conv.external_id, + **out, + ) + return {"id": conv.external_id, **out} + @app.post("/conversations/{public_id}/bind") async def post_bind(public_id: str, request: Request) -> dict[str, Any]: await require_token(request, runtime, scope=SCOPE) diff --git a/tests/test_api.py b/tests/test_api.py index 5d05112..bcfee45 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -16,7 +16,7 @@ from claude_agent_sdk import ( UserMessage, ) from httpx import ASGITransport, AsyncClient -from test_conversations import ScriptedClient, World +from test_conversations import ScriptedClient, World, _three_turns from beaver_gateway.frontends.markdown.history import rewrite_messages from beaver_gateway.security.auth import TokenStore @@ -444,6 +444,52 @@ async def test_sessions_show_their_processes_and_take_orders(world: World) -> No assert kinds == ["session_interrupt", "session_close"] +async def test_rewind_lists_turns_and_needs_admin_scope(world: World) -> None: + api = Api(world) + api.store.grant("api-only", "weak", scope="api") + weak = {"Authorization": "Bearer weak"} + sid, entries = await _three_turns(world) + conv = await world.conversations.create( + kind="master", agent="a", origin="test", session_id=sid + ) + path = f"/conversations/{conv.external_id}" + + listed = (await api.http.get(f"{path}/turns", headers=weak)).json() + assert [t["text"] for t in listed["turns"]] == ["two", "three"] + assert (await api.get(f"{path}/turns", {"limit": 1}))["turns"] == [ + listed["turns"][1] + ] + body = {"turn": listed["turns"][0]["id"]} + denied = await api.http.post(f"{path}/rewind", json=body, headers=weak) + assert denied.status_code == 403 + missing = await api.http.post( + f"{path}/rewind", json={"turn": entries[0]["uuid"]}, headers=HEADERS + ) + assert missing.status_code == 404 + + done = await api.http.post(f"{path}/rewind", json=body, headers=HEADERS) + assert done.status_code == 200, done.text + out = done.json() + assert out["turns"] == 2 and out["previous_session"] == sid + assert (await api.get(path))["session_id"] == out["session_id"] != sid + assert (await api.get(f"{path}/turns"))["turns"] == [] + messages = (await api.get(f"{path}/history"))["messages"] + assert [m["content"] for m in messages] == [ + "one", + [{"type": "text", "text": "done one"}], + ] + (record,) = (await api.get("/audit"))["records"] + assert record["kind"] == "api_rewind" + assert record["detail"]["previous_session"] == sid + + ScriptedClient.hold = asyncio.Event() + await world.conversations.post(conv, "busy") + await asyncio.sleep(0.2) + busy = await api.http.post(f"{path}/rewind", json=body, headers=HEADERS) + assert busy.status_code == 409 + ScriptedClient.hold.set() + + async def test_tokens_and_audit_need_admin_scope(world: World) -> None: api = Api(world) api.store.grant("api-only", "weak", scope="api") diff --git a/tests/test_conversations.py b/tests/test_conversations.py index 30e4bca..c54b9ab 100644 --- a/tests/test_conversations.py +++ b/tests/test_conversations.py @@ -504,6 +504,105 @@ async def test_copy_seed_forks_parent_with_window(world: World) -> None: assert (await world.conversations.get(child.external_id)).flags["seed"] is None +async def _three_turns(world: World) -> tuple[str, list[dict[str, Any]]]: + sid = str(uuid.uuid4()) + history = [ + {"role": "user", "content": "one"}, + {"role": "assistant", "content": "done one"}, + {"role": "user", "content": "two"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "t1", "name": "Read", "input": {}}], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "t1", "content": "x"}], + }, + {"role": "assistant", "content": "done two"}, + {"role": "user", "content": "three"}, + {"role": "assistant", "content": "refused"}, + ] + entries = build_entries(history, session_id=sid, cwd=str(world.root), model="m") + entries[-1]["isApiErrorMessage"] = True + note = { + **entries[2], + "uuid": str(uuid.uuid4()), + "message": { + "role": "user", + "content": [{"type": "text", "text": "[Request interrupted by user]"}], + }, + } + del note["promptSource"] + entries.insert(6, note) + await world.store.append(world.key(sid), entries) + return sid, entries + + +async def test_rewind_moves_the_thread_to_a_fork_that_ends_before_the_turn( + world: World, +) -> None: + sid, entries = await _three_turns(world) + conv = await world.conversations.create( + kind="master", agent="a", origin="test", session_id=sid + ) + await world.conversations.post(conv, "live") + await world.settle(conv, 1) + assert world.pool.get(conv.external_id) is not None + + turns = await world.conversations.turns(conv) + assert [(t["text"], t["reply"], t["failed"]) for t in turns] == [ + ("two", "done two", False), + ("three", "refused", True), + ] + assert [t["id"] for t in turns] == [entries[2]["uuid"], entries[7]["uuid"]] + + result = await world.conversations.rewind(conv, turns[1]["id"]) + + assert result.turns == 1 and result.previous_session == sid + moved = result.conversation.session_id + assert moved != sid + assert (await world.conversations.get(conv.external_id)).session_id == moved + assert await world.store.load(world.key(sid)) == entries + kept = [ + e["message"]["content"] + for e in await world.store.load(world.key(moved)) + if e["type"] in ("user", "assistant") + ] + assert len(kept) == 7 and kept[0] == "one" + assert kept[-1] == [{"type": "text", "text": "[Request interrupted by user]"}] + assert world.pool.get(conv.external_id) is None + + await world.conversations.run_text_turn(conv, "after", origin="user") + assert conv.session_id == sid + assert ScriptedClient.instances[-1].options.resume == moved + + (two,) = await world.conversations.turns(result.conversation) + assert two["text"] == "two" + again = await world.conversations.rewind(conv, two["id"]) + assert again.turns == 1 and again.previous_session == moved + assert await world.conversations.turns(again.conversation) == [] + first = (await world.store.load(world.key(again.conversation.session_id)))[0] + with pytest.raises(LookupError): + await world.conversations.rewind(conv, first["uuid"]) + with pytest.raises(LookupError): + await world.conversations.rewind(conv, str(uuid.uuid4())) + + +async def test_rewind_waits_for_no_turn(world: World) -> None: + sid, entries = await _three_turns(world) + conv = await world.conversations.create( + kind="master", agent="a", origin="test", session_id=sid + ) + ScriptedClient.hold = asyncio.Event() + await world.conversations.post(conv, "busy") + await asyncio.sleep(0.2) + with pytest.raises(RuntimeError, match="mid-turn"): + await world.conversations.rewind(conv, entries[7]["uuid"]) + ScriptedClient.hold.set() + await world.settle(conv, 1) + assert (await world.conversations.get(conv.external_id)).session_id == sid + + async def test_merge_injects_summary_into_parent(world: World) -> None: sid = str(uuid.uuid4()) await world.store.append( diff --git a/ui/src/lib/api/client.ts b/ui/src/lib/api/client.ts index 2c83977..79979b0 100644 --- a/ui/src/lib/api/client.ts +++ b/ui/src/lib/api/client.ts @@ -16,6 +16,7 @@ import type { SearchResponse, SessionsResponse, TokenRow, + TurnSummary, UsageGroup, UsageResponse, VaultGraph, @@ -187,6 +188,14 @@ export class ApiClient { return this.get(`/api/conversations/${id}/history`); } + turns(id: string, limit: number): Promise<{ turns: TurnSummary[] }> { + return this.get(`/api/conversations/${id}/turns`, { limit }); + } + + rewind(id: string, turn: string): Promise<{ turns: number }> { + return this.post(`/api/conversations/${id}/rewind`, { turn }); + } + context(id: string, prompt = false): Promise { return this.get(`/api/conversations/${id}/context`, { prompt: prompt ? 1 : undefined, diff --git a/ui/src/lib/api/types.ts b/ui/src/lib/api/types.ts index 0b6a99c..af42548 100644 --- a/ui/src/lib/api/types.ts +++ b/ui/src/lib/api/types.ts @@ -113,6 +113,14 @@ export interface HistoryMessage { ts?: string | null; } +export interface TurnSummary { + failed: boolean; + id: string; + reply: string; + text: string; + ts: string | null; +} + export interface EntriesPage { entries: Record[]; id: string; diff --git a/ui/src/lib/format.ts b/ui/src/lib/format.ts index 9791572..91c9892 100644 --- a/ui/src/lib/format.ts +++ b/ui/src/lib/format.ts @@ -109,6 +109,13 @@ export function fmtTime(iso: string | null | undefined): string { : "–"; } +export function fmtMinute(iso: string | null | undefined): string { + const date = parseDate(iso); + return date + ? date.toLocaleTimeString(undefined, { hour: "2-digit", minute: "2-digit" }) + : "–"; +} + export function fmtDateTime(iso: string | null | undefined): string { const date = parseDate(iso); return date diff --git a/ui/src/lib/panel/actions.svelte.ts b/ui/src/lib/panel/actions.svelte.ts index cb1a971..f715daf 100644 --- a/ui/src/lib/panel/actions.svelte.ts +++ b/ui/src/lib/panel/actions.svelte.ts @@ -9,6 +9,7 @@ import FileTextIcon from "@lucide/svelte/icons/file-text"; import GitBranchIcon from "@lucide/svelte/icons/git-branch"; import GitMergeIcon from "@lucide/svelte/icons/git-merge"; import PencilIcon from "@lucide/svelte/icons/pencil"; +import RewindIcon from "@lucide/svelte/icons/rewind"; import RotateCcwIcon from "@lucide/svelte/icons/rotate-ccw"; import SendIcon from "@lucide/svelte/icons/send"; import type { Component } from "svelte"; @@ -38,6 +39,7 @@ export interface ActionScope { onChanged: () => void; onOpen: (id: string) => void; onRename: () => void; + onRewind: () => void; } const TELEGRAM = "telegram"; @@ -103,6 +105,12 @@ export function conversationActions( label: "Branch…", run: scope.onBranch, }); + actions.push({ + disabled: !info.session_id, + icon: RewindIcon, + label: "Rewind…", + run: scope.onRewind, + }); actions.push({ checked: remembers, gap: true, diff --git a/ui/src/lib/panel/conversation-menu.svelte b/ui/src/lib/panel/conversation-menu.svelte index c8c343b..543ac25 100644 --- a/ui/src/lib/panel/conversation-menu.svelte +++ b/ui/src/lib/panel/conversation-menu.svelte @@ -13,6 +13,7 @@ import { usePanelHost } from "./host"; import MenuItems, { type MenuKit } from "./menu-items.svelte"; import RenameDialog from "./rename-dialog.svelte"; + import RewindDialog from "./rewind-dialog.svelte"; let { client, @@ -47,6 +48,7 @@ let failure = $state(null); let branchOpen = $state(false); let renameOpen = $state(false); + let rewindOpen = $state(false); const loaded = $derived(info ?? fetched); @@ -77,6 +79,9 @@ onRename: () => { renameOpen = true; }, + onRewind: () => { + rewindOpen = true; + }, }); const actions = $derived(loaded ? conversationActions(loaded, scope) : []); @@ -141,4 +146,5 @@ bind:open={branchOpen} /> + {/if} diff --git a/ui/src/lib/panel/rewind-dialog.svelte b/ui/src/lib/panel/rewind-dialog.svelte new file mode 100644 index 0000000..85e311c --- /dev/null +++ b/ui/src/lib/panel/rewind-dialog.svelte @@ -0,0 +1,210 @@ + + + + + + Rewind {name} + + Pick the turn to cut from. The model forgets it and every turn after it. + The old session stays in the database. + + +
+ {#if failure} + + {:else if turns === null} +
+ {#each [0, 1, 2] as row (row)} +
+ + +
+ {/each} +
+ {:else if turns.length === 0} + + {:else} +
+ {#each turns as turn, index (turn.id)} + {@const gone = picked !== null && index >= picked} + + {/each} +
+ {/if} + {#if info.running_turn} +

A turn is running. Interrupt it first.

+ {/if} +
+ + + {#if turns?.length} + + {/if} + +
+