From bce1f2f618809238a4398a747ebe3b731fa0877d Mon Sep 17 00:00:00 2001 From: galuszkm Date: Wed, 23 Sep 2026 23:50:05 +0200 Subject: [PATCH] fix(renderers): keep concurrent agent output in one block per agent - Track render state per agent so a neighbour's tokens cannot split a block - Hand the terminal to a mid-stream agent so it never buffers invisibly - Refuse the terminal to an agent with outstanding tool calls - Reset state on flush so a reused renderer reprints headers --- src/strands_compose/renderers/ansi.py | 348 +++++++++++++++++--------- tests/runtime/test_renderers.py | 73 ++++++ 2 files changed, 299 insertions(+), 122 deletions(-) diff --git a/src/strands_compose/renderers/ansi.py b/src/strands_compose/renderers/ansi.py index 921007b..83fb28d 100644 --- a/src/strands_compose/renderers/ansi.py +++ b/src/strands_compose/renderers/ansi.py @@ -19,12 +19,38 @@ import sys import time from collections.abc import Callable +from dataclasses import dataclass, field from typing import Any from ..types import EventType, SessionManifest, StreamEvent from .base import EventRenderer +@dataclass +class _Branch: + """Render state for one agent's output. + + strands runs the tool calls of a single turn concurrently, so agents + delegated in the same turn interleave their events. Per-agent state keeps + a neighbour's tokens from splitting an agent's block. + + Attributes: + mode: Open stream mode (``"reasoning"``, ``"responding"``, or ``None`` + when the next chunk needs a fresh header). + in_stream: Whether output is mid-line in a token/reasoning stream. + ended: Whether the agent emitted its terminal event, so its buffered + block is complete. + pending_tools: Number of the agent's tool calls still running. + buffer: Rendered output held back while another agent is live. + """ + + mode: str | None = None + in_stream: bool = False + ended: bool = False + pending_tools: int = 0 + buffer: list[str] = field(default_factory=list) + + class AnsiRenderer(EventRenderer): """Renders events using raw ANSI escape codes.""" @@ -42,8 +68,13 @@ def __init__( Token and reasoning events are written inline (no trailing newline) so they appear as a continuous stream. A separator line is printed - when the mode changes between reasoning and responding, or when the - active agent changes. + when an agent's mode changes between reasoning and responding. + + Concurrent agents are serialised rather than interleaved: one agent + holds the terminal and streams live while the others' output is + buffered, then written as a contiguous block once the live agent + yields. Each agent keeps a single header and an unbroken block, at + the cost of that block appearing later than it was produced. Args: file: Output stream (defaults to ``sys.stdout``). @@ -54,11 +85,12 @@ def __init__( the effect entirely with no overhead. """ self._out = file or sys.stdout - self._in_stream = False - self._mode: str | None = None # "reasoning" | "responding" | None - self._active_agent: str | None = None self._typewriter_delay = typewriter_delay + # Insertion order decides the order buffered blocks are written in. + self._branches: dict[str, _Branch] = {} + self._owner: str | None = None + # Cache TTY state and pre-compute ANSI escape strings. is_tty = hasattr(self._out, "isatty") and self._out.isatty() self._dim = "\033[2m" if is_tty else "" @@ -101,187 +133,260 @@ def render(self, event: StreamEvent) -> None: handler(event) def flush(self) -> None: - if self._in_stream: - self._out.write("\n") - self._out.flush() - self._in_stream = False + """Close the live stream, write every buffered block, and reset state. + + An agent that never emitted a terminal event still has output to write, + so all remaining buffers are drained here. State is cleared because a + renderer is reused across turns. + """ + if self._owner is not None: + self._break(self._owner) + for agent, branch in self._branches.items(): + self._write_buffer(agent) + if branch.in_stream: + self._out.write("\n") + self._out.flush() + self._branches.clear() + self._owner = None # -- Per-event-type handlers ------------------------------------------- def _handle_session_start(self, event: StreamEvent) -> None: - self._break() manifest = SessionManifest.model_validate(event.data["manifest"]) agent_names = ", ".join(a.name for a in manifest.agents) or "—" orch_names = ", ".join(o.name for o in manifest.orchestrations) or "—" - self._out.write(self._separator(manifest.entry.name, "SESSION START", color=self._cyan)) - self._out.write( - f" {self._dim}entry: {manifest.entry.name} ({manifest.entry.kind})\n" - f" agents: {agent_names}\n" + text = ( + self._separator(manifest.entry.name, "SESSION START", color=self._cyan) + + f" {self._dim}entry: {manifest.entry.name} ({manifest.entry.kind})\n" + + f" agents: {agent_names}\n" ) if manifest.orchestrations: - self._out.write(f" orchestrations: {orch_names}\n") - self._out.write(self._reset) - self._out.flush() + text += f" orchestrations: {orch_names}\n" + self._write_direct(text + self._reset) def _handle_session_end(self, event: StreamEvent) -> None: - self._break() session_id = event.data.get("session_id") sid_str = session_id if session_id else "—" - self._out.write(self._separator(event.agent_name, "SESSION END", color=self._cyan)) - self._out.write(f" {self._dim}session_id: {sid_str}{self._reset}\n") - self._out.flush() + self._write_direct( + self._separator(event.agent_name, "SESSION END", color=self._cyan) + + f" {self._dim}session_id: {sid_str}{self._reset}\n" + ) def _handle_token(self, event: StreamEvent) -> None: text = event.data.get("text", "") - already_streaming = self._mode == "responding" and self._active_agent == event.agent_name - if not already_streaming and not text.strip(): + branch = self._claim(event.agent_name) + if branch.mode != "responding" and not text.strip(): return self._ensure_mode(event.agent_name, "responding") - self._write_with_delay(text) - self._in_stream = True + self._emit(event.agent_name, text, stream=True) + branch.in_stream = True def _handle_reasoning(self, event: StreamEvent) -> None: text = event.data.get("text", "") - already_streaming = self._mode == "reasoning" and self._active_agent == event.agent_name - if not already_streaming and not text.strip(): + branch = self._claim(event.agent_name) + if branch.mode != "reasoning" and not text.strip(): return self._ensure_mode(event.agent_name, "reasoning") - self._out.write(self._yellow) - self._write_with_delay(text) - self._out.write(self._reset) - self._out.flush() - self._in_stream = True + self._emit(event.agent_name, self._yellow) + self._emit(event.agent_name, text, stream=True) + self._emit(event.agent_name, self._reset) + branch.in_stream = True def _handle_agent_start(self, event: StreamEvent) -> None: - self._break() - self._mode = None - self._active_agent = None - self._out.write(self._separator(event.agent_name, "AGENT START", color=self._magenta)) - self._out.write(f"{self._cyan}{self._bold}[{event.agent_name}]{self._reset} starting…\n") - self._out.flush() + agent = event.agent_name + self._start_line(agent) + self._emit( + agent, + self._separator(agent, "AGENT START", color=self._magenta) + + f"{self._cyan}{self._bold}[{agent}]{self._reset} starting…\n", + ) def _handle_tool_start(self, event: StreamEvent) -> None: - self._break() - self._mode = None - self._active_agent = None - data = event.data - label = data.get("tool_name", "unknown") - tool_input = data.get("tool_input", {}) - preview = str(tool_input)[:80] + ("…" if len(str(tool_input)) > 80 else "") - self._out.write(self._separator(event.agent_name, "TOOL USE", color=self._magenta)) - self._out.write( - f" {self._yellow}⚙{self._reset} [{event.agent_name}] -> {label!r}" - f" {self._dim}{preview}{self._reset}\n" + agent = event.agent_name + self._start_line(agent) + label = event.data.get("tool_name", "unknown") + tool_input = str(event.data.get("tool_input", {})) + preview = tool_input[:80] + ("…" if len(tool_input) > 80 else "") + self._emit( + agent, + self._separator(agent, "TOOL USE", color=self._magenta) + + f" {self._yellow}⚙{self._reset} [{agent}] -> {label!r}" + f" {self._dim}{preview}{self._reset}\n", ) - self._out.flush() + # The agent is blocked on the tool, so the sub-agent it called runs next. + self._branches[agent].pending_tools += 1 + self._release(agent) def _handle_tool_end(self, event: StreamEvent) -> None: - self._break() - self._mode = None - status = event.data.get("status", "?") agent = event.agent_name - if status == "error": - self._out.write( - f" {self._red}✗{self._reset} [{agent}] tool error: {event.data.get('error')}\n" + self._start_line(agent) + branch = self._branches[agent] + # A resumed interrupt can deliver a tool end whose start was a turn ago. + branch.pending_tools = max(branch.pending_tools - 1, 0) + if event.data.get("status", "?") == "error": + self._emit( + agent, + f" {self._red}✗{self._reset} [{agent}] tool error: {event.data.get('error')}\n", ) else: - self._out.write(f" {self._green}✓{self._reset} [{agent}] tool done\n") - self._out.flush() + self._emit(agent, f" {self._green}✓{self._reset} [{agent}] tool done\n") def _handle_agent_complete(self, event: StreamEvent) -> None: - self._break() - self._mode = None - self._active_agent = None + agent = event.agent_name + self._start_line(agent) usage = event.data.get("usage", {}) in_tokens = usage.get("input_tokens", 0) out_tokens = usage.get("output_tokens", 0) - self._out.write( - f" {self._dim}✅ [{event.agent_name}] complete ({in_tokens} input, {out_tokens} output tokens){self._reset}\n" + self._emit( + agent, + f" {self._dim}✅ [{agent}] complete " + f"({in_tokens} input, {out_tokens} output tokens){self._reset}\n", ) - self._out.flush() + self._end(agent) def _handle_interrupt(self, event: StreamEvent) -> None: - self._break() - self._mode = None - self._active_agent = None + agent = event.agent_name + self._start_line(agent) data = event.data name = data.get("name") or "—" reason = data.get("reason") or "no reason given" - self._out.write(self._separator(event.agent_name, "INTERRUPT", color=self._yellow)) - self._out.write( - f" {self._yellow}⏸{self._reset} [{event.agent_name}] awaiting input for {name!r}: {reason}\n" - f" {self._dim}interrupt_id: {data.get('interrupt_id') or '—'}{self._reset}\n" + self._emit( + agent, + self._separator(agent, "INTERRUPT", color=self._yellow) + + f" {self._yellow}⏸{self._reset} [{agent}] awaiting input for {name!r}: {reason}\n" + + f" {self._dim}interrupt_id: {data.get('interrupt_id') or '—'}{self._reset}\n", ) - self._out.flush() def _handle_error(self, event: StreamEvent) -> None: - self._break() - self._mode = None - self._out.write(self._separator(event.agent_name, "ERROR", color=self._red)) + agent = event.agent_name + self._start_line(agent) msg = event.data.get("text", "unknown error") exc_type = event.data.get("exception_type") prefix = f"{exc_type}: " if exc_type else "" - self._out.write(f" {self._red}✗ [{event.agent_name}] ERROR: {prefix}{msg}{self._reset}\n") - self._out.flush() + self._emit( + agent, + self._separator(agent, "ERROR", color=self._red) + + f" {self._red}✗ [{agent}] ERROR: {prefix}{msg}{self._reset}\n", + ) + # An ERROR suppresses the agent's AGENT_COMPLETE, so nothing more follows. + self._end(agent) def _handle_node_start(self, event: StreamEvent) -> None: - self._break() - self._out.write( + self._write_direct( f"\n{self._cyan}->{self._reset} node '{event.data.get('node_id')}' starting\n" ) - self._out.flush() def _handle_node_stop(self, event: StreamEvent) -> None: - self._break() - self._out.write(f"{self._cyan}←{self._reset} node '{event.data.get('node_id')}' done\n") - self._out.flush() + self._write_direct( + f"{self._cyan}←{self._reset} node '{event.data.get('node_id')}' done\n" + ) def _handle_handoff(self, event: StreamEvent) -> None: - self._break() to_ids = ", ".join(event.data.get("to_node_ids", [])) - self._out.write(f" {self._cyan}↪{self._reset} handoff -> {to_ids}\n") - self._out.flush() + self._write_direct(f" {self._cyan}↪{self._reset} handoff -> {to_ids}\n") def _handle_multiagent_start(self, event: StreamEvent) -> None: - self._break() kind = event.data.get("multiagent_type", "") - self._out.write(f"\n{self._cyan}⊕{self._reset} {kind} orchestration starting\n") - self._out.flush() + self._write_direct(f"\n{self._cyan}⊕{self._reset} {kind} orchestration starting\n") def _handle_multiagent_complete(self, event: StreamEvent) -> None: - self._break() kind = event.data.get("multiagent_type", "") - self._out.write(f"{self._cyan}⊗{self._reset} {kind} orchestration complete\n\n") - self._out.flush() + self._write_direct(f"{self._cyan}⊗{self._reset} {kind} orchestration complete\n\n") # -- Internal helpers -------------------------------------------------- - def _write_with_delay(self, text: str) -> None: - """Write *text* to the output stream, sleeping between printable characters. + def _claim(self, agent: str) -> _Branch: + """Give *agent* the terminal when it is free and return its state. + + An agent taking the terminal writes the backlog it built up while + another agent was live, keeping its block contiguous. An agent still + waiting on tool calls is refused, so it cannot hold the terminal while + the sub-agent it called is producing output. + """ + branch = self._branches.setdefault(agent, _Branch()) + if self._owner is None and not branch.pending_tools: + self._owner = agent + self._write_buffer(agent) + return branch + + def _start_line(self, agent: str) -> None: + """Claim the terminal for *agent* and start a structured line. + + Closes an open stream and forces a fresh header on the next chunk. + """ + self._claim(agent) + self._break(agent) + self._branches[agent].mode = None + + def _end(self, agent: str) -> None: + """Mark *agent* finished and hand back the terminal it held.""" + self._branches[agent].ended = True + self._release(agent) + + def _release(self, agent: str) -> None: + """Free the terminal held by *agent* and write out the blocks now due. - When :attr:`_typewriter_delay` is ``0.0`` the text is written in a - single call with no per-character overhead. Otherwise each printable - character is written individually, flushed, and followed by a - ``time.sleep`` call so that the typewriter effect is visible. - Whitespace and control characters are written without a delay to - avoid stutter on word boundaries. + Finished blocks are written first, then the terminal goes to an agent + that is mid-stream. Without that hand-off an idle agent would claim it + and the streaming agent would stay buffered and invisible. """ - if self._typewriter_delay <= 0.0: - self._out.write(text) - self._out.flush() + if self._owner != agent: return - for char in text: - self._out.write(char) - self._out.flush() - if char.isprintable() and not char.isspace(): - time.sleep(self._typewriter_delay) - - def _break(self) -> None: - """Insert a newline if we are mid-token/reasoning stream.""" - if self._in_stream: - self._out.write("\n") - self._out.flush() - self._in_stream = False + self._owner = None + for name, branch in self._branches.items(): + if branch.ended: + self._write_buffer(name) + for name, branch in self._branches.items(): + if not branch.ended and branch.in_stream: + self._owner = name + self._write_buffer(name) + return + + def _emit(self, agent: str, text: str, *, stream: bool = False) -> None: + """Write *text* for *agent*, buffering it while another agent is live. + + Args: + agent: Agent the output belongs to. + text: Rendered text to write. + stream: Whether to apply the typewriter delay. Whitespace and + control characters are never delayed, to avoid stutter on word + boundaries, and buffered blocks always replay at full speed. + """ + if agent != self._owner: + self._branches[agent].buffer.append(text) + return + if stream and self._typewriter_delay > 0.0: + for char in text: + self._out.write(char) + self._out.flush() + if char.isprintable() and not char.isspace(): + time.sleep(self._typewriter_delay) + return + self._out.write(text) + self._out.flush() + + def _write_direct(self, text: str) -> None: + """Write *text* that belongs to no single agent, breaking the live stream.""" + if self._owner is not None: + self._break(self._owner) + self._out.write(text) + self._out.flush() + + def _write_buffer(self, agent: str) -> None: + """Write *agent*'s buffered block straight to the output stream.""" + branch = self._branches[agent] + if not branch.buffer: + return + self._out.write("".join(branch.buffer)) + branch.buffer.clear() + self._out.flush() + + def _break(self, agent: str) -> None: + """Insert a newline if *agent* is mid-token/reasoning stream.""" + branch = self._branches[agent] + if branch.in_stream: + self._emit(agent, "\n") + branch.in_stream = False def _separator(self, agent: str, label: str, color: str | None = None) -> str: """Build a separator line like ``── agent — REASONING ──``.""" @@ -293,13 +398,12 @@ def _separator(self, agent: str, label: str, color: str | None = None) -> str: return f"\n{c}{'─' * left}{inner}{'─' * right}{self._reset}\n" def _ensure_mode(self, agent: str, mode: str) -> None: - """Print a separator when the agent or mode changes.""" - if mode == self._mode and agent == self._active_agent: + """Print a separator when *agent*'s mode changes.""" + branch = self._branches[agent] + if branch.mode == mode: return - self._break() - self._active_agent = agent - self._mode = mode + self._break(agent) + branch.mode = mode label = "REASONING" if mode == "reasoning" else "RESPONDING" color = self._yellow if mode == "reasoning" else self._cyan - self._out.write(self._separator(agent, label, color=color)) - self._out.flush() + self._emit(agent, self._separator(agent, label, color=color)) diff --git a/tests/runtime/test_renderers.py b/tests/runtime/test_renderers.py index df8437c..a660a33 100644 --- a/tests/runtime/test_renderers.py +++ b/tests/runtime/test_renderers.py @@ -2,6 +2,9 @@ Renders into a StringIO and asserts the renderer surfaces the event's key data (agent name, node id, status). Formatting/colour is not pinned. + +Agents delegated in the same turn run concurrently and interleave their events, +so the last tests cover serialising them back into one block per agent. """ from __future__ import annotations @@ -27,6 +30,11 @@ def _ev(kind, agent="worker", **data) -> StreamEvent: return StreamEvent(type=kind, agent_name=agent, data=data) +def _headers(out: str, agent: str, label: str) -> int: + """Count the separator headers rendered for *agent* / *label*.""" + return out.count(f" {agent} \u2014 {label} ") + + def test_token_text_is_written(): assert "hello" in _render(_ev(EventType.TOKEN, text="hello")) @@ -141,3 +149,68 @@ def test_every_event_type_renders_something(kind): """No event type may be silently dropped by the renderer.""" out = _render(StreamEvent(type=kind, agent_name="worker", data=_MINIMAL_DATA[kind])) assert out != "" + + +def test_concurrent_agents_get_one_header_and_a_contiguous_block(): + out = _render( + _ev(EventType.TOKEN, agent="researcher", text="aaa"), + _ev(EventType.TOKEN, agent="writer", text="xxx"), + _ev(EventType.TOKEN, agent="researcher", text="bbb"), + _ev(EventType.TOKEN, agent="writer", text="yyy"), + _ev(EventType.AGENT_COMPLETE, agent="researcher", usage={}), + _ev(EventType.AGENT_COMPLETE, agent="writer", usage={}), + ) + assert "aaabbb" in out and "xxxyyy" in out + assert _headers(out, "researcher", "RESPONDING") == 1 + assert _headers(out, "writer", "RESPONDING") == 1 + + +def test_streaming_sub_agent_takes_over_when_the_live_one_finishes(): + out = _render( + _ev(EventType.TOOL_START, agent="main", tool_name="researcher", tool_input={}), + _ev(EventType.TOOL_START, agent="main", tool_name="writer", tool_input={}), + _ev(EventType.TOKEN, agent="researcher", text="facts"), + _ev(EventType.TOKEN, agent="writer", text="draft"), + _ev(EventType.AGENT_COMPLETE, agent="researcher", usage={}), + _ev(EventType.TOOL_END, agent="main", status="success"), + _ev(EventType.TOKEN, agent="writer", text=" continued"), + _ev(EventType.AGENT_COMPLETE, agent="writer", usage={}), + _ev(EventType.TOOL_END, agent="main", status="success"), + _ev(EventType.TOKEN, agent="main", text="synthesis"), + _ev(EventType.AGENT_COMPLETE, agent="main", usage={}), + ) + assert "draft continued" in out + assert out.index("facts") < out.index("draft") < out.index("synthesis") + + +def test_delegator_waiting_on_tools_does_not_take_the_terminal(): + out = _render( + _ev(EventType.TOOL_START, agent="main", tool_name="researcher", tool_input={}), + _ev(EventType.TOOL_START, agent="main", tool_name="writer", tool_input={}), + _ev(EventType.TOKEN, agent="researcher", text="facts"), + _ev(EventType.AGENT_COMPLETE, agent="researcher", usage={}), + _ev(EventType.TOOL_END, agent="main", status="success"), + _ev(EventType.TOKEN, agent="writer", text="draft"), + _ev(EventType.AGENT_COMPLETE, agent="writer", usage={}), + _ev(EventType.TOOL_END, agent="main", status="success"), + _ev(EventType.TOKEN, agent="main", text="synthesis"), + _ev(EventType.AGENT_COMPLETE, agent="main", usage={}), + ) + assert out.index("facts") < out.index("draft") < out.index("synthesis") + + +def test_buffered_output_survives_a_missing_terminal_event(): + out = _render( + _ev(EventType.TOKEN, agent="researcher", text="aaa"), + _ev(EventType.TOKEN, agent="writer", text="xxx"), + ) + assert "xxx" in out + + +def test_flush_resets_state_between_turns(): + buf = io.StringIO() + renderer = AnsiRenderer(file=buf, separator_width=40) + for _ in range(2): + renderer.render(_ev(EventType.TOKEN, agent="main", text="hi")) + renderer.flush() + assert _headers(buf.getvalue(), "main", "RESPONDING") == 2