From 8019c76e53a325803d98745792521ffbb8006b4d Mon Sep 17 00:00:00 2001 From: h Date: Fri, 4 Sep 2026 20:23:33 +0200 Subject: [PATCH] fix(backends,conversations,telegram,ui): a background subagent reports back as its own turn --- src/beaver_gateway/backends/claude_sdk.py | 286 ++++++++++++++---- src/beaver_gateway/backends/sessions.py | 27 +- src/beaver_gateway/conversations/state.py | 6 + src/beaver_gateway/conversations/texts.py | 2 + src/beaver_gateway/conversations/turns.py | 101 ++++++- .../frontends/telegram/drafts.py | 16 +- .../frontends/telegram/frontend.py | 2 +- tests/test_claude_sdk_backend.py | 12 + tests/test_conversations.py | 100 ++++++ tests/test_telegram.py | 19 +- ui/src/lib/panel/activity-feed.svelte | 2 +- ui/src/lib/panel/activity.svelte.ts | 87 +++++- ui/src/lib/panel/chat-view.svelte | 10 +- ui/src/lib/panel/tool-node.svelte | 2 +- ui/src/lib/panel/turn-card.svelte | 27 +- 15 files changed, 608 insertions(+), 91 deletions(-) diff --git a/src/beaver_gateway/backends/claude_sdk.py b/src/beaver_gateway/backends/claude_sdk.py index e8052b6..a4d0efd 100644 --- a/src/beaver_gateway/backends/claude_sdk.py +++ b/src/beaver_gateway/backends/claude_sdk.py @@ -37,6 +37,7 @@ from claude_agent_sdk import ( PermissionResultDeny, ResultMessage, StreamEvent, + TaskNotificationMessage, TextBlock, ThinkingBlock, ToolResultBlock, @@ -72,7 +73,14 @@ from beaver_gateway.events.stream import ( ) if TYPE_CHECKING: - from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Sequence + from collections.abc import ( + AsyncIterator, + Awaitable, + Callable, + Coroutine, + Iterable, + Sequence, + ) from anthropic.types import MessageParam from claude_agent_sdk import ( @@ -201,8 +209,16 @@ class ClaudeSdkBackend: tool_server: Callable[[str, str], McpSdkServerConfig | None] | None = None, asker: Callable[[str, dict[str, Any]], Awaitable[str]] | None = None, audit_sink: Callable[[policy_mod.ToolAudit], Awaitable[None]] | None = None, + waker: Callable[[str, str], Coroutine[Any, Any, None]] | None = None, ) -> None: + """One backend per Claude agent. + + ``waker(key, note)`` is called when a session starts a turn by itself + (a background task reported back); the caller then attaches. + """ self._agent = agent + self.waker = waker + self._tasks: set[asyncio.Task[None]] = set() self._audit_sink = audit_sink self._store = session_store self._runner = runner or RunnerConfig() @@ -256,6 +272,51 @@ class ClaudeSdkBackend: def live(self, key: str) -> Session | None: return self._pool.get(key) + def has_pending(self, key: str) -> bool: + """A turn the session started on its own is waiting for ``attach``.""" + live = self._pool.get(key) + return live is not None and bool(live.pending) + + def pending_note(self, key: str) -> str | None: + live = self._pool.get(key) + return live.note if live is not None else None + + async def attach( + self, + *, + key: str, + capture: TurnCapture | None = None, + observer: Callable[[Any], None] | None = None, + turn_id: str | None = None, + ) -> AsyncIterator[MessageStreamEvent]: + """Consume the turn a live session started by itself as a normal turn. + + The CLI runs one when a background task (a subagent) reports back: + no prompt goes out, the events are what the model said in reply to + the report, the turn ends on that result. + """ + live = self._pool.get(key) + if live is None or not live.pending: + return + message_id = f"msg_{uuid.uuid4().hex}" + yield build_message_start(message_id=message_id, model=self._agent.model) + turn = _Turn() + async with live.lock: + live.running_turn = turn_id or message_id + live.last_used = time.monotonic() + async for event in self._run_turn(live, None, turn, observer, capture): + yield event + live.turns += 1 + live.last_used = time.monotonic() + live.running_turn = None + usage, _ = await self._after_turn( + live, turn, conversation_id=key, history=[], capture=capture + ) + yield build_message_delta( + stop_reason=turn.stop_reason, usage=_wire_usage(usage) + ) + yield build_message_stop() + async def repair_session( self, session_id: str, *, text: str = "interrupted" ) -> int: @@ -404,70 +465,71 @@ class ClaudeSdkBackend: async def _run_turn( self, live: Session, - prompt: str | list[dict[str, Any]], + prompt: str | list[dict[str, Any]] | None, turn: _Turn, observer: Callable[[Any], None] | None = None, capture: TurnCapture | None = None, ) -> AsyncIterator[MessageStreamEvent]: - """Run one prompt, yielding wire events as they arrive. + """Run one prompt, or attach to a self-started one, yielding events. ``turn`` is filled in place (result, synthesized history, count of events already yielded) so the caller can finish bookkeeping - and decide whether a retry is still possible - after a failure. The session id lands in ``capture`` with the first frame, so a turn cut by a restart still leaves a resumable session behind. + + A result whose ``origin`` is not the human's ends a turn the CLI ran + on its own; inside a prompted turn it is passed over, the events it + came with stay in this turn rather than getting lost. """ - streaming = self._agent.options.include_partial_messages raw: list[Any] = [] - next_index = 0 - offset = 0 - await live.client.query(prompt if isinstance(prompt, str) else _stream(prompt)) - async for message in live.client.receive_response(): - if observer is not None: - observer(message) - session_id = getattr(message, "session_id", None) - if isinstance(session_id, str) and session_id and live.session_id is None: - live.session_id = session_id - if capture is not None: - capture.session_id = session_id - if isinstance(message, MirrorErrorMessage): - live.dirty = True - _log.error( - "session %s: mirror error, marked dirty: %s", - live.session_id, - message.error, + cursor = _Cursor() + live.observer = observer + live.attached = True + for queued in live.pending: + live.inbox.put_nowait(queued) + live.pending.clear() + try: + if prompt is not None: + await live.client.query( + prompt if isinstance(prompt, str) else _stream(prompt) ) - continue - if getattr(message, "parent_tool_use_id", None) is not None: - continue - if isinstance(message, StreamEvent): - event = message.event - if event.get("type") == "message_start": - offset = next_index - turn.context_tokens = _context_of(event.get("message")) + while True: + message = await live.inbox.get() + if message is _END: + if prompt is not None: + msg = f"claude exited before answering ({live.session_id})" + raise RuntimeError(msg) + break + if isinstance(message, BaseException): + raise message + if observer is not None: + observer(message) + if not self._note(live, message, capture): continue - index = event.get("index") - if isinstance(index, int): - next_index = max(next_index, offset + index + 1) - if streaming: - for out in _emit_stream_event(event, offset + index): - turn.events += 1 - yield out - elif isinstance(message, AssistantMessage): - raw.append(message) - if not streaming: - for block in message.content: - for out in _emit_block(block, next_index): - turn.events += 1 - yield out - next_index += 1 - elif isinstance(message, UserMessage): - raw.append(message) - elif isinstance(message, ResultMessage): - turn.result = message - turn.stop_reason = _STOP_REASONS.get( - message.stop_reason or "", "end_turn" - ) + if isinstance(message, ResultMessage): + if prompt is not None and _self_started(message): + _log.info( + "session %s: a self-started turn ended inside a " + "prompted one, folded in", + live.session_id, + ) + continue + turn.result = message + turn.stop_reason = _STOP_REASONS.get( + message.stop_reason or "", "end_turn" + ) + break + for out in self._emit(message, turn, cursor, raw): + turn.events += 1 + yield out + if turn.events == 0: + for out in _emit_whole(raw, cursor): + turn.events += 1 + yield out + finally: + live.attached = False + live.woken = False turn.synthesized = synthesize_turn_messages(raw) _log.info( "turn: agent=%s session=%s events=%d synthesized=%d stop=%s", @@ -478,6 +540,88 @@ class ClaudeSdkBackend: turn.stop_reason, ) + @staticmethod + def _note(live: Session, message: Any, capture: TurnCapture | None) -> bool: + """Session bookkeeping for one message; False if not the main thread's.""" + session_id = getattr(message, "session_id", None) + if isinstance(session_id, str) and session_id and live.session_id is None: + live.session_id = session_id + if capture is not None: + capture.session_id = session_id + if isinstance(message, MirrorErrorMessage): + live.dirty = True + _log.error( + "session %s: mirror error, marked dirty: %s", + live.session_id, + message.error, + ) + return False + if isinstance(message, TaskNotificationMessage): + live.note = _note_of(message) + return False + return getattr(message, "parent_tool_use_id", None) is None + + def _emit( + self, message: Any, turn: _Turn, cursor: _Cursor, raw: list[Any] + ) -> Iterable[MessageStreamEvent]: + streaming = self._agent.options.include_partial_messages + if isinstance(message, StreamEvent): + event = message.event + if event.get("type") == "message_start": + cursor.offset = cursor.next_index + turn.context_tokens = _context_of(event.get("message")) + return + index = event.get("index") + if isinstance(index, int): + cursor.next_index = max(cursor.next_index, cursor.offset + index + 1) + if streaming: + yield from _emit_stream_event(event, cursor.offset + index) + elif isinstance(message, AssistantMessage): + raw.append(message) + if not streaming: + for block in message.content: + yield from _emit_block(block, cursor.next_index) + cursor.next_index += 1 + elif isinstance(message, UserMessage): + raw.append(message) + + async def _read(self, live: Session) -> None: + """Read the client for the session's whole life. + + During a turn every message goes to the turn's inbox. Between turns + a subagent still running is shown to the last observer, and the + first frame of the main thread means the CLI started a turn by + itself: it is buffered in ``pending`` and the waker is told once. + """ + try: + async for message in live.client.receive_messages(): + if live.attached: + live.inbox.put_nowait(message) + continue + if isinstance(message, TaskNotificationMessage): + live.note = _note_of(message) + if _main_thread(message): + live.pending.append(message) + if not live.woken: + live.woken = True + self._wake(live) + elif live.observer is not None: + live.observer(message) + live.inbox.put_nowait(_END) + except asyncio.CancelledError: + raise + except Exception as exc: # noqa: BLE001 + _log.error("session %s: reader died: %s", live.session_id, exc) + live.dirty = True + live.inbox.put_nowait(exc) + + def _wake(self, live: Session) -> None: + if self.waker is None: + return + task = asyncio.create_task(self.waker(live.key, live.note or "")) + self._tasks.add(task) + task.add_done_callback(self._tasks.discard) + async def _acquire( self, key: str, @@ -541,7 +685,7 @@ class ClaudeSdkBackend: spec.tools, self._runner.user, ) - return Session( + live = Session( key=key, agent=self._agent.name, kind=spec.kind, @@ -550,6 +694,8 @@ class ClaudeSdkBackend: resumed=resume is not None, pinned=spec.pinned, ) + live.reader = asyncio.create_task(self._read(live)) + return live def _default_factory(self, options: ClaudeAgentOptions) -> SessionClient: return _RunnerClient(options, uid=self._uid) @@ -739,6 +885,15 @@ class _SessionSpec: tools: bool +_END = object() + + +@dataclass +class _Cursor: + next_index: int = 0 + offset: int = 0 + + @dataclass class _Turn: events: int = 0 @@ -866,6 +1021,33 @@ async def _stream(content: list[dict[str, Any]]) -> AsyncIterator[dict[str, Any] } +def _emit_whole(raw: Iterable[Any], cursor: _Cursor) -> Iterable[MessageStreamEvent]: + """Blocks of a turn that streamed no deltas, e.g. one reported whole.""" + for message in raw: + if isinstance(message, AssistantMessage): + for block in message.content: + yield from _emit_block(block, cursor.next_index) + cursor.next_index += 1 + + +def _self_started(result: ResultMessage) -> bool: + origin = result.origin + return origin is not None and origin.get("kind") != "human" + + +def _main_thread(message: Any) -> bool: + if getattr(message, "parent_tool_use_id", None) is not None: + return False + return isinstance( + message, (StreamEvent, AssistantMessage, UserMessage, ResultMessage) + ) + + +def _note_of(message: TaskNotificationMessage) -> str: + summary = message.summary or message.data.get("description") or message.task_id + return f"{summary} ({message.status})" + + def synthesize_turn_messages(raw: Iterable[Any]) -> list[dict[str, Any]]: out: list[dict[str, Any]] = [] for message in raw: diff --git a/src/beaver_gateway/backends/sessions.py b/src/beaver_gateway/backends/sessions.py index b85bf4a..feba836 100644 --- a/src/beaver_gateway/backends/sessions.py +++ b/src/beaver_gateway/backends/sessions.py @@ -20,7 +20,13 @@ from typing import TYPE_CHECKING, Any, Protocol import psutil if TYPE_CHECKING: - from collections.abc import AsyncIterable, AsyncIterator, Iterator, Mapping + from collections.abc import ( + AsyncIterable, + AsyncIterator, + Callable, + Iterator, + Mapping, + ) __all__ = ["DEFAULT_TTL", "Session", "SessionClient", "SessionPool", "cgroup_limit"] @@ -42,7 +48,7 @@ _RSS_HEADROOM = 0.8 class SessionClient(Protocol): async def connect(self) -> None: ... async def query(self, prompt: str | AsyncIterable[dict[str, Any]]) -> None: ... - def receive_response(self) -> AsyncIterator[Any]: ... + def receive_messages(self) -> AsyncIterator[Any]: ... async def interrupt(self) -> None: ... async def disconnect(self) -> None: ... @@ -67,6 +73,20 @@ class Session: state: dict[str, Any] = field(default_factory=dict) """Scratch for policy rules (``agents/policy.py``); dies with the process.""" + reader: asyncio.Task[None] | None = None + """Reads the client for the session's whole life, not only during a turn.""" + inbox: asyncio.Queue[Any] = field(default_factory=asyncio.Queue) + """Where the reader puts messages while a turn is attached.""" + pending: list[Any] = field(default_factory=list) + """A turn the CLI started on its own (a background task reported back), + buffered until a consumer attaches.""" + attached: bool = False + woken: bool = False + note: str | None = None + """Summary of the last task notification, the pending turn's prompt.""" + observer: Callable[[Any], None] | None = None + """The last turn's observer; sees subagent traffic between turns.""" + @property def busy(self) -> bool: return self.lock.locked() or self.running_turn is not None @@ -192,6 +212,8 @@ class SessionPool: if session is None: return _log.info("closing session %s (%s, %s)", session.session_id, session.kind, key) + if session.reader is not None and session.reader is not asyncio.current_task(): + session.reader.cancel() try: await session.client.disconnect() except Exception: # noqa: BLE001 @@ -238,6 +260,7 @@ class SessionPool: "turns": s.turns, "busy": s.busy, "running_turn": s.running_turn, + "pending": len(s.pending), "pending_question": s.pending_question, "pinned": s.pinned, "dirty": s.dirty, diff --git a/src/beaver_gateway/conversations/state.py b/src/beaver_gateway/conversations/state.py index cc7b493..9066a7e 100644 --- a/src/beaver_gateway/conversations/state.py +++ b/src/beaver_gateway/conversations/state.py @@ -127,6 +127,12 @@ class State: self._tasks = set() self._idle_task = None self.scheduler = None + for backend in backends.values(): + if hasattr(backend, "waker"): + backend.waker = self._woken + + async def _woken(self, key: str, note: str) -> None: + raise NotImplementedError @property def db(self) -> Database: diff --git a/src/beaver_gateway/conversations/texts.py b/src/beaver_gateway/conversations/texts.py index aed7797..43b73cc 100644 --- a/src/beaver_gateway/conversations/texts.py +++ b/src/beaver_gateway/conversations/texts.py @@ -77,6 +77,8 @@ class ConversationTexts: ) inject_header: Callable[[injects.InjectContext], str] = injects.inject_header bundle_header: str = "[injects accumulated since {since}; not the user]" + task_prompt: str = "[subagent reported back: {note}]" + """Stands in for the user's words in mirrors of a turn a subagent started.""" interrupted: str = "interrupted" answered: str = "The user answered: {answer}" unanswered: str = ( diff --git a/src/beaver_gateway/conversations/turns.py b/src/beaver_gateway/conversations/turns.py index 1df73b0..132c933 100644 --- a/src/beaver_gateway/conversations/turns.py +++ b/src/beaver_gateway/conversations/turns.py @@ -54,8 +54,15 @@ class Turns(Seeds): tools: bool = True, turn_id: str | None = None, item_origin: str | None = None, + shown: str | None = None, + attach: bool = False, ) -> AsyncIterator[MessageStreamEvent]: - """Run one turn under the conversation's lock; the only path to the backend.""" + """Run one turn under the conversation's lock; the only path to the backend. + + ``shown`` is what frontends display as the turn's prompt (the user's + words, not the envelope). ``attach`` sends nothing: the turn is one + the live session started by itself, a subagent reporting back. + """ row_id = cast("int", conv.id) runner = self._runner(row_id) backend = self._backend(conv.agent_name) @@ -65,7 +72,7 @@ class Turns(Seeds): async with runner.lock: runner.turn_id = turn_id runner.origin = origin - runner.text = _prompt_preview(messages) + 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) @@ -81,18 +88,28 @@ class Turns(Seeds): stop = "error" cut = False try: - events = backend.complete( - agent=self._claude_agent(conv.agent_name), - messages=messages, - conversation_id=conv.external_id, - session_id=resume if use_session else None, - reseed=not use_session, - capture=capture, - kind=conv.kind, - pinned=conv.kind == "master", - tools=tools, - observer=self._observer(conv, runner, turn_id, origin), - turn_id=turn_id, + observer = self._observer(conv, runner, turn_id, origin) + events = ( + backend.attach( + key=conv.external_id, + capture=capture, + observer=observer, + turn_id=turn_id, + ) + if attach + else backend.complete( + agent=self._claude_agent(conv.agent_name), + messages=messages, + conversation_id=conv.external_id, + session_id=resume if use_session else None, + reseed=not use_session, + capture=capture, + kind=conv.kind, + pinned=conv.kind == "master", + tools=tools, + observer=observer, + turn_id=turn_id, + ) ) async for event in events: yield event @@ -128,6 +145,8 @@ class Turns(Seeds): turn_id: str | None = None, item_origin: str | None = None, attachments: Sequence[dict[str, Any]] | None = None, + shown: str | None = None, + attach: bool = False, ) -> tuple[str, TurnCapture]: capture = TurnCapture() acc = StreamAccumulator() @@ -140,6 +159,8 @@ class Turns(Seeds): tools=tools, turn_id=turn_id, item_origin=item_origin, + shown=shown, + attach=attach, ): acc.feed(event) message = acc.finalize(model=agent.model) @@ -160,6 +181,47 @@ class Turns(Seeds): runner = self._runners.get(cast("int", conv.id)) return runner.origin if runner is not None and runner.turn_id else None + async def _woken(self, key: str, note: str) -> None: + """The live session of ``key`` started a turn by itself; run it next.""" + conv = await self.get(key) + if conv is None: + _log.warning("self-started turn on unknown conversation %s: %s", key, note) + return + _log.info("conversation %s woke by itself: %s", key, note) + self._ensure_worker(cast("int", conv.id)) + + async def _run_pending(self, conv: Conversation) -> None: + turn_id = f"turn_{uuid4().hex[:12]}" + note = self._backend(conv.agent_name).pending_note(conv.external_id) or "" + try: + text, _ = await self.run_text_turn( + conv, + note, + origin="task", + turn_id=turn_id, + item_origin="task", + shown=note, + attach=True, + ) + except Exception: # noqa: BLE001 + _log.exception( + "self-started turn %s on %s failed", turn_id, conv.external_id + ) + return + if not text.strip(): + return + self._bus.publish( + "reply", + conversation_id=conv.external_id, + turn_id=turn_id, + item=None, + item_origin="task", + source="task", + prompt=self._texts.task_prompt.format(note=note), + user_text=None, + text=text, + ) + def _ensure_worker(self, row_id: int) -> None: runner = self._runner(row_id) runner.wake.set() @@ -170,6 +232,10 @@ class Turns(Seeds): async def _worker(self, row_id: int) -> None: runner = self._runner(row_id) while True: + conv = await self.get_row(row_id) + if conv is not None and self._has_pending(conv): + await self._run_pending(conv) + continue items = await self._queue.pending(row_id) batch, wait = self._pick(items) if batch is None: @@ -233,6 +299,7 @@ class Turns(Seeds): turn_id=turn_id, item_origin=head.origin, attachments=head.attachments if origin == "user" else None, + shown=head.text if origin == "user" else None, ) except Exception: # noqa: BLE001 _log.exception("turn %s on %s failed", turn_id, conv.external_id) @@ -254,6 +321,12 @@ class Turns(Seeds): text=text, ) + def _has_pending(self, conv: Conversation) -> bool: + try: + return self._backend(conv.agent_name).has_pending(conv.external_id) + except LookupError: + return False + def _bundle(self, items: Sequence[InjectQueueItem]) -> str: lines = [self._texts.bundle_header.format(since=iso(items[0].created_at))] lines.extend(f"- [{i.origin}] {i.text}" for i in items) diff --git a/src/beaver_gateway/frontends/telegram/drafts.py b/src/beaver_gateway/frontends/telegram/drafts.py index 18fff98..f841c39 100644 --- a/src/beaver_gateway/frontends/telegram/drafts.py +++ b/src/beaver_gateway/frontends/telegram/drafts.py @@ -47,6 +47,7 @@ class Draft: self._interval = interval self.status = status self.text = "" + self._paused = False self._dirty = True self._broken = False self._last_sent = 0.0 @@ -65,11 +66,22 @@ class Draft: if status != self.status: self.status = status self._dirty = True + self._paused = bool(self.text.strip()) def append(self, text: str) -> None: - if text: + """Add streamed text; a tool call in between starts a new paragraph. + + The model ends a sentence, calls a tool, then goes on without a + break - the same text the final message shows as paragraphs. + """ + if not text: + return + if self._paused and not self.text.endswith("\n"): + self.text = self.text.rstrip() + "\n\n" + text.lstrip() + else: self.text += text - self._dirty = True + self._paused = False + self._dirty = True async def stop(self) -> None: """Stop pushing and wait for the push in flight. diff --git a/src/beaver_gateway/frontends/telegram/frontend.py b/src/beaver_gateway/frontends/telegram/frontend.py index 47d5fbd..fab4931 100644 --- a/src/beaver_gateway/frontends/telegram/frontend.py +++ b/src/beaver_gateway/frontends/telegram/frontend.py @@ -777,7 +777,7 @@ class TelegramFrontend(Frontend): return match kind: case "turn.start": - if event.get("origin") == "user": + if event.get("origin") in ("user", "task"): await self._open_draft(key, event, target) case "stream": self._feed_draft(key, event) diff --git a/tests/test_claude_sdk_backend.py b/tests/test_claude_sdk_backend.py index c841ca7..4920727 100644 --- a/tests/test_claude_sdk_backend.py +++ b/tests/test_claude_sdk_backend.py @@ -86,6 +86,7 @@ class FakeClient: self.prompts: list[str] = [] self.connected = False self.session_id = options.resume or "fresh-session" + self.asked = asyncio.Event() FakeClient.instances.append(self) async def connect(self) -> None: @@ -93,6 +94,17 @@ class FakeClient: async def query(self, prompt: str) -> None: self.prompts.append(prompt) + self.asked.set() + + async def receive_messages(self): + served = 0 + while True: + while len(self.prompts) <= served: + self.asked.clear() + await self.asked.wait() + served += 1 + async for message in self.receive_response(): + yield message async def interrupt(self) -> None: self.interrupted = True diff --git a/tests/test_conversations.py b/tests/test_conversations.py index 977e687..a7bad0e 100644 --- a/tests/test_conversations.py +++ b/tests/test_conversations.py @@ -8,6 +8,7 @@ from zoneinfo import ZoneInfo import pytest from claude_agent_sdk import ( + TaskNotificationMessage, AssistantMessage, InMemorySessionStore, ResultMessage, @@ -92,6 +93,8 @@ class ScriptedClient: self.session_id = options.resume or str(uuid.uuid4()) self.interrupted = False self.connected = False + self.asked = asyncio.Event() + self.extra: list[Any] = [] ScriptedClient.instances.append(self) async def connect(self) -> None: @@ -106,6 +109,26 @@ class ScriptedClient: b.get("text", "") for b in content if b.get("type") == "text" ) self.prompts.append(prompt) + self.asked.set() + + async def receive_messages(self): + """One response per prompt, forever; ``extra`` is what the CLI sends + on its own between prompts (a subagent reporting back).""" + served = 0 + while True: + while len(self.prompts) <= served: + self.asked.clear() + if self.extra: + for message in self.extra: + if isinstance(message, asyncio.Event): + await message.wait() + continue + yield message + self.extra = [] + await self.asked.wait() + served += 1 + async for message in self.receive_response(): + yield message async def receive_response(self): prompt = self.prompts[-1] @@ -858,3 +881,80 @@ async def test_inject_into_a_closed_master_lands_in_the_open_one(world: World) - assert item.conversation_id == new.id assert await world.statuses(old) == [] assert await world.statuses(new) == [("normal", "queued")] + + +def _self_started(session_id: str, text: str) -> list[Any]: + """What the CLI sends by itself once a background subagent reports back.""" + return [ + AssistantMessage( + content=[ToolUseBlock(id="tu_dig", name="Grep", input={"pattern": "au"})], + model="m", + parent_tool_use_id="tu_bg", + ), + TaskNotificationMessage( + subtype="task_notification", + data={}, + task_id="t1", + status="completed", + output_file="", + summary='Agent "dig" finished', + uuid="n1", + session_id=session_id, + ), + AssistantMessage(content=[TextBlock(text=text)], model="m"), + ResultMessage( + subtype="success", + duration_ms=1, + duration_api_ms=1, + is_error=False, + num_turns=1, + session_id=session_id, + stop_reason="end_turn", + total_cost_usd=0.0, + usage={"input_tokens": 1, "output_tokens": 1}, + origin={"kind": "task-notification"}, + ), + ] + + +async def test_subagent_report_between_turns_is_its_own_turn(world: World) -> None: + conv = await world.conversations.create(kind="master", agent="a", origin="test") + seen: list[dict[str, Any]] = [] + + async def collect() -> None: + async for event in world.bus.stream(conversation_id=conv.external_id): + seen.append(event) + + task = asyncio.create_task(collect()) + await world.conversations.post(conv, "go") + await world.settle(conv, 1) + client = ScriptedClient.instances[0] + client.extra = _self_started(client.session_id, "the agent found gold") + client.asked.set() + await asyncio.sleep(0.3) + task.cancel() + replies = [e for e in seen if e["type"] == "reply"] + assert [r["text"] for r in replies] == ["ok:go", "the agent found gold"] + assert replies[1]["item_origin"] == "task" + starts = [e for e in seen if e["type"] == "turn.start"] + assert [s["origin"] for s in starts] == ["user", "task"] + assert starts[1]["text"] == 'Agent "dig" finished (completed)' + assert len(client.prompts) == 1 + first_end = next(i for i, e in enumerate(seen) if e["type"] == "turn.end") + later = [e for e in seen[first_end + 1 :] if e["type"] in ("stream", "tool")] + assert later and all(e["turn_id"] == starts[0]["turn_id"] for e in later[:1]) + assert all(e["turn_id"] == starts[1]["turn_id"] for e in later[1:]) + assert (await world.conversations.get(conv.external_id)).running_turn is None + + +async def test_subagent_result_inside_a_turn_does_not_end_it(world: World) -> None: + conv = await world.conversations.create(kind="master", agent="a", origin="test") + await world.conversations.post(conv, "one") + await world.settle(conv, 1) + client = ScriptedClient.instances[0] + client.extra = _self_started(client.session_id, "late report") + await world.conversations.post(conv, "two") + await world.settle(conv, 2) + assert client.prompts[-1].startswith("two") + rows = await world.conversations.queue.recent(conv.id) + assert [r.status for r in rows] == ["done", "done"] diff --git a/tests/test_telegram.py b/tests/test_telegram.py index 30b3b19..9831984 100644 --- a/tests/test_telegram.py +++ b/tests/test_telegram.py @@ -4,7 +4,7 @@ import contextlib import tempfile from pathlib import Path from types import SimpleNamespace -from typing import Any +from typing import Any, cast import pytest from aiogram.exceptions import TelegramBadRequest, TelegramNetworkError @@ -16,6 +16,7 @@ from beaver_gateway.app import McpRegistry from beaver_gateway.backends.transcript import build_entries from beaver_gateway.frontends.base import GatewayRuntime from beaver_gateway.frontends.telegram import Attachments, TelegramFrontend +from beaver_gateway.frontends.telegram.drafts import Draft from beaver_gateway.frontends.telegram.render import ( LIMIT, chunks, @@ -719,6 +720,22 @@ def test_chunks_split_blockquote_and_stay_under_limit() -> None: assert " ".join(parts).split() == long_line.split() +def test_draft_breaks_a_paragraph_after_a_tool_call() -> None: + draft = Draft(cast("Any", None), chat_id=1, thread_id=None, turn_id="t") + draft.append("Checking the balance.") + draft.set_status("⏳ Bash") + draft.set_status("✍️ writing") + draft.append("It is 20 378.") + assert draft.text == "Checking the balance.\n\nIt is 20 378." + draft.set_status("⏳ Read") + draft.append("\n\nAnd a list:") + assert draft.text.endswith("It is 20 378.\n\nAnd a list:") + fresh = Draft(cast("Any", None), chat_id=1, thread_id=None, turn_id="u") + fresh.set_status("⏳ Bash") + fresh.append("First words") + assert fresh.text == "First words" + + def test_to_html_tail_inside_fence() -> None: full = ( "intro\n```py\n" + "\n".join(f"code {i} <" for i in range(200)) + "\n```\nafter" diff --git a/ui/src/lib/panel/activity-feed.svelte b/ui/src/lib/panel/activity-feed.svelte index e13e9e7..8dabd7b 100644 --- a/ui/src/lib/panel/activity-feed.svelte +++ b/ui/src/lib/panel/activity-feed.svelte @@ -23,7 +23,7 @@ onMount(() => { const timer = setInterval(() => { - if (model.running) { + if (model.live.length > 0) { now = Date.now(); } }, TICK_MS); diff --git a/ui/src/lib/panel/activity.svelte.ts b/ui/src/lib/panel/activity.svelte.ts index 20b4fb7..76ed785 100644 --- a/ui/src/lib/panel/activity.svelte.ts +++ b/ui/src/lib/panel/activity.svelte.ts @@ -27,6 +27,9 @@ export interface Turn { itemOrigin: string | null; nodes: Record; origin: string; + // A tool call came between two pieces of text; the next piece starts a + // paragraph, as the final message will show it. + paused: boolean; resultSubtype: string | null; roots: string[]; says: string[]; @@ -51,6 +54,20 @@ function str(value: unknown): string | null { return typeof value === "string" ? value : null; } +export function joinText(head: string, tail: string, paused: boolean): string { + if (paused && head.trim() && !head.endsWith("\n")) { + return `${head.trimEnd()}\n\n${tail.trimStart()}`; + } + return head + tail; +} + +export function isLive(turn: Turn): boolean { + return ( + turn.status === "running" || + Object.values(turn.nodes).some((node) => node.status === "running") + ); +} + function stopStatus(turn: Turn, stop: string | null): TurnStatus { if (stop === "end_turn") { return turn.status === "error" ? "error" : "done"; @@ -65,6 +82,9 @@ export class ActivityModel { turns = $state([]); question = $state(null); + // A subagent still working after its turn ended: nothing on this page + // knows the turn, so its tool calls hang under the launching Agent node + // of a turn named after that node. private readonly handlers: Record void> = { "conversation.created": (c) => this.onConversation(c), "conversation.updated": (c) => this.onConversation(c), @@ -89,6 +109,12 @@ export class ActivityModel { return this.turns.find((turn) => turn.status === "running") ?? null; } + // A turn is live while it runs or while a subagent it launched still + // works: the CLI reports subagent tool calls after the turn's own end. + get live(): Turn[] { + return this.turns.filter((turn) => isLive(turn)); + } + setConversation(info: ConversationInfo): void { this.conversation = info; this.question = info.question; @@ -151,6 +177,42 @@ export class ActivityModel { turn.startedAt = ts; turn.itemOrigin = str(event.item_origin); turn.userText = str(event.text) ?? turn.userText; + if (turn.origin === "task") { + this.settleAgents(ts); + } + } + + // The report is in: whatever Agent node still runs has finished. + private settleAgents(ts: string): void { + for (const turn of this.turns) { + if (turn.origin === "agent" && turn.status === "running") { + turn.status = "done"; + turn.endedAt = ts; + } + for (const node of Object.values(turn.nodes)) { + const agent = SUBAGENT_TOOLS.has(node.name) || node.name === "?"; + if (node.status === "running" && node.parent === null && agent) { + node.status = "done"; + node.endedAt = ts; + } + } + } + } + + // Subagent traffic belongs to the turn that launched the subagent, whatever + // turn the gateway was in when it arrived. + private turnFor(turnId: string, ts: string, parent: string | null): Turn { + if (parent) { + const owner = this.turns.find((turn) => parent in turn.nodes); + if (owner) { + return owner; + } + const current = this.turns.find((turn) => turn.id === turnId); + if (current?.status !== "running") { + return this.ensureTurn(`agent:${parent}`, ts, "agent"); + } + } + return this.ensureTurn(turnId, ts, null); } private onStream({ event, ts, turnId, parent }: Cursor): void { @@ -158,7 +220,7 @@ export class ActivityModel { if (!(turnId && raw && typeof raw === "object")) { return; } - const turn = this.ensureTurn(turnId, ts, null); + const turn = this.turnFor(turnId, ts, parent); const sdk = raw as Record; if (sdk.type === "content_block_start") { const block = sdk.content_block as Record | undefined; @@ -172,7 +234,8 @@ export class ActivityModel { if (sdk.type === "content_block_delta" && parent === null) { const delta = sdk.delta as Record | undefined; if (delta?.type === "text_delta" && typeof delta.text === "string") { - turn.text += delta.text; + turn.text = joinText(turn.text, delta.text, turn.paused); + turn.paused = false; } } } @@ -181,7 +244,10 @@ export class ActivityModel { if (!(turnId && typeof event.tool_use_id === "string")) { return; } - const turn = this.ensureTurn(turnId, ts, null); + const turn = this.turnFor(turnId, ts, parent); + if (parent === null) { + turn.paused = true; + } const node = this.ensureNode( turn, event.tool_use_id, @@ -196,7 +262,7 @@ export class ActivityModel { if (!(turnId && typeof event.tool_use_id === "string")) { return; } - const turn = this.ensureTurn(turnId, ts, null); + const turn = this.turnFor(turnId, ts, parent); const node = this.ensureNode(turn, event.tool_use_id, "?", parent, ts); node.status = event.is_error ? "error" : "done"; node.endedAt = ts; @@ -266,11 +332,13 @@ export class ActivityModel { }; } + // The turn's own tool calls die with it; a subagent's keep going until + // their own results arrive. private closeTurn(turn: Turn, status: TurnStatus, ts: string): void { turn.status = status; turn.endedAt = ts; for (const node of Object.values(turn.nodes)) { - if (node.status === "running") { + if (node.status === "running" && node.parent === null) { node.status = "aborted"; node.endedAt = ts; } @@ -291,6 +359,7 @@ export class ActivityModel { itemOrigin: null, nodes: {}, origin: origin ?? "?", + paused: false, resultSubtype: null, roots: [], says: [], @@ -394,3 +463,11 @@ export function summarizeInput(name: string, input: unknown): string { export function toolLabel(name: string): string { return name.startsWith("mcp__") ? name.slice(5).replace("__", " · ") : name; } + +// How a turn was started, as a word for the header. +export const ORIGIN_LABELS: Record = { + agent: "subagent at work", + inject: "inject", + task: "subagent report", + user: "user", +}; diff --git a/ui/src/lib/panel/chat-view.svelte b/ui/src/lib/panel/chat-view.svelte index 4fee77d..0ed9641 100644 --- a/ui/src/lib/panel/chat-view.svelte +++ b/ui/src/lib/panel/chat-view.svelte @@ -7,7 +7,7 @@ import { clip, fmtDateTime, fmtTime } from "$lib/format"; import { cn } from "$lib/utils"; import type { ActivityModel } from "./activity.svelte"; - import { summarizeInput, toolLabel } from "./activity.svelte"; + import { isLive, summarizeInput, toolLabel } from "./activity.svelte"; import { cacheableHistory, historyKey } from "./history-cache"; import { usePanelHost } from "./host"; import Markdown from "./markdown.svelte"; @@ -48,7 +48,7 @@ const tail = $derived( model.turns - .filter((turn) => turn.status === "running" || turn.startedAt > loadedAt) + .filter((turn) => isLive(turn) || turn.startedAt > loadedAt) .reverse() ); const tailSize = $derived( @@ -150,14 +150,16 @@ $effect(() => { const timer = setInterval(() => { - if (model.running) { + if (model.live.length > 0) { now = Date.now(); } }, TICK_MS); return () => clearInterval(timer); }); - const SYSTEM_HEAD = /^\[[^\]\n]+\]/; + // Text the gateway or the CLI put in the user's seat: an inject header, + // an envelope, a subagent's report. + const SYSTEM_HEAD = /^(\[[^\]\n]+\]|)/; const SECONDS = /:\d{2}$/; const DAY_TIME = /,?\s*\d{2}:\d{2}$/; let openSystem = $state>(new Set()); diff --git a/ui/src/lib/panel/tool-node.svelte b/ui/src/lib/panel/tool-node.svelte index 610ab58..7e46a9c 100644 --- a/ui/src/lib/panel/tool-node.svelte +++ b/ui/src/lib/panel/tool-node.svelte @@ -53,7 +53,7 @@ - {toolLabel(node.name)} + {subagent || node.name === "?" ? "Agent" : toolLabel(node.name)} {#if summary} {summary} diff --git a/ui/src/lib/panel/turn-card.svelte b/ui/src/lib/panel/turn-card.svelte index d138fa6..15042ef 100644 --- a/ui/src/lib/panel/turn-card.svelte +++ b/ui/src/lib/panel/turn-card.svelte @@ -9,7 +9,7 @@ fmtTokens, } from "$lib/format"; import { cn } from "$lib/utils"; - import type { Turn } from "./activity.svelte"; + import { isLive, ORIGIN_LABELS, type Turn } from "./activity.svelte"; import Markdown from "./markdown.svelte"; import ToolNodeView from "./tool-node.svelte"; @@ -18,10 +18,14 @@ const duration = $derived( turn.usage?.duration_ms ?? elapsedMs(turn.startedAt, turn.endedAt, now) ); + const origin = $derived(ORIGIN_LABELS[turn.origin] ?? turn.origin); const originLabel = $derived( turn.itemOrigin && turn.itemOrigin !== turn.origin - ? `${turn.origin} · ${turn.itemOrigin}` - : turn.origin + ? `${origin} · ${turn.itemOrigin}` + : origin + ); + const pill = $derived( + turn.status !== "running" && isLive(turn) ? "running" : turn.status ); const toolCount = $derived(Object.keys(turn.nodes).length); @@ -29,15 +33,20 @@
- + {#if turn.origin !== "agent"} + + {/if} {originLabel} @@ -65,7 +74,9 @@ {/if}
- {#if turn.userText} + {#if turn.userText && turn.origin === "task"} +

{turn.userText}

+ {:else if turn.userText} - {:else if turn.status === "running" && turn.roots.length === 0 && turn.thinking === 0} + {:else if turn.status === "running" && turn.origin !== "agent" && turn.roots.length === 0 && turn.thinking === 0}

Waiting for the model…