From 0d125acb31decda935f28588692973011ef31a27 Mon Sep 17 00:00:00 2001 From: John Toman Date: Wed, 10 Jun 2026 12:07:32 -0700 Subject: [PATCH 1/2] Adds Live, claude style console interface A live updating region of active agents, subagents, and tool calls --- composer/cli/live_autoprove.py | 39 +++ composer/io/graph_runner.py | 73 ++++-- composer/ui/autoprove_live.py | 205 +++++++++++++++ composer/ui/conversation_client.py | 1 - composer/ui/live_display.py | 400 +++++++++++++++++++++++++++++ composer/ui/tool_display.py | 48 +++- 6 files changed, 735 insertions(+), 31 deletions(-) create mode 100644 composer/cli/live_autoprove.py create mode 100644 composer/ui/autoprove_live.py create mode 100644 composer/ui/live_display.py diff --git a/composer/cli/live_autoprove.py b/composer/cli/live_autoprove.py new file mode 100644 index 00000000..fd7ed112 --- /dev/null +++ b/composer/cli/live_autoprove.py @@ -0,0 +1,39 @@ +"""Entry point for the auto-prove pipeline — live-display mode. + +Renders the pipeline as a single inline ``rich.live.Live`` region +showing the tree of active agents (one root per phase task, sub-agents +nested below). Scrollback above the Live region captures phase +boundaries, prover lifecycle terminals, and per-rule analysis events. + +Sister to ``tui-autoprove`` (Textual full-screen TUI) and +``console-autoprove`` (plain ``print``). +""" + +import asyncio + +import composer.bind as _ + +from composer.diagnostics.timing import RunSummary +from composer.spec.source.autoprove_common import _entry_point +from composer.ui.autoprove_live import AutoProveLiveHandler + + +async def _main() -> int: + summary = RunSummary() + async with _entry_point(summary) as run: + async with AutoProveLiveHandler() as handler: + result = await run(handler.make_handler) + print(f"\n{'=' * 60}") + print(summary.format()) + print(f"\n Components: {result.n_components}") + print(f" Properties: {result.n_properties}") + if result.failures: + print(f" Failures: {len(result.failures)}") + for f in result.failures: + print(f" - {f}") + print(f"{'=' * 60}") + return 0 + + +def main() -> int: + return asyncio.run(_main()) diff --git a/composer/io/graph_runner.py b/composer/io/graph_runner.py index 4db36b8e..46539566 100644 --- a/composer/io/graph_runner.py +++ b/composer/io/graph_runner.py @@ -24,6 +24,51 @@ from langchain_core.runnables import RunnableConfig +def _normalize_updates_payload(payload: dict) -> dict: + """Flatten langgraph's mixed-mode tools-node payload to ``dict[node, dict]``. + + When a ``ToolNode`` batch contains tools whose returns differ in shape + (some return raw ``ToolMessage`` / ``str`` / dicts; some return + ``Command(update=...)`` with extra state-channel deltas), langgraph's + ``_combine_tool_outputs`` emits the node's value as a *list* of separate + updates instead of a single merged dict — see + ``langgraph/prebuilt/tool_node.py``. Every downstream ``log_state_update`` + handler in this codebase expects ``dict[node, dict_with_messages]`` and + silently drops the list shape (either via ``isinstance(update, dict)`` + guards or ``"messages" in update`` membership checks, which return + ``False`` against a list of dicts/Commands). The result is that any + ToolMessage produced by a parallel batch with at least one + ``Command``-returning tool never reaches the renderer — the tool widgets + stay "pending" until something else forces a redraw. + + Merge here so handlers stay simple. ``messages`` lists are concatenated; + other state keys retain last-write-wins semantics (matches what + langgraph's reducers would do for non-message channels, which the + handlers in this codebase don't introspect for ingest purposes). + """ + out: dict[str, Any] = {} + for node_name, value in payload.items(): + if isinstance(value, dict) or not isinstance(value, list): + out[node_name] = value + continue + merged: dict[str, Any] = {} + for item in value: + d = item.update if isinstance(item, Command) and isinstance(item.update, dict) else item + if not isinstance(d, dict): + continue + for k, v in d.items(): + if ( + k == "messages" + and isinstance(merged.get(k), list) + and isinstance(v, list) + ): + merged[k] = [*merged[k], *v] + else: + merged[k] = v + out[node_name] = merged + return out + + class SinkProtocol(Protocol): """Write-only event sink. Synchronous — must not block.""" def __call__(self, event: GraphEvents) -> None: @@ -111,29 +156,13 @@ async def run_graph[H, S: StateLike, I: StateLike, C: StateLike | None]( graph_input = Command(resume=human_response) interrupted = True break - elif ty == "custom": - event_sink( - CustomUpdate(payload, thread_id=tid, checkpoint_id=curr_checkpoint) # pyright: ignore[reportPossiblyUnboundVariable] - ) - else: - assert ty == "updates" - if "__interrupt__" in payload: - assert human_handler is not None - if "configurable" in curr_config and "checkpoint_id" in curr_config["configurable"]: - del curr_config["configurable"]["checkpoint_id"] - interrupt_data = cast(H, payload["__interrupt__"][0].value) - curr_state = cast(S, (await graph.aget_state({"configurable": {"thread_id": tid}})).values) - human_response = await human_handler(interrupt_data, curr_state) - graph_input = Command(resume=human_response) - interrupted = True - break - event_sink( - StateUpdate( - payload, thread_id=tid - ) + event_sink( + StateUpdate( + _normalize_updates_payload(payload), thread_id=tid ) - if interrupted: - continue + ) + if interrupted: + continue result_state = (await graph.aget_state({"configurable": {"thread_id": tid}})).values return cast(S, result_state) diff --git a/composer/ui/autoprove_live.py b/composer/ui/autoprove_live.py new file mode 100644 index 00000000..bb80ad34 --- /dev/null +++ b/composer/ui/autoprove_live.py @@ -0,0 +1,205 @@ +"""Live-display handler for the auto-prove pipeline. + +Drives an inline ``rich.live.Live`` region that shows the tree of +active agents (one root per phase task; sub-agents nest below). +Scrollback above the Live region accumulates phase headers, prover +lifecycle terminals, and errors. Per-line subprocess output streams +(``prover_output``, ``cloud_polling``, ``auto_setup_output``) are +dropped — there's no real way to surface a continuous stdout stream +in a Live region without fighting its redraw cycle. + +Sister to: + + - ``AutoProveApp`` (Textual full-screen TUI) — for users who want + per-task panels and scrollable log widgets. + - ``AutoProveConsoleHandler`` (plain ``print``) — for users who + want raw log output (CI, redirected pipes, etc.). + +Pattern follows ``AutoProveConsoleHandler``: one handler instance per +pipeline run, ``make_handler`` returns a ``TaskHandle`` that wraps the +same instance for every phase so all roots feed into a single Live +region. +""" + +from contextlib import asynccontextmanager +from typing import AsyncIterator, Callable, cast, override + +from rich.console import Console, RenderableType +from rich.text import Text + +from composer.io.conversation import ConversationClient +from composer.io.event_handler import NullEventHandler +from composer.io.multi_job import TaskHandle, TaskInfo +from composer.spec.source.prover import ProverEvents +from composer.spec.source.autosetup import AutoSetupEvents +from composer.ui.autoprove_app import AUTOPROVE_PHASE_LABELS, AutoProvePhase +from composer.ui.conversation_client import ConsoleConversationClient +from composer.ui.live_display import LiveDisplayHandler +from composer.ui.tool_display import ToolDisplayConfig + + +class AutoProveLiveHandler(LiveDisplayHandler[None], NullEventHandler): + """``IOHandler[None]`` + ``EventHandler`` + ``HandlerFactory`` for + the auto-prove pipeline, rendered as a single rich.live region. + + Use as:: + + async with AutoProveLiveHandler() as handler: + await executor(handler.make_handler) + """ + + def __init__( + self, + *, + tool_display_config: ToolDisplayConfig | None = None, + console: Console | None = None, + ) -> None: + super().__init__( + tool_display_config=tool_display_config or ToolDisplayConfig(), + console=console, + ) + # Tracks which phases have already produced a header. The first + # task within a phase to fire ``on_start`` emits the banner; + # subsequent parallel tasks in the same phase don't repeat it. + # Safe without a lock: ``on_start`` is sync and asyncio is + # single-threaded, so the check-and-add doesn't race. + self._phases_seen: set[AutoProvePhase] = set() + + # ── IOHandler hook: HITL — autoprove has none ──────────────────── + + @override + async def handle_human_interaction( + self, ty: None, debug_thunk: Callable[[], None] + ) -> str: + raise RuntimeError( + "Unexpected HITL interrupt in auto-prove live handler" + ) + + # ── EventHandler hooks ────────────────────────────────────────── + + @override + async def handle_event( + self, payload: dict, path: list[str], checkpoint_id: str + ) -> None: + evt = cast(ProverEvents, payload) + thread_id = path[-1] + match evt["type"]: + case "prover_run": + self._update_progress(thread_id, "running prover") + case "prover_link": + self._scrollback( + Text( + f"[{self._label(path)}] prover link → {evt['link']}", + style="dim", + ) + ) + case "prover_result": + self._update_progress(thread_id, None) + self._scrollback( + Text(f"[{self._label(path)}] prover complete", style="green") + ) + case "cex_analysis": + self._update_progress( + thread_id, f"analyzing CEX: {evt['rule_name']}" + ) + case "rule_analysis": + self._scrollback( + Text( + f"[{self._label(path)}] rule analysis → {evt['rule']}", + style="dim", + ) + ) + case "prover_output" | "cloud_polling": + # Per-line subprocess stream — no good way to surface a + # continuous stdout flow in a Live region. Use AutoProveApp + # or AutoProveConsoleHandler for raw prover output. + pass + + @override + async def handle_progress_event(self, payload: dict) -> None: + evt = cast(AutoSetupEvents, payload) + match evt["type"]: + case "auto_setup_start": + self._scrollback(Text("AutoSetup: started", style="dim")) + case "auto_setup_complete": + self._scrollback( + Text( + f"AutoSetup: complete (rc={evt['return_code']})", + style="dim", + ) + ) + case "auto_setup_output": + # Per-line subprocess stream — same reasoning as + # prover_output / cloud_polling above. + pass + + # ── ConversationProvider ──────────────────────────────────────── + + @asynccontextmanager + async def _start_conversation( + self, initial: RenderableType + ) -> AsyncIterator[ConversationClient]: + # Pause the Live region for the duration of the conversation so + # ``prompt_toolkit`` owns the terminal; restore on exit. The + # client itself is shared with ``AutoProveConsoleHandler`` — + # same prompt-toolkit prompt, same progress-event drainer. + # + # The status provider closes over our live agent table so the + # prompt's bottom toolbar shows "Background: N agents running" + # while the user composes input. prompt-toolkit's + # ``refresh_interval`` keeps it ticking. + async with self._paused_live(): + client = ConsoleConversationClient( + initial + ) + async with client: + yield client + + def _background_status(self) -> str | None: + n = len(self._agents) + if n == 0: + return None + return f"Background: {n} agent{'' if n == 1 else 's'} running" + + # ── HandlerFactory ────────────────────────────────────────────── + + async def make_handler( + self, info: TaskInfo[AutoProvePhase] + ) -> TaskHandle[None]: + # ``run_task`` fires ``on_start`` / ``on_done`` per-task, not + # per-phase — a single phase like CVL_GEN spawns one + # "Invariant CVL" task plus one per-component batch, all + # sharing the same ``AutoProvePhase``. Use ``_phases_seen`` to + # collapse the per-task callback into a single phase header + # emitted the first time any task of that phase actually starts + # work. Subsequent tasks in the same phase don't repeat it. + # No ``on_done`` — top-level task completion is already + # captured by ``format_end``. + + def _on_start() -> None: + if info.phase in self._phases_seen: + return + self._phases_seen.add(info.phase) + label = AUTOPROVE_PHASE_LABELS.get(info.phase, info.phase.name) + self._scrollback( + Text( + "─" * 60 + f"\nPhase: {label}\n" + "─" * 60, + style="bold", + ) + ) + + async def _on_error(exc: Exception, tb: str) -> None: + self._scrollback( + Text( + f"[{info.label}] ERROR — {type(exc).__name__}: {exc}\n{tb}", + style="red", + ) + ) + + return TaskHandle( + handler=self, + event_handler=self, + on_start=_on_start, + on_error=_on_error, + conversation_provider=self._start_conversation, + ) diff --git a/composer/ui/conversation_client.py b/composer/ui/conversation_client.py index bab1d558..3862b4d6 100644 --- a/composer/ui/conversation_client.py +++ b/composer/ui/conversation_client.py @@ -102,4 +102,3 @@ async def __aexit__(self, exc_type, exc, tb): await self.drain_task except Exception: print("Conversation cleanup failed") - diff --git a/composer/ui/live_display.py b/composer/ui/live_display.py new file mode 100644 index 00000000..f3449c8b --- /dev/null +++ b/composer/ui/live_display.py @@ -0,0 +1,400 @@ +"""Common base for the codegen + autoprove live-display IOHandlers. + +Drives an inline ``rich.live.Live`` region whose contents are derived +from a per-thread ``AgentDisplayState`` map. Scrollback (the area +above the Live region) is append-only; ``log_start`` / ``log_end`` +always land there. The Live region itself shows a graceful-degradation +view of the active agent tree. + +Subclasses fill in: + +- ``handle_human_interaction`` — abstract. +- ``progress_update`` / ``handle_event`` — domain event ingest. Both + ultimately call ``_update_progress`` and ``_scrollback``. +- ``derive_status`` — usually inherited as-is; override only if a + workflow has its own descriptor logic that the default + composition doesn't capture. +- ``format_start`` / ``format_end`` / ``format_state_update`` / + ``handle_overflow`` — scrollback + overflow policy hooks. +""" + +import asyncio +from contextlib import asynccontextmanager +from dataclasses import dataclass, field +from typing import Any, AsyncIterator, Callable, Self + +from langchain_core.messages import AIMessage, ToolCall as LC_ToolCall, ToolMessage + +from rich.console import Console, RenderableType +from rich.live import Live +from rich.text import Text + +from composer.ui.tool_display import GroupedTool, ToolDisplayConfig + + +# ---------------- data model ---------------- + + +@dataclass +class ToolCall: + id: str + name: str + args: dict[str, Any] + + +@dataclass +class LiveToolGroup: + """Per-thread tool-call grouping state. + + Mirrors ``ToolGroupState`` from ``tool_call_renderer`` but persists + across tool *results* — only a non-matching tool call resets it. + The live region wants accumulated history across an iteration phase + (e.g. "Reading: a, b, c, d" across four rounds of get_file), which + the TUI's per-result reset would erase. + """ + group: GroupedTool | None = None + items: list[str] = field(default_factory=list) + + def reset(self) -> None: + self.group, self.items = None, [] + + +@dataclass +class AgentDisplayState: + thread_id: str + description: str + parent_id: str | None + spawn_tool_call_id: str | None + children: dict[str, "AgentDisplayState"] = field(default_factory=dict) + pending_tool_calls: dict[str, ToolCall] = field(default_factory=dict) + tool_group: LiveToolGroup = field(default_factory=LiveToolGroup) + progress_override: str | None = None + + +# ---------------- handler ---------------- + + +class LiveDisplayHandler[H]: + """IOHandler[H] base that drives an inline rich.live.Live region.""" + + def __init__( + self, + *, + tool_display_config: ToolDisplayConfig, + reserved_rows: int | None = None, + console: Console | None = None, + ) -> None: + self._tool_display_config = tool_display_config + self._console = console or Console() + self._reserved_rows = reserved_rows or max(6, self._console.size.height // 3) + + self._descriptions: dict[str, str] = {} + self._agents: dict[str, AgentDisplayState] = {} + self._roots: list[str] = [] + + self._live: Live | None = None + self._mutate_lock = asyncio.Lock() + self._suppress_scrollback = False + + # ---------- lifecycle ---------- + + async def __aenter__(self) -> Self: + self._live = Live( + self._render(), + console=self._console, + refresh_per_second=10, + transient=False, + ) + self._live.start() + return self + + async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + if self._live is not None: + self._live.stop() + self._live = None + + @asynccontextmanager + async def _paused_live(self) -> AsyncIterator[None]: + # Whoever runs inside _paused_live (HITL conversation, etc.) + # owns the terminal. Suppress scrollback so background graph + # activity in other threads doesn't scramble their rendering — + # rich and prompt-toolkit would otherwise race for stdout. + live = self._live + if live is not None: + live.stop() + prev_suppress = self._suppress_scrollback + self._suppress_scrollback = True + try: + yield + finally: + self._suppress_scrollback = prev_suppress + if live is not None: + live.start() + live.update(self._render()) + + # ---------- IOHandler[H] surface ---------- + + async def log_checkpoint_id(self, *, path: list[str], checkpoint_id: str) -> None: + pass + + async def log_start( + self, *, path: list[str], description: str, tool_id: str | None + ) -> None: + async with self._mutate_lock: + thread_id = path[-1] + parent_id = path[-2] if len(path) >= 2 else None + + self._descriptions[thread_id] = description + agent = AgentDisplayState( + thread_id=thread_id, + description=description, + parent_id=parent_id, + spawn_tool_call_id=tool_id, + ) + self._agents[thread_id] = agent + parent = self._agents.get(parent_id) if parent_id is not None else None + if parent is not None: + parent.children[thread_id] = agent + else: + self._roots.append(thread_id) + + self._scrollback(self.format_start(path, description, tool_id)) + self._refresh() + + async def log_end(self, path: list[str]) -> None: + async with self._mutate_lock: + thread_id = path[-1] + agent = self._agents.pop(thread_id, None) + if agent is None: + self._scrollback(self.format_end(path)) + return + + parent = ( + self._agents.get(agent.parent_id) + if agent.parent_id is not None + else None + ) + if parent is not None: + parent.children.pop(thread_id, None) + if agent.spawn_tool_call_id is not None: + parent.pending_tool_calls.pop(agent.spawn_tool_call_id, None) + elif thread_id in self._roots: + self._roots.remove(thread_id) + + self._scrollback(self.format_end(path)) + self._refresh() + + async def log_state_update(self, path: list[str], st: dict) -> None: + async with self._mutate_lock: + agent = self._agents.get(path[-1]) + if agent is not None: + self._ingest_state_update(agent, st) + + rendered = self.format_state_update(path, st) + if rendered is not None: + self._scrollback(rendered) + self._refresh() + + async def human_interaction( + self, + ty: H, + debug_thunk: Callable[[], None], + ) -> str: + async with self._paused_live(): + return await self.handle_human_interaction(ty, debug_thunk) + + # ---------- pluggable hooks ---------- + + async def handle_human_interaction( + self, + ty: H, + debug_thunk: Callable[[], None], + ) -> str: + raise NotImplementedError + + def derive_status(self, agent: AgentDisplayState) -> str: + if agent.progress_override is not None: + return agent.progress_override + + parts: list[str] = [] + regular_clause = self._render_regular_slots(agent) + if regular_clause: + parts.append(regular_clause) + if agent.children: + n = len(agent.children) + parts.append(f"waiting on {n} agent{'' if n == 1 else 's'}") + return "; ".join(parts) if parts else "thinking" + + def format_start( + self, path: list[str], description: str, tool_id: str | None + ) -> RenderableType | None: + # All transient activity is already visible in the Live region + # as the tree mounts a new node; a per-agent "start" line in + # scrollback just pushes interesting events off the top. + return None + + def format_end(self, path: list[str]) -> RenderableType | None: + # Only top-level completions go to scrollback. Sub-agents finish + # inside their parent's tree subtree — their disappearance from + # the Live region is the signal that they're done. + if len(path) > 1: + return None + return Text(f"[{self._label(path)}] ✓", style="green") + + def format_state_update( + self, path: list[str], st: dict + ) -> RenderableType | None: + return None + + def handle_overflow( + self, roots: list[AgentDisplayState], avail: int + ) -> RenderableType: + keep = max(1, avail - 1) + out = Text() + for r in roots[:keep]: + out.append(f"● {r.description} — {self.derive_status(r)}\n") + extra = len(roots) - keep + if extra > 0: + out.append(f"+{extra} more\n", style="dim") + return out + + # ---------- helpers for subclasses ---------- + + def _update_progress(self, thread_id: str, descriptor: str | None) -> None: + agent = self._agents.get(thread_id) + if agent is None: + return + agent.progress_override = descriptor + self._refresh() + + def _scrollback(self, renderable: RenderableType | None) -> None: + if renderable is None or self._suppress_scrollback: + return + if self._live is not None: + self._live.console.print(renderable) + else: + self._console.print(renderable) + + def _label(self, path: list[str]) -> str: + return " / ".join(self._descriptions.get(tid, tid) for tid in path) + + def _refresh(self) -> None: + if self._live is not None: + self._live.update(self._render()) + + # ---------- internal: state update ingest ---------- + + def _ingest_state_update(self, agent: AgentDisplayState, st: dict) -> None: + # ``run_graph._normalize_updates_payload`` guarantees each node's + # value is a dict by the time it reaches us, even when langgraph's + # ToolNode produced a list-of-Commands payload internally. + for update in st.values(): + if not isinstance(update, dict): + continue + for msg in update.get("messages", []): + if not isinstance(msg, AIMessage) and not isinstance(msg, ToolMessage): + continue + if isinstance(msg, ToolMessage): + agent.pending_tool_calls.pop(msg.tool_call_id, None) + continue + tc_list = msg.tool_calls + if tc_list: + # New AIMessage with tool calls — the prior tool's + # custom progress descriptor (e.g. "running prover", + # "analyzing CEX: ...") is stale by definition. + agent.progress_override = None + for tc in tc_list: + self._ingest_tool_call(agent, tc) + + def _ingest_tool_call(self, agent: AgentDisplayState, tc: LC_ToolCall) -> None: + name = tc["name"] + args = tc.get("args", {}) or {} + tcid = tc["id"] + + if tcid is None: + return + + agent.pending_tool_calls[tcid] = ToolCall(id=tcid, name=name, args=args) + + grouped = self._tool_display_config.get_group(name) + group = agent.tool_group + if grouped is None: + group.reset() + return + + raw = grouped.extract_group_items(args) + new_items = [raw] if isinstance(raw, str) else list(raw) + if group.group is not None and grouped.group_id == group.group.group_id: + group.items.extend(new_items) + else: + group.group = grouped + group.items = list(new_items) + + # ---------- internal: regular-slot rendering ---------- + + def _render_regular_slots(self, agent: AgentDisplayState) -> str: + spawn_ids = {c.spawn_tool_call_id for c in agent.children.values()} + parts: list[str] = [] + + group = agent.tool_group + if group.group is not None and group.items: + parts.append(group.group.render_group(group.items)) + + for tc in agent.pending_tool_calls.values(): + if tc.id in spawn_ids: + continue + if self._tool_display_config.get_group(tc.name) is not None: + continue + parts.append(self._tool_display_config.format_tool_call(tc.name, tc.args, short=True)) + + return ", ".join(parts) + + # ---------- internal: live region rendering ---------- + + def _root_list(self) -> list[AgentDisplayState]: + return [self._agents[t] for t in self._roots if t in self._agents] + + def _render(self) -> RenderableType: + roots = self._root_list() + if not roots: + return Text("(no active agents)\n", style="dim") + + avail = self._reserved_rows + if (expanded := self._render_expanded(roots, avail)) is not None: + return expanded + if (tight := self._render_tight(roots, avail)) is not None: + return tight + return self.handle_overflow(roots, avail) + + def _render_expanded( + self, roots: list[AgentDisplayState], budget: int + ) -> Text | None: + if self._expanded_line_count(roots) > budget: + return None + out = Text() + def walk(agent: AgentDisplayState, depth: int) -> None: + prefix = (" " * (depth - 1) + "└─ ") if depth else "" + out.append( + f"{prefix}● {agent.description} — {self.derive_status(agent)}\n" + ) + for child in agent.children.values(): + walk(child, depth + 1) + for r in roots: + walk(r, 0) + return out + + def _render_tight( + self, roots: list[AgentDisplayState], budget: int + ) -> Text | None: + if len(roots) > budget: + return None + out = Text() + for r in roots: + out.append(f"● {r.description} — {self.derive_status(r)}\n") + return out + + @staticmethod + def _expanded_line_count(roots: list[AgentDisplayState]) -> int: + def count(agent: AgentDisplayState) -> int: + return 1 + sum(count(c) for c in agent.children.values()) + return sum(count(r) for r in roots) diff --git a/composer/ui/tool_display.py b/composer/ui/tool_display.py index 13fe910e..11d632f9 100644 --- a/composer/ui/tool_display.py +++ b/composer/ui/tool_display.py @@ -15,8 +15,11 @@ class ToolDisplay: display_name: DisplayLabelTy """ - Label shown when the tool is called. ``str`` for a static name, callable - ``(input) -> str`` to vary based on the concrete arguments. + The "long" label shown when the tool is called — used when the renderer + has the budget for verbose detail (e.g. an expanded line in a transcript). + ``str`` for a static name, callable ``(input) -> str`` to vary based on + the concrete arguments. Free to inline argument fields verbatim; if the + renderer needs a length-bounded variant it should ask for ``short``. """ result: ResultOutputTy @@ -30,6 +33,17 @@ class ToolDisplay: to override both. """ + short_display_name: DisplayLabelTy | None = None + """ + Optional length-bounded variant of ``display_name``. Renderers that + need a compact line (status bar, collapsed group header, single-line + progress indicator) request this via ``format_tool_call(..., short=True)``. + When ``None`` (the default), ``format_tool_call`` falls back to + ``display_name`` — preserving prior behavior. Provide this whenever + ``display_name`` inlines free-text argument fields (research questions, + explanations, reasons) that can blow up a single line. + """ + @dataclass class GroupedTool: @@ -123,7 +137,11 @@ class CommonTools: ) result = ToolDisplay("Delivering result", suppress_ack("Result")) - code_explorer = ToolDisplay(lambda q: f"Code Exploration Request: {q["question"]}", "Code Explorer Answer") + code_explorer = ToolDisplay( + lambda q: f"Code Exploration Request: {q["question"]}", + "Code Explorer Answer", + short_display_name="Code exploration", + ) get_file = GroupedTool( "read", @@ -150,6 +168,7 @@ class CommonTools: ) cvl_research = ToolDisplay( lambda p: f"Researching CVL: {p.get('question', '?')}", "Research result", + short_display_name="CVL research", ) scan_knowledge_base = ToolDisplay("Scanning knowledge base", "KB scan results") get_knowledge_base_article = ToolDisplay("Reading KB article", "KB article") @@ -254,12 +273,18 @@ def _find_formatter(self, name) -> ToolDisplay | GroupedTool | None: return _graphcore_global_tools[name] return None - def format_tool_call(self, name: str, input: dict) -> str: - """Return a user-friendly label for a tool invocation.""" + def format_tool_call(self, name: str, input: dict, *, short: bool = False) -> str: + """Return a user-friendly label for a tool invocation. + + When ``short=True``, the entry's ``short_display_name`` is used if + provided, otherwise we fall back to ``display_name``. The decision + of which variant to use lives with the renderer — annotations + declare both shapes and let the caller pick. + """ entry = self._find_formatter(name) if entry is None or isinstance(entry, GroupedTool): return f"Tool: {name}" - nm = entry.display_name + nm = entry.short_display_name if short and entry.short_display_name is not None else entry.display_name if isinstance(nm, str): return nm return nm(input) @@ -397,6 +422,13 @@ def new_as_tool( def tool_display( label: DisplayLabelTy, - result: ResultOutputTy + result: ResultOutputTy, + short_label: DisplayLabelTy | None = None, ) -> Callable[[T_VAR], T_VAR]: - return tool_display_of(ToolDisplay(display_name=label, result=result)) + return tool_display_of( + ToolDisplay( + display_name=label, + result=result, + short_display_name=short_label, + ) + ) From 1a964309c29c30030b5af527f093abf8fe986c40 Mon Sep 17 00:00:00 2001 From: John Toman Date: Tue, 16 Jun 2026 12:42:31 -0700 Subject: [PATCH 2/2] better typing --- composer/io/graph_runner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/composer/io/graph_runner.py b/composer/io/graph_runner.py index 46539566..80722e8a 100644 --- a/composer/io/graph_runner.py +++ b/composer/io/graph_runner.py @@ -24,7 +24,7 @@ from langchain_core.runnables import RunnableConfig -def _normalize_updates_payload(payload: dict) -> dict: +def _normalize_updates_payload(payload: dict[str, Any]) -> dict: """Flatten langgraph's mixed-mode tools-node payload to ``dict[node, dict]``. When a ``ToolNode`` batch contains tools whose returns differ in shape