From 3afdd98bd159b512d6aec0cbf399fee264b09e50 Mon Sep 17 00:00:00 2001 From: Ali Khokhar Date: Sun, 28 Jun 2026 00:36:36 -0700 Subject: [PATCH] Refactor messaging transcript into package (#930) --- ARCHITECTURE.md | 7 + messaging/transcript.py | 581 ---------------------- messaging/transcript/__init__.py | 6 + messaging/transcript/buffer.py | 222 +++++++++ messaging/transcript/context.py | 20 + messaging/transcript/renderer.py | 67 +++ messaging/transcript/segments.py | 178 +++++++ messaging/transcript/subagents.py | 110 ++++ pyproject.toml | 2 +- tests/contracts/test_import_boundaries.py | 21 + tests/messaging/test_transcript.py | 76 +-- uv.lock | 2 +- 12 files changed, 680 insertions(+), 612 deletions(-) delete mode 100644 messaging/transcript.py create mode 100644 messaging/transcript/__init__.py create mode 100644 messaging/transcript/buffer.py create mode 100644 messaging/transcript/context.py create mode 100644 messaging/transcript/renderer.py create mode 100644 messaging/transcript/segments.py create mode 100644 messaging/transcript/subagents.py diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 3043338e..f8215f5b 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -565,6 +565,13 @@ lifecycle for queued nodes: parent-session fork/resume, session registration, CLI event parsing, transcript/status updates, cancellation, error propagation, and session cleanup. +[messaging/event_parser.py](messaging/event_parser.py) normalizes managed Claude +JSON events into low-level transcript events. +[messaging/transcript/](messaging/transcript/) owns transcript assembly and +rendering: open content-block tracking, Task/subagent display state, segment +models, render context, and truncation. Platform markdown details stay in +[messaging/rendering/](messaging/rendering/). + [messaging/command_context.py](messaging/command_context.py) defines the typed dependency surface for `/stop`, `/clear`, and `/stats`; commands should not depend on the concrete workflow object or on platform SDK runtimes. diff --git a/messaging/transcript.py b/messaging/transcript.py deleted file mode 100644 index f1a015d0..00000000 --- a/messaging/transcript.py +++ /dev/null @@ -1,581 +0,0 @@ -"""Ordered transcript builder for messaging UIs (Telegram, etc.). - -This module maintains an ordered list of "segments" that represent what the user -should see in the chat transcript: thinking, tool calls, tool results, subagent -headers, and assistant text. It is designed for in-place message editing where -the transcript grows over time and older content must be truncated. -""" - -from __future__ import annotations - -import json -from abc import ABC, abstractmethod -from collections import deque -from collections.abc import Callable, Iterable -from dataclasses import dataclass, field -from typing import Any - -from loguru import logger - - -def _safe_json_dumps(obj: Any) -> str: - try: - return json.dumps(obj, indent=2, ensure_ascii=False, sort_keys=True) - except Exception: - return str(obj) - - -@dataclass -class Segment(ABC): - kind: str - - @abstractmethod - def render(self, ctx: RenderCtx) -> str: ... - - -@dataclass -class ThinkingSegment(Segment): - def __init__(self) -> None: - super().__init__(kind="thinking") - self._parts: list[str] = [] - - def append(self, t: str) -> None: - if t: - self._parts.append(t) - - @property - def text(self) -> str: - return "".join(self._parts) - - def render(self, ctx: RenderCtx) -> str: - raw = self.text or "" - if ctx.thinking_tail_max is not None and len(raw) > ctx.thinking_tail_max: - raw = "..." + raw[-(ctx.thinking_tail_max - 3) :] - inner = ctx.escape_code(raw) - return f"💭 {ctx.bold('Thinking')}\n```\n{inner}\n```" - - -@dataclass -class TextSegment(Segment): - def __init__(self) -> None: - super().__init__(kind="text") - self._parts: list[str] = [] - - def append(self, t: str) -> None: - if t: - self._parts.append(t) - - @property - def text(self) -> str: - return "".join(self._parts) - - def render(self, ctx: RenderCtx) -> str: - raw = self.text or "" - if ctx.text_tail_max is not None and len(raw) > ctx.text_tail_max: - raw = "..." + raw[-(ctx.text_tail_max - 3) :] - return ctx.render_markdown(raw) - - -@dataclass -class ToolCallSegment(Segment): - tool_use_id: str - name: str - closed: bool = False - indent_level: int = 0 - - def __init__(self, tool_use_id: str, name: str, *, indent_level: int = 0) -> None: - super().__init__(kind="tool_call") - self.tool_use_id = str(tool_use_id or "") - self.name = str(name or "tool") - self.indent_level = max(0, int(indent_level)) - - def render(self, ctx: RenderCtx) -> str: - name = ctx.code_inline(self.name) - # Per UX requirement: do not display tool args/results, only the tool call. - prefix = " " * self.indent_level - return f"{prefix}🛠 {ctx.bold('Tool call:')} {name}" - - -@dataclass -class ToolResultSegment(Segment): - tool_use_id: str - name: str | None - content_text: str - is_error: bool = False - - def __init__( - self, - tool_use_id: str, - content: Any, - *, - name: str | None = None, - is_error: bool = False, - ) -> None: - super().__init__(kind="tool_result") - self.tool_use_id = str(tool_use_id or "") - self.name = str(name) if name is not None else None - self.is_error = bool(is_error) - if isinstance(content, str): - self.content_text = content - else: - self.content_text = _safe_json_dumps(content) - - def render(self, ctx: RenderCtx) -> str: - raw = self.content_text or "" - if ctx.tool_output_tail_max is not None and len(raw) > ctx.tool_output_tail_max: - raw = "..." + raw[-(ctx.tool_output_tail_max - 3) :] - inner = ctx.escape_code(raw) - label = "Tool error:" if self.is_error else "Tool result:" - maybe_name = f" {ctx.code_inline(self.name)}" if self.name else "" - return f"📤 {ctx.bold(label)}{maybe_name}\n```\n{inner}\n```" - - -@dataclass -class SubagentSegment(Segment): - description: str - tool_calls: int = 0 - tools_used: set[str] = field(default_factory=set) - current_tool: ToolCallSegment | None = None - - def __init__(self, description: str) -> None: - super().__init__(kind="subagent") - self.description = str(description or "Subagent") - self.tool_calls = 0 - self.tools_used = set() - self.current_tool = None - - def set_current_tool_call(self, tool_use_id: str, name: str) -> ToolCallSegment: - tool_use_id = str(tool_use_id or "") - name = str(name or "tool") - self.tools_used.add(name) - self.tool_calls += 1 - self.current_tool = ToolCallSegment(tool_use_id, name, indent_level=1) - return self.current_tool - - def render(self, ctx: RenderCtx) -> str: - inner_prefix = " " - - lines: list[str] = [ - f"🤖 {ctx.bold('Subagent:')} {ctx.code_inline(self.description)}" - ] - - if self.current_tool is not None: - try: - rendered = self.current_tool.render(ctx) - except Exception: - rendered = "" - if rendered: - lines.append(rendered) - - tools_used = sorted(self.tools_used) - tools_set_raw = "{{{}}}".format(", ".join(tools_used)) if tools_used else "{}" - - # Keep braces inside a code entity so MarkdownV2 doesn't require escaping them. - lines.append( - f"{inner_prefix}{ctx.bold('Tools used:')} {ctx.code_inline(tools_set_raw)}" - ) - lines.append( - f"{inner_prefix}{ctx.bold('Tool calls:')} {ctx.code_inline(str(self.tool_calls))}" - ) - return "\n".join(lines) - - -@dataclass -class ErrorSegment(Segment): - message: str - - def __init__(self, message: str) -> None: - super().__init__(kind="error") - self.message = str(message or "Unknown error") - - def render(self, ctx: RenderCtx) -> str: - return f"⚠️ {ctx.bold('Error:')} {ctx.code_inline(self.message)}" - - -@dataclass -class RenderCtx: - bold: Callable[[str], str] - code_inline: Callable[[str], str] - escape_code: Callable[[str], str] - escape_text: Callable[[str], str] - render_markdown: Callable[[str], str] - - thinking_tail_max: int | None = 1000 - tool_input_tail_max: int | None = 1200 - tool_output_tail_max: int | None = 1600 - text_tail_max: int | None = 2000 - - -class TranscriptBuffer: - """Maintains an ordered, truncatable transcript of events.""" - - def __init__( - self, - *, - show_tool_results: bool = True, - debug_subagent_stack: bool = False, - ) -> None: - self._segments: list[Segment] = [] - self._open_thinking_by_index: dict[int, ThinkingSegment] = {} - self._open_text_by_index: dict[int, TextSegment] = {} - - # content_block index -> tool call segment (for streaming tool args) - self._open_tools_by_index: dict[int, ToolCallSegment] = {} - - # tool_use_id -> tool name (for tool_result labeling) - self._tool_name_by_id: dict[str, str] = {} - - self._show_tool_results = bool(show_tool_results) - - # subagent context stack. Each entry is the Task tool_use_id we are waiting to close. - self._subagent_stack: list[str] = [] - # Parallel stack of segments for rendering nested subagents. - self._subagent_segments: list[SubagentSegment] = [] - self._debug_subagent_stack = debug_subagent_stack - - def _in_subagent(self) -> bool: - return bool(self._subagent_stack) - - def _subagent_current(self) -> SubagentSegment | None: - return self._subagent_segments[-1] if self._subagent_segments else None - - def _task_heading_from_input(self, inp: Any) -> str: - # We never display full JSON args; only extract a short heading. - if isinstance(inp, dict): - desc = str(inp.get("description", "") or "").strip() - if desc: - return desc - subagent_type = str(inp.get("subagent_type", "") or "").strip() - if subagent_type: - return subagent_type - typ = str(inp.get("type", "") or "").strip() - if typ: - return typ - return "Subagent" - - def _subagent_push(self, tool_id: str, seg: SubagentSegment) -> None: - # Some providers can omit ids; still track depth for UI suppression. - tool_id = ( - str(tool_id or "").strip() or f"__task_{len(self._subagent_stack) + 1}" - ) - self._subagent_stack.append(tool_id) - self._subagent_segments.append(seg) - if self._debug_subagent_stack: - logger.debug( - "SUBAGENT_STACK: push id=%r depth=%d heading=%r", - tool_id, - len(self._subagent_stack), - getattr(seg, "description", None), - ) - - def _subagent_pop(self, tool_id: str) -> bool: - tool_id = str(tool_id or "").strip() - if not self._subagent_stack: - return False - - def _ids_roughly_match(stack_id: str, result_id: str) -> bool: - if not stack_id or not result_id: - return False - if stack_id == result_id: - return True - # Some providers emit Task result ids with a suffix/prefix variant. - # Treat those as the same logical Task invocation. - return stack_id.startswith(result_id) or result_id.startswith(stack_id) - - if tool_id: - # O(1) common case: LIFO - top of stack matches. - if _ids_roughly_match(self._subagent_stack[-1], tool_id): - self._subagent_stack.pop() - if self._subagent_segments: - self._subagent_segments.pop() - if self._debug_subagent_stack: - logger.debug( - "SUBAGENT_STACK: pop id=%r depth=%d (LIFO)", - tool_id, - len(self._subagent_stack), - ) - return True - # Pop to the matching id (defensive against non-LIFO emissions). - idx = -1 - for i in range(len(self._subagent_stack) - 1, -1, -1): - if _ids_roughly_match(self._subagent_stack[i], tool_id): - idx = i - break - if idx < 0: - return False - while len(self._subagent_stack) > idx: - popped = self._subagent_stack.pop() - if self._subagent_segments: - self._subagent_segments.pop() - if self._debug_subagent_stack: - logger.debug( - "SUBAGENT_STACK: pop id=%r depth=%d (matched=%r)", - popped, - len(self._subagent_stack), - tool_id, - ) - return True - - # No id in result; only close if we have a synthetic top marker. - if self._subagent_stack and self._subagent_stack[-1].startswith("__task_"): - popped = self._subagent_stack.pop() - if self._subagent_segments: - self._subagent_segments.pop() - if self._debug_subagent_stack: - logger.debug( - "SUBAGENT_STACK: pop id=%r depth=%d (synthetic)", - popped, - len(self._subagent_stack), - ) - return True - return False - - def _ensure_thinking(self) -> ThinkingSegment: - seg = ThinkingSegment() - self._segments.append(seg) - return seg - - def _ensure_text(self) -> TextSegment: - seg = TextSegment() - self._segments.append(seg) - return seg - - def apply(self, ev: dict[str, Any]) -> None: - """Apply a parsed event to the transcript.""" - et = ev.get("type") - - # Subagent rules: inside a Task/subagent, we only show tool calls/results. - if self._in_subagent() and et in ( - "thinking_start", - "thinking_delta", - "thinking_chunk", - "text_start", - "text_delta", - "text_chunk", - ): - return - - if et == "thinking_start": - idx = int(ev.get("index", -1)) - if idx >= 0: - # Defensive: if a provider reuses indices without emitting a stop, - # close the previous open segment first. - self.apply({"type": "block_stop", "index": idx}) - seg = self._ensure_thinking() - if idx >= 0: - self._open_thinking_by_index[idx] = seg - return - if et in ("thinking_delta", "thinking_chunk"): - idx = int(ev.get("index", -1)) - seg = self._open_thinking_by_index.get(idx) - if seg is None: - seg = self._ensure_thinking() - if idx >= 0: - self._open_thinking_by_index[idx] = seg - seg.append(str(ev.get("text", ""))) - return - if et == "thinking_stop": - idx = int(ev.get("index", -1)) - if idx >= 0: - self._open_thinking_by_index.pop(idx, None) - return - - if et == "text_start": - idx = int(ev.get("index", -1)) - if idx >= 0: - self.apply({"type": "block_stop", "index": idx}) - seg = self._ensure_text() - if idx >= 0: - self._open_text_by_index[idx] = seg - return - if et in ("text_delta", "text_chunk"): - idx = int(ev.get("index", -1)) - seg = self._open_text_by_index.get(idx) - if seg is None: - seg = self._ensure_text() - if idx >= 0: - self._open_text_by_index[idx] = seg - seg.append(str(ev.get("text", ""))) - return - if et == "text_stop": - idx = int(ev.get("index", -1)) - if idx >= 0: - self._open_text_by_index.pop(idx, None) - return - - if et == "tool_use_start": - idx = int(ev.get("index", -1)) - if idx >= 0: - self.apply({"type": "block_stop", "index": idx}) - tool_id = str(ev.get("id", "") or "").strip() - name = str(ev.get("name", "") or "tool") - if tool_id: - self._tool_name_by_id[tool_id] = name - - # Task tool indicates subagent. - if name == "Task": - heading = self._task_heading_from_input(ev.get("input")) - seg = SubagentSegment(heading) - self._segments.append(seg) - self._subagent_push(tool_id, seg) - return - - # Normal tool call. - if self._in_subagent(): - parent = self._subagent_current() - if parent is not None: - seg = parent.set_current_tool_call(tool_id, name) - else: - seg = ToolCallSegment(tool_id, name) - self._segments.append(seg) - else: - seg = ToolCallSegment(tool_id, name) - self._segments.append(seg) - - if idx >= 0: - self._open_tools_by_index[idx] = seg - return - - if et == "tool_use_delta": - # Track open tool by index for tool_use_stop (closing state). - return - - if et == "tool_use_stop": - idx = int(ev.get("index", -1)) - seg = self._open_tools_by_index.pop(idx, None) - if seg is not None: - seg.closed = True - return - - if et == "block_stop": - idx = int(ev.get("index", -1)) - if idx in self._open_tools_by_index: - self.apply({"type": "tool_use_stop", "index": idx}) - return - if idx in self._open_thinking_by_index: - self.apply({"type": "thinking_stop", "index": idx}) - return - if idx in self._open_text_by_index: - self.apply({"type": "text_stop", "index": idx}) - return - return - - if et == "tool_use": - tool_id = str(ev.get("id", "") or "").strip() - name = str(ev.get("name", "") or "tool") - if tool_id: - self._tool_name_by_id[tool_id] = name - - if name == "Task": - heading = self._task_heading_from_input(ev.get("input")) - seg = SubagentSegment(heading) - self._segments.append(seg) - self._subagent_push(tool_id, seg) - return - - if self._in_subagent(): - parent = self._subagent_current() - if parent is not None: - seg = parent.set_current_tool_call(tool_id, name) - else: - seg = ToolCallSegment(tool_id, name) - self._segments.append(seg) - else: - seg = ToolCallSegment(tool_id, name) - self._segments.append(seg) - - seg.closed = True - return - - if et == "tool_result": - tool_id = str(ev.get("tool_use_id", "") or "").strip() - name = self._tool_name_by_id.get(tool_id) - - # If this was the Task tool result, close subagent context. - if self._subagent_stack: - popped = self._subagent_pop(tool_id) - top = self._subagent_stack[-1] if self._subagent_stack else "" - looks_like_task_id = "task" in tool_id.lower() - # Some streams omit Task tool_use ids (synthetic stack ids), but include - # a real Task id on tool_result (e.g. "functions.Task:0"). Reconcile that. - if ( - not popped - and tool_id - and top.startswith("__task_") - and (name in (None, "Task")) - and looks_like_task_id - ): - self._subagent_pop("") - - if not self._show_tool_results: - return - - seg = ToolResultSegment( - tool_id, - ev.get("content"), - name=name, - is_error=bool(ev.get("is_error", False)), - ) - self._segments.append(seg) - return - - if et == "error": - self._segments.append(ErrorSegment(str(ev.get("message", "")))) - return - - def render(self, ctx: RenderCtx, *, limit_chars: int, status: str | None) -> str: - """Render transcript with truncation (drop oldest segments).""" - # Filter out empty rendered segments. - rendered: list[str] = [] - for seg in self._segments: - try: - out = seg.render(ctx) - except Exception: - continue - if out: - rendered.append(out) - - status_text = f"\n\n{status}" if status else "" - prefix_marker = ctx.escape_text("... (truncated)\n") - - def _join(parts: Iterable[str], add_marker: bool) -> str: - body = "\n".join(parts) - if add_marker and body: - body = prefix_marker + body - return body + status_text if (body or status_text) else status_text - - # Fast path. - candidate = _join(rendered, add_marker=False) - if len(candidate) <= limit_chars: - return candidate - - # Drop oldest segments until under limit (keep the tail). - # Use deque for O(1) popleft; list.pop(0) would be O(n) per iteration. - parts: deque[str] = deque(rendered) - dropped = False - last_part: str | None = None - while parts: - candidate = _join(parts, add_marker=True) - if len(candidate) <= limit_chars: - return candidate - last_part = parts.popleft() - dropped = True - - # Nothing fits - preserve tail of last segment instead of only marker+status. - if dropped and last_part: - budget = limit_chars - len(prefix_marker) - len(status_text) - if budget > 20: - if len(last_part) > budget: - tail = "..." + last_part[-(budget - 3) :] - else: - tail = last_part - candidate = prefix_marker + tail + status_text - if len(candidate) <= limit_chars: - return candidate - - # Fallback: marker + status only. - if dropped: - minimal = prefix_marker + status_text.lstrip("\n") - if len(minimal) <= limit_chars: - return minimal - return status or "" diff --git a/messaging/transcript/__init__.py b/messaging/transcript/__init__.py new file mode 100644 index 00000000..881befcd --- /dev/null +++ b/messaging/transcript/__init__.py @@ -0,0 +1,6 @@ +"""Public transcript API for messaging UI rendering.""" + +from .buffer import TranscriptBuffer +from .context import RenderCtx + +__all__ = ["RenderCtx", "TranscriptBuffer"] diff --git a/messaging/transcript/buffer.py b/messaging/transcript/buffer.py new file mode 100644 index 00000000..6af43b7f --- /dev/null +++ b/messaging/transcript/buffer.py @@ -0,0 +1,222 @@ +"""Transcript event application and open-block tracking.""" + +from __future__ import annotations + +from typing import Any + +from .context import RenderCtx +from .renderer import render_segments +from .segments import ( + ErrorSegment, + Segment, + SubagentSegment, + TextSegment, + ThinkingSegment, + ToolCallSegment, + ToolResultSegment, +) +from .subagents import SubagentState, task_heading_from_input + +_SUBAGENT_SUPPRESSED_EVENTS = frozenset( + { + "thinking_start", + "thinking_delta", + "thinking_chunk", + "text_start", + "text_delta", + "text_chunk", + } +) + + +class TranscriptBuffer: + """Maintains an ordered, truncatable transcript of parsed CLI events.""" + + def __init__( + self, + *, + show_tool_results: bool = True, + debug_subagent_stack: bool = False, + ) -> None: + self._segments: list[Segment] = [] + self._open_thinking_by_index: dict[int, ThinkingSegment] = {} + self._open_text_by_index: dict[int, TextSegment] = {} + self._open_tools_by_index: dict[int, ToolCallSegment] = {} + self._tool_name_by_id: dict[str, str] = {} + self._show_tool_results = bool(show_tool_results) + self._subagents = SubagentState(debug=debug_subagent_stack) + + def apply(self, event: dict[str, Any]) -> None: + """Apply a parsed CLI transcript event.""" + event_type = event.get("type") + if self._subagents.in_subagent() and event_type in _SUBAGENT_SUPPRESSED_EVENTS: + return + + if event_type == "thinking_start": + self._start_thinking(_event_index(event)) + return + if event_type in ("thinking_delta", "thinking_chunk"): + self._append_thinking(_event_index(event), str(event.get("text", ""))) + return + if event_type == "thinking_stop": + self._open_thinking_by_index.pop(_event_index(event), None) + return + + if event_type == "text_start": + self._start_text(_event_index(event)) + return + if event_type in ("text_delta", "text_chunk"): + self._append_text(_event_index(event), str(event.get("text", ""))) + return + if event_type == "text_stop": + self._open_text_by_index.pop(_event_index(event), None) + return + + if event_type == "tool_use_start": + self._start_tool_use(event) + return + if event_type == "tool_use_delta": + return + if event_type == "tool_use_stop": + segment = self._open_tools_by_index.pop(_event_index(event), None) + if segment is not None: + segment.closed = True + return + + if event_type == "block_stop": + self._close_block(_event_index(event)) + return + if event_type == "tool_use": + self._append_complete_tool_use(event) + return + if event_type == "tool_result": + self._append_tool_result(event) + return + if event_type == "error": + self._segments.append(ErrorSegment(str(event.get("message", "")))) + + def render(self, ctx: RenderCtx, *, limit_chars: int, status: str | None) -> str: + return render_segments( + self._segments, + ctx, + limit_chars=limit_chars, + status=status, + ) + + def _start_thinking(self, index: int) -> None: + if index >= 0: + self._close_block(index) + segment = ThinkingSegment() + self._segments.append(segment) + if index >= 0: + self._open_thinking_by_index[index] = segment + + def _append_thinking(self, index: int, text: str) -> None: + segment = self._open_thinking_by_index.get(index) + if segment is None: + segment = ThinkingSegment() + self._segments.append(segment) + if index >= 0: + self._open_thinking_by_index[index] = segment + segment.append(text) + + def _start_text(self, index: int) -> None: + if index >= 0: + self._close_block(index) + segment = TextSegment() + self._segments.append(segment) + if index >= 0: + self._open_text_by_index[index] = segment + + def _append_text(self, index: int, text: str) -> None: + segment = self._open_text_by_index.get(index) + if segment is None: + segment = TextSegment() + self._segments.append(segment) + if index >= 0: + self._open_text_by_index[index] = segment + segment.append(text) + + def _start_tool_use(self, event: dict[str, Any]) -> None: + index = _event_index(event) + if index >= 0: + self._close_block(index) + + tool_id = _event_tool_id(event, "id") + name = str(event.get("name", "") or "tool") + if tool_id: + self._tool_name_by_id[tool_id] = name + + if name == "Task": + segment = SubagentSegment(task_heading_from_input(event.get("input"))) + self._segments.append(segment) + self._subagents.push(tool_id, segment) + return + + segment = self._append_tool_call(tool_id, name) + if index >= 0: + self._open_tools_by_index[index] = segment + + def _append_complete_tool_use(self, event: dict[str, Any]) -> None: + tool_id = _event_tool_id(event, "id") + name = str(event.get("name", "") or "tool") + if tool_id: + self._tool_name_by_id[tool_id] = name + + if name == "Task": + segment = SubagentSegment(task_heading_from_input(event.get("input"))) + self._segments.append(segment) + self._subagents.push(tool_id, segment) + return + + segment = self._append_tool_call(tool_id, name) + segment.closed = True + + def _append_tool_call(self, tool_id: str, name: str) -> ToolCallSegment: + if self._subagents.in_subagent(): + parent = self._subagents.current_segment() + if parent is not None: + return parent.set_current_tool_call(tool_id, name) + + segment = ToolCallSegment(tool_id, name) + self._segments.append(segment) + return segment + + def _append_tool_result(self, event: dict[str, Any]) -> None: + tool_id = _event_tool_id(event, "tool_use_id") + name = self._tool_name_by_id.get(tool_id) + + if self._subagents.in_subagent(): + self._subagents.close_for_tool_result(tool_id, tool_name=name) + + if not self._show_tool_results: + return + + self._segments.append( + ToolResultSegment( + tool_id, + event.get("content"), + name=name, + is_error=bool(event.get("is_error", False)), + ) + ) + + def _close_block(self, index: int) -> None: + if index in self._open_tools_by_index: + segment = self._open_tools_by_index.pop(index, None) + if segment is not None: + segment.closed = True + return + if index in self._open_thinking_by_index: + self._open_thinking_by_index.pop(index, None) + return + if index in self._open_text_by_index: + self._open_text_by_index.pop(index, None) + + +def _event_index(event: dict[str, Any]) -> int: + return int(event.get("index", -1)) + + +def _event_tool_id(event: dict[str, Any], key: str) -> str: + return str(event.get(key, "") or "").strip() diff --git a/messaging/transcript/context.py b/messaging/transcript/context.py new file mode 100644 index 00000000..3e1f76f3 --- /dev/null +++ b/messaging/transcript/context.py @@ -0,0 +1,20 @@ +"""Rendering context used by transcript segments.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + + +@dataclass +class RenderCtx: + bold: Callable[[str], str] + code_inline: Callable[[str], str] + escape_code: Callable[[str], str] + escape_text: Callable[[str], str] + render_markdown: Callable[[str], str] + + thinking_tail_max: int | None = 1000 + tool_input_tail_max: int | None = 1200 + tool_output_tail_max: int | None = 1600 + text_tail_max: int | None = 2000 diff --git a/messaging/transcript/renderer.py b/messaging/transcript/renderer.py new file mode 100644 index 00000000..b4b35f76 --- /dev/null +++ b/messaging/transcript/renderer.py @@ -0,0 +1,67 @@ +"""Render and truncate ordered transcript segments.""" + +from __future__ import annotations + +from collections import deque +from collections.abc import Iterable + +from .context import RenderCtx +from .segments import Segment + + +def render_segments( + segments: Iterable[Segment], + ctx: RenderCtx, + *, + limit_chars: int, + status: str | None, +) -> str: + rendered: list[str] = [] + for segment in segments: + try: + output = segment.render(ctx) + except Exception: + continue + if output: + rendered.append(output) + + status_text = f"\n\n{status}" if status else "" + prefix_marker = ctx.escape_text("... (truncated)\n") + + def _join(parts: Iterable[str], add_marker: bool) -> str: + body = "\n".join(parts) + if add_marker and body: + body = prefix_marker + body + return body + status_text if (body or status_text) else status_text + + candidate = _join(rendered, add_marker=False) + if len(candidate) <= limit_chars: + return candidate + + parts: deque[str] = deque(rendered) + dropped = False + last_part: str | None = None + while parts: + candidate = _join(parts, add_marker=True) + if len(candidate) <= limit_chars: + return candidate + last_part = parts.popleft() + dropped = True + + if dropped and last_part: + budget = limit_chars - len(prefix_marker) - len(status_text) + if budget > 20: + tail = ( + "..." + last_part[-(budget - 3) :] + if len(last_part) > budget + else last_part + ) + candidate = prefix_marker + tail + status_text + if len(candidate) <= limit_chars: + return candidate + + if dropped: + minimal = prefix_marker + status_text.lstrip("\n") + if len(minimal) <= limit_chars: + return minimal + return status or "" diff --git a/messaging/transcript/segments.py b/messaging/transcript/segments.py new file mode 100644 index 00000000..cf3d62a7 --- /dev/null +++ b/messaging/transcript/segments.py @@ -0,0 +1,178 @@ +"""Transcript segment types for messaging UI output.""" + +from __future__ import annotations + +import json +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Any + +from .context import RenderCtx + + +def safe_json_dumps(obj: Any) -> str: + try: + return json.dumps(obj, indent=2, ensure_ascii=False, sort_keys=True) + except Exception: + return str(obj) + + +@dataclass +class Segment(ABC): + kind: str + + @abstractmethod + def render(self, ctx: RenderCtx) -> str: ... + + +@dataclass +class ThinkingSegment(Segment): + def __init__(self) -> None: + super().__init__(kind="thinking") + self._parts: list[str] = [] + + def append(self, text: str) -> None: + if text: + self._parts.append(text) + + @property + def text(self) -> str: + return "".join(self._parts) + + def render(self, ctx: RenderCtx) -> str: + raw = self.text or "" + if ctx.thinking_tail_max is not None and len(raw) > ctx.thinking_tail_max: + raw = "..." + raw[-(ctx.thinking_tail_max - 3) :] + inner = ctx.escape_code(raw) + return f"💭 {ctx.bold('Thinking')}\n```\n{inner}\n```" + + +@dataclass +class TextSegment(Segment): + def __init__(self) -> None: + super().__init__(kind="text") + self._parts: list[str] = [] + + def append(self, text: str) -> None: + if text: + self._parts.append(text) + + @property + def text(self) -> str: + return "".join(self._parts) + + def render(self, ctx: RenderCtx) -> str: + raw = self.text or "" + if ctx.text_tail_max is not None and len(raw) > ctx.text_tail_max: + raw = "..." + raw[-(ctx.text_tail_max - 3) :] + return ctx.render_markdown(raw) + + +@dataclass +class ToolCallSegment(Segment): + tool_use_id: str + name: str + closed: bool = False + indent_level: int = 0 + + def __init__(self, tool_use_id: str, name: str, *, indent_level: int = 0) -> None: + super().__init__(kind="tool_call") + self.tool_use_id = str(tool_use_id or "") + self.name = str(name or "tool") + self.closed = False + self.indent_level = max(0, int(indent_level)) + + def render(self, ctx: RenderCtx) -> str: + name = ctx.code_inline(self.name) + prefix = " " * self.indent_level + return f"{prefix}🛠 {ctx.bold('Tool call:')} {name}" + + +@dataclass +class ToolResultSegment(Segment): + tool_use_id: str + name: str | None + content_text: str + is_error: bool = False + + def __init__( + self, + tool_use_id: str, + content: Any, + *, + name: str | None = None, + is_error: bool = False, + ) -> None: + super().__init__(kind="tool_result") + self.tool_use_id = str(tool_use_id or "") + self.name = str(name) if name is not None else None + self.is_error = bool(is_error) + self.content_text = ( + content if isinstance(content, str) else safe_json_dumps(content) + ) + + def render(self, ctx: RenderCtx) -> str: + raw = self.content_text or "" + if ctx.tool_output_tail_max is not None and len(raw) > ctx.tool_output_tail_max: + raw = "..." + raw[-(ctx.tool_output_tail_max - 3) :] + inner = ctx.escape_code(raw) + label = "Tool error:" if self.is_error else "Tool result:" + maybe_name = f" {ctx.code_inline(self.name)}" if self.name else "" + return f"📤 {ctx.bold(label)}{maybe_name}\n```\n{inner}\n```" + + +@dataclass +class SubagentSegment(Segment): + description: str + tool_calls: int = 0 + tools_used: set[str] = field(default_factory=set) + current_tool: ToolCallSegment | None = None + + def __init__(self, description: str) -> None: + super().__init__(kind="subagent") + self.description = str(description or "Subagent") + self.tool_calls = 0 + self.tools_used = set() + self.current_tool = None + + def set_current_tool_call(self, tool_use_id: str, name: str) -> ToolCallSegment: + tool_use_id = str(tool_use_id or "") + name = str(name or "tool") + self.tools_used.add(name) + self.tool_calls += 1 + self.current_tool = ToolCallSegment(tool_use_id, name, indent_level=1) + return self.current_tool + + def render(self, ctx: RenderCtx) -> str: + inner_prefix = " " + lines = [f"🤖 {ctx.bold('Subagent:')} {ctx.code_inline(self.description)}"] + + if self.current_tool is not None: + try: + rendered = self.current_tool.render(ctx) + except Exception: + rendered = "" + if rendered: + lines.append(rendered) + + tools_used = sorted(self.tools_used) + tools_set_raw = "{{{}}}".format(", ".join(tools_used)) if tools_used else "{}" + lines.append( + f"{inner_prefix}{ctx.bold('Tools used:')} {ctx.code_inline(tools_set_raw)}" + ) + lines.append( + f"{inner_prefix}{ctx.bold('Tool calls:')} {ctx.code_inline(str(self.tool_calls))}" + ) + return "\n".join(lines) + + +@dataclass +class ErrorSegment(Segment): + message: str + + def __init__(self, message: str) -> None: + super().__init__(kind="error") + self.message = str(message or "Unknown error") + + def render(self, ctx: RenderCtx) -> str: + return f"⚠️ {ctx.bold('Error:')} {ctx.code_inline(self.message)}" diff --git a/messaging/transcript/subagents.py b/messaging/transcript/subagents.py new file mode 100644 index 00000000..1ec561bc --- /dev/null +++ b/messaging/transcript/subagents.py @@ -0,0 +1,110 @@ +"""Task/subagent display state for messaging transcripts.""" + +from __future__ import annotations + +from typing import Any + +from loguru import logger + +from .segments import SubagentSegment + + +class SubagentState: + """Track active Task tool calls that suppress nested text/thinking output.""" + + def __init__(self, *, debug: bool = False) -> None: + self._stack: list[str] = [] + self._segments: list[SubagentSegment] = [] + self._debug = debug + + @property + def open_ids(self) -> tuple[str, ...]: + return tuple(self._stack) + + def in_subagent(self) -> bool: + return bool(self._stack) + + def current_segment(self) -> SubagentSegment | None: + return self._segments[-1] if self._segments else None + + def push(self, tool_id: str, segment: SubagentSegment) -> None: + marker = str(tool_id or "").strip() or f"__task_{len(self._stack) + 1}" + self._stack.append(marker) + self._segments.append(segment) + if self._debug: + logger.debug( + "SUBAGENT_STACK: push id=%r depth=%d heading=%r", + marker, + len(self._stack), + segment.description, + ) + + def close_for_tool_result(self, tool_id: str, *, tool_name: str | None) -> bool: + tool_id = str(tool_id or "").strip() + popped = self._pop(tool_id) + top = self._stack[-1] if self._stack else "" + looks_like_task_id = "task" in tool_id.lower() + + if ( + not popped + and tool_id + and top.startswith("__task_") + and tool_name in (None, "Task") + and looks_like_task_id + ): + return self._pop("") + return popped + + def _pop(self, tool_id: str) -> bool: + tool_id = str(tool_id or "").strip() + if not self._stack: + return False + + if tool_id: + if _ids_roughly_match(self._stack[-1], tool_id): + self._pop_to_depth(len(self._stack) - 1, tool_id, "LIFO") + return True + + for idx in range(len(self._stack) - 1, -1, -1): + if _ids_roughly_match(self._stack[idx], tool_id): + self._pop_to_depth(idx, tool_id, "matched") + return True + return False + + if self._stack[-1].startswith("__task_"): + self._pop_to_depth(len(self._stack) - 1, self._stack[-1], "synthetic") + return True + return False + + def _pop_to_depth(self, idx: int, requested_id: str, reason: str) -> None: + while len(self._stack) > idx: + popped = self._stack.pop() + if self._segments: + self._segments.pop() + if self._debug: + logger.debug( + "SUBAGENT_STACK: pop id=%r depth=%d (%s=%r)", + popped, + len(self._stack), + reason, + requested_id, + ) + + +def task_heading_from_input(input_value: Any) -> str: + if isinstance(input_value, dict): + for key in ("description", "subagent_type", "type"): + value = str(input_value.get(key, "") or "").strip() + if value: + return value + return "Subagent" + + +def _ids_roughly_match(stack_id: str, result_id: str) -> bool: + if not stack_id or not result_id: + return False + return ( + stack_id == result_id + or stack_id.startswith(result_id) + or result_id.startswith(stack_id) + ) diff --git a/pyproject.toml b/pyproject.toml index 7465ed2e..95a0ba11 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "free-claude-code" -version = "2.3.20" +version = "2.3.21" description = "Middleware between Claude Code CLI (Anthropic API) and NVIDIA NIM" readme = "README.md" requires-python = ">=3.14.0" diff --git a/tests/contracts/test_import_boundaries.py b/tests/contracts/test_import_boundaries.py index 5564bc18..4014fdc9 100644 --- a/tests/contracts/test_import_boundaries.py +++ b/tests/contracts/test_import_boundaries.py @@ -408,6 +408,27 @@ def test_admin_config_uses_package_owners_and_catalog_manifest() -> None: ) +def test_messaging_transcript_uses_package_owners() -> None: + repo_root = Path(__file__).resolve().parents[2] + messaging_root = repo_root / "messaging" + transcript_root = messaging_root / "transcript" + + assert not (messaging_root / "transcript.py").exists() + for filename in { + "__init__.py", + "buffer.py", + "context.py", + "renderer.py", + "segments.py", + "subagents.py", + }: + assert (transcript_root / filename).exists() + + init_text = (transcript_root / "__init__.py").read_text(encoding="utf-8") + assert "TranscriptBuffer" in init_text + assert "RenderCtx" in init_text + + def test_messaging_workflow_uses_split_runtime_owners() -> None: repo_root = Path(__file__).resolve().parents[2] messaging_root = repo_root / "messaging" diff --git a/tests/messaging/test_transcript.py b/tests/messaging/test_transcript.py index 09ae7938..fd2b7ac6 100644 --- a/tests/messaging/test_transcript.py +++ b/tests/messaging/test_transcript.py @@ -1,5 +1,3 @@ -from unittest.mock import patch - from messaging.rendering.telegram_markdown import ( escape_md_v2, escape_md_v2_code, @@ -8,6 +6,9 @@ from messaging.rendering.telegram_markdown import ( render_markdown_to_mdv2, ) from messaging.transcript import RenderCtx, TranscriptBuffer +from messaging.transcript.renderer import render_segments +from messaging.transcript.segments import Segment, SubagentSegment +from messaging.transcript.subagents import SubagentState def _ctx() -> RenderCtx: @@ -34,6 +35,17 @@ def test_transcript_order_thinking_tool_text(): assert out.find("think1") < out.find("Tool call:") < out.find("done") +def test_transcript_can_hide_tool_results(): + t = TranscriptBuffer(show_tool_results=False) + t.apply({"type": "tool_use", "id": "tool_1", "name": "ls", "input": {"path": "."}}) + t.apply({"type": "tool_result", "tool_use_id": "tool_1", "content": "secret"}) + + out = t.render(_ctx(), limit_chars=3900, status=None) + assert "Tool call:" in out + assert "Tool result:" not in out + assert "secret" not in out + + def test_transcript_subagent_suppresses_thinking_and_text_inside(): t = TranscriptBuffer() @@ -133,18 +145,11 @@ def test_transcript_subagent_closes_on_task_result_id_suffix_match(): def test_transcript_unmatched_non_task_tool_result_does_not_pop_subagent(): - t = TranscriptBuffer() - t.apply( - { - "type": "tool_use", - "id": "task_1", - "name": "Task", - "input": {"description": "Outer"}, - } - ) - t.apply({"type": "tool_result", "tool_use_id": "totally_unrelated", "content": "x"}) + state = SubagentState() + state.push("task_1", SubagentSegment("Outer")) - assert t._subagent_stack == ["task_1"] + assert not state.close_for_tool_result("totally_unrelated", tool_name=None) + assert state.open_ids == ("task_1",) def test_transcript_sequential_tasks_mismatched_results_no_depth_drift(): @@ -175,11 +180,12 @@ def test_transcript_sequential_tasks_mismatched_results_no_depth_drift(): "input": {"description": "C"}, } ) + t.apply({"type": "text_chunk", "text": "still hidden inside task three"}) out = t.render(_ctx(), limit_chars=3900, status=None) assert "🤖 *Subagent:* `A`\n 🤖 *Subagent:* `B`" not in out assert "\n 🤖 *Subagent:* `C`" not in out - assert t._subagent_stack == ["task_3"] + assert "still hidden inside task three" not in out def test_transcript_synthetic_task_start_closes_on_functions_task_result_id(): @@ -221,8 +227,10 @@ def test_transcript_synthetic_task_not_closed_by_unknown_non_task_result_id(): } ) t.apply({"type": "tool_result", "tool_use_id": "call_deadbeef", "content": "x"}) + t.apply({"type": "text_chunk", "text": "hidden while synthetic task is open"}) - assert t._subagent_stack == ["__task_1"] + out = t.render(_ctx(), limit_chars=3900, status=None) + assert "hidden while synthetic task is open" not in out def test_transcript_overlapping_tasks_are_flat_not_nested(): @@ -289,33 +297,43 @@ def test_transcript_reused_index_closes_previous_open_block(): t = TranscriptBuffer() # Open a text block at index 0, but never close it. t.apply({"type": "text_start", "index": 0}) - t.apply({"type": "text_delta", "index": 0, "text": "a"}) + t.apply({"type": "text_delta", "index": 0, "text": "alpha visible"}) # Provider reuses index 0 for a new tool block without a stop. t.apply( {"type": "tool_use_start", "index": 0, "id": "t1", "name": "ls", "input": {}} ) - # Old open text should have been closed. - assert 0 not in t._open_text_by_index - assert 0 in t._open_tools_by_index + t.apply({"type": "text_delta", "index": 0, "text": "omega visible"}) + + out = t.render(_ctx(), limit_chars=3900, status=None) + assert out.find("alpha") < out.find("Tool call:") < out.find("omega") def test_transcript_render_segment_exception_skipped(): """When a segment's render() raises, that segment is skipped and rest is rendered.""" - t = TranscriptBuffer() - t.apply({"type": "thinking_chunk", "text": "before"}) - t.apply({"type": "text_chunk", "text": "middle"}) - t.apply({"type": "text_chunk", "text": "after"}) - bad_segment = t._segments[1] + class StaticSegment(Segment): + def __init__(self, text: str) -> None: + super().__init__(kind="static") + self._text = text - def _raising_render(self, ctx): - raise ValueError("render failed") + def render(self, ctx: RenderCtx) -> str: + return self._text - with patch.object(bad_segment, "render", _raising_render): - out = t.render(_ctx(), limit_chars=3900, status=None) + class BrokenSegment(Segment): + def __init__(self) -> None: + super().__init__(kind="broken") + + def render(self, ctx: RenderCtx) -> str: + raise ValueError("render failed") + + out = render_segments( + [StaticSegment("before"), BrokenSegment(), StaticSegment("after")], + _ctx(), + limit_chars=3900, + status=None, + ) assert "before" in out assert "after" in out - assert "middle" not in out def test_transcript_render_status_only_exceeds_limit(): diff --git a/uv.lock b/uv.lock index 625240c1..82d62331 100644 --- a/uv.lock +++ b/uv.lock @@ -561,7 +561,7 @@ wheels = [ [[package]] name = "free-claude-code" -version = "2.3.20" +version = "2.3.21" source = { editable = "." } dependencies = [ { name = "aiohttp" },