diff --git a/messaging/event_parser.py b/messaging/event_parser.py index 5de5e348..6e748ec7 100644 --- a/messaging/event_parser.py +++ b/messaging/event_parser.py @@ -1,6 +1,7 @@ """CLI event parser for Claude Code CLI output. -Extracted from cli.parser to avoid tight coupling between messaging and cli packages. +This parser emits an ordered stream of low-level events suitable for building a +Claude Code-like transcript in messaging UIs. """ import logging @@ -23,12 +24,14 @@ def parse_cli_event(event: Any) -> List[Dict]: return [] etype = event.get("type") - results = [] + results: List[Dict[str, Any]] = [] - # 1. Handle full messages (assistant or result) + # 1. Handle full messages (assistant/user or result) msg_obj = None if etype == "assistant": msg_obj = event.get("message") + elif etype == "user": + msg_obj = event.get("message") elif etype == "result": res = event.get("result") if isinstance(res, dict): @@ -39,40 +42,35 @@ def parse_cli_event(event: Any) -> List[Dict]: if msg_obj and isinstance(msg_obj, dict): content = msg_obj.get("content", []) if isinstance(content, list): - parts = [] - thinking_parts = [] - tool_calls = [] + # Preserve order exactly as content blocks appear. for c in content: if not isinstance(c, dict): continue ctype = c.get("type") if ctype == "text": - parts.append(c.get("text", "")) + results.append({"type": "text_chunk", "text": c.get("text", "")}) elif ctype == "thinking": - thinking_parts.append(c.get("thinking", "")) + results.append( + {"type": "thinking_chunk", "text": c.get("thinking", "")} + ) elif ctype == "tool_use": - tool_calls.append(c) - - # Prioritize thinking first - if thinking_parts: - results.append({"type": "thinking", "text": "\n".join(thinking_parts)}) - - # Then tools or subagents - if tool_calls: - # Check for subagents (Task tool) - subagents = [ - t.get("input", {}).get("description", "Subagent") - for t in tool_calls - if t.get("name") == "Task" - ] - if subagents: - results.append({"type": "subagent_start", "tasks": subagents}) - else: - results.append({"type": "tool_start", "tools": tool_calls}) - - # Then text content if any - if parts: - results.append({"type": "content", "text": "".join(parts)}) + results.append( + { + "type": "tool_use", + "id": c.get("id", ""), + "name": c.get("name", ""), + "input": c.get("input"), + } + ) + elif ctype == "tool_result": + results.append( + { + "type": "tool_result", + "tool_use_id": c.get("tool_use_id", ""), + "content": c.get("content"), + "is_error": bool(c.get("is_error", False)), + } + ) if results: return results @@ -82,18 +80,53 @@ def parse_cli_event(event: Any) -> List[Dict]: delta = event.get("delta", {}) if isinstance(delta, dict): if delta.get("type") == "text_delta": - return [{"type": "content", "text": delta.get("text", "")}] + return [ + { + "type": "text_delta", + "index": event.get("index", -1), + "text": delta.get("text", ""), + } + ] if delta.get("type") == "thinking_delta": - return [{"type": "thinking", "text": delta.get("thinking", "")}] + return [ + { + "type": "thinking_delta", + "index": event.get("index", -1), + "text": delta.get("thinking", ""), + } + ] + if delta.get("type") == "input_json_delta": + return [ + { + "type": "tool_use_delta", + "index": event.get("index", -1), + "partial_json": delta.get("partial_json", ""), + } + ] # 3. Handle tool usage start if etype == "content_block_start": block = event.get("content_block", {}) - if isinstance(block, dict) and block.get("type") == "tool_use": - if block.get("name") == "Task": - desc = block.get("input", {}).get("description", "Subagent") - return [{"type": "subagent_start", "tasks": [desc]}] - return [{"type": "tool_start", "tools": [block]}] + if isinstance(block, dict): + btype = block.get("type") + if btype == "thinking": + return [{"type": "thinking_start", "index": event.get("index", -1)}] + if btype == "text": + return [{"type": "text_start", "index": event.get("index", -1)}] + if btype == "tool_use": + return [ + { + "type": "tool_use_start", + "index": event.get("index", -1), + "id": block.get("id", ""), + "name": block.get("name", ""), + "input": block.get("input"), + } + ] + + # 3.5 Handle block stop (to close open streaming segments) + if etype == "content_block_stop": + return [{"type": "block_stop", "index": event.get("index", -1)}] # 4. Handle errors and exit if etype == "error": diff --git a/messaging/handler.py b/messaging/handler.py index 1c5a4e85..e028266a 100644 --- a/messaging/handler.py +++ b/messaging/handler.py @@ -19,6 +19,7 @@ from .models import IncomingMessage from .session import SessionStore from .tree_queue import TreeQueueManager, MessageNode, MessageState, MessageTree from .event_parser import parse_cli_event +from .transcript import TranscriptBuffer, RenderCtx logger = logging.getLogger(__name__) @@ -76,7 +77,8 @@ def _normalize_gfm_tables(text: str) -> str: and _TABLE_SEP_RE.match(lines[idx + 1]) ): if out_lines and out_lines[-1].strip() != "": - indent = re.match(r"^(\s*)", line).group(1) + m = re.match(r"^(\s*)", line) + indent = m.group(1) if m else "" # A line of only whitespace counts as a blank line and preserves # list indentation contexts (tables inside list items). out_lines.append(indent) @@ -614,17 +616,18 @@ class ClaudeMessageHandler: if tree: await tree.update_state(node_id, MessageState.IN_PROGRESS) - # Components for structured display - components = { - "thinking": [], - "tools": [], - "subagents": [], - "content": [], - "errors": [], - } + transcript = TranscriptBuffer(show_tool_results=False) + render_ctx = RenderCtx( + bold=mdv2_bold, + code_inline=mdv2_code_inline, + escape_code=escape_md_v2_code, + escape_text=escape_md_v2, + render_markdown=render_markdown_to_mdv2, + ) last_ui_update = 0.0 last_displayed_text = None + had_transcript_events = False captured_session_id = None temp_session_id = None @@ -645,7 +648,7 @@ class ClaudeMessageHandler: return last_ui_update = now - display = self._build_message(components, status) + display = transcript.render(render_ctx, limit_chars=3900, status=status) if display and display != last_displayed_text: last_displayed_text = display await self.platform.queue_edit_message( @@ -667,7 +670,7 @@ class ClaudeMessageHandler: else: captured_session_id = session_or_temp_id except RuntimeError as e: - components["errors"].append(str(e)) + transcript.apply({"type": "error", "message": str(e)}) await update_ui( format_status("⏳", "Session limit reached"), force=True ) @@ -718,28 +721,43 @@ class ClaudeMessageHandler: logger.debug(f"HANDLER: Parsed {len(parsed_list)} events from CLI") for parsed in parsed_list: - if parsed["type"] == "thinking": - components["thinking"].append(parsed["text"]) + ptype = parsed.get("type") + if ptype in ( + "thinking_start", + "thinking_delta", + "thinking_chunk", + "thinking_stop", + "text_start", + "text_delta", + "text_chunk", + "text_stop", + "tool_use_start", + "tool_use_delta", + "tool_use_stop", + "tool_use", + "tool_result", + "block_stop", + "error", + ): + transcript.apply(parsed) + had_transcript_events = True + + if ptype in ("thinking_start", "thinking_delta", "thinking_chunk"): await update_ui(format_status("🧠", "Claude is thinking...")) - - elif parsed["type"] == "content": - if parsed.get("text"): - components["content"].append(parsed["text"]) + elif ptype in ("text_start", "text_delta", "text_chunk"): await update_ui(format_status("🧠", "Claude is working...")) - - elif parsed["type"] == "tool_start": - names = [t.get("name") for t in parsed.get("tools", [])] - components["tools"].extend(names) + elif ptype in ("tool_use_start", "tool_use_delta", "tool_use"): + if parsed.get("name") == "Task": + await update_ui(format_status("🤖", "Subagent working...")) + else: + await update_ui(format_status("⏳", "Executing tools...")) + elif ptype == "tool_result": await update_ui(format_status("⏳", "Executing tools...")) - elif parsed["type"] == "subagent_start": - tasks = parsed.get("tasks", []) - components["subagents"].extend(tasks) - await update_ui(format_status("🤖", "Subagent working...")) - elif parsed["type"] == "complete": - if not any(components.values()): - components["content"].append("Done.") + # If nothing happened (rare), still show a completion marker. + if not had_transcript_events: + transcript.apply({"type": "text_chunk", "text": "Done."}) logger.info("HANDLER: Task complete, updating UI") await update_ui(format_status("✅", "Complete"), force=True) @@ -757,7 +775,6 @@ class ClaudeMessageHandler: logger.error( f"HANDLER: Error event received: {error_msg[:200]}" ) - components["errors"].append(error_msg) logger.info("HANDLER: Updating UI with error status") await update_ui(format_status("❌", "Error"), force=True) if tree: @@ -774,7 +791,7 @@ class ClaudeMessageHandler: if cancel_reason == "stop": await update_ui(format_status("⏹", "Stopped."), force=True) else: - components["errors"].append("Task was cancelled") + transcript.apply({"type": "error", "message": "Task was cancelled"}) await update_ui(format_status("❌", "Cancelled"), force=True) # Do not propagate cancellation to children; a reply-scoped "/stop" @@ -788,16 +805,14 @@ class ClaudeMessageHandler: f"HANDLER: Task failed with exception: {type(e).__name__}: {e}" ) error_msg = str(e)[:200] - components["errors"].append(error_msg) + transcript.apply({"type": "error", "message": error_msg}) await update_ui(format_status("💥", "Task Failed"), force=True) if tree: await self._propagate_error_to_children( node_id, error_msg, "Parent task failed" ) finally: - logger.info( - f"HANDLER: _process_node completed for node {node_id}, errors={len(components['errors'])}" - ) + logger.info(f"HANDLER: _process_node completed for node {node_id}") # Free the session-manager slot. Session IDs are persisted in the tree and # can be resumed later by ID; we don't need to keep a CLISession instance # around after this node completes. @@ -830,87 +845,6 @@ class ClaudeMessageHandler: ) ) - def _build_message( - self, - components: dict, - status: Optional[str] = None, - ) -> str: - """ - Build unified message with specific order. - Handles truncation while preserving markdown structure (closing code blocks). - """ - lines = [] - - # 1. Thinking - if components["thinking"]: - thinking_text = "".join(components["thinking"]) - # Truncate thinking if too long, it's usually less critical than final content - if len(thinking_text) > 1000: - thinking_text = "..." + thinking_text[-995:] - - lines.append( - f"💭 {mdv2_bold('Thinking:')}\n```\n{escape_md_v2_code(thinking_text)}\n```" - ) - - # 2. Tools - if components["tools"]: - unique_tools = [] - seen = set() - for t in components["tools"]: - if t and t not in seen: - unique_tools.append(str(t)) - seen.add(t) - if unique_tools: - lines.append( - f"🛠 {mdv2_bold('Tools:')} {mdv2_code_inline(', '.join(unique_tools))}" - ) - - # 3. Subagents - if components["subagents"]: - for task in components["subagents"]: - lines.append(f"🤖 {mdv2_bold('Subagent:')} {mdv2_code_inline(task)}") - - # 4. Content - if components["content"]: - lines.append(render_markdown_to_mdv2("".join(components["content"]))) - - # 5. Errors - if components["errors"]: - for err in components["errors"]: - lines.append(f"⚠️ {mdv2_bold('Error:')} {mdv2_code_inline(err)}") - - if not any(lines) and not status: - return format_status("⏳", "Claude is working...") - - # Telegram character limit is 4096. We leave buffer for status updates. - LIMIT = 3900 - - # Filter out empty lines first for a clean join - lines = [l for l in lines if l] - - main_text = "\n".join(lines) - status_text = f"\n\n{status}" if status else "" - - if len(main_text) + len(status_text) <= LIMIT: - return ( - main_text + status_text - if main_text + status_text - else format_status("⏳", "Claude is working...") - ) - - # If too long, truncate the start of the content (keep the end) - available_limit = LIMIT - len(status_text) - 20 # 20 for truncation marker - raw_truncated = main_text[-available_limit:].lstrip() - - # Check for unbalanced code blocks - prefix = escape_md_v2("... (truncated)\n") - if raw_truncated.count("```") % 2 != 0: - prefix += "```\n" - - truncated_main = prefix + raw_truncated - - return truncated_main + status_text - def _get_initial_status( self, tree: Optional[object], diff --git a/messaging/transcript.py b/messaging/transcript.py new file mode 100644 index 00000000..99c58403 --- /dev/null +++ b/messaging/transcript.py @@ -0,0 +1,402 @@ +"""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 dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional + + +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: + kind: str + + def render(self, ctx: "RenderCtx") -> str: + raise NotImplementedError + + +@dataclass +class ThinkingSegment(Segment): + text: str = "" + + def __init__(self) -> None: + super().__init__(kind="thinking") + + def append(self, t: str) -> None: + if t: + self.text += t + + 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): + text: str = "" + + def __init__(self) -> None: + super().__init__(kind="text") + + def append(self, t: str) -> None: + if t: + self.text += t + + 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 + input_text: str = "" + closed: bool = False + + def __init__(self, tool_use_id: str, name: str) -> None: + super().__init__(kind="tool_call") + self.tool_use_id = str(tool_use_id or "") + self.name = str(name or "tool") + + def set_initial_input(self, inp: Any) -> None: + if inp is None: + return + if isinstance(inp, str): + self.input_text = inp + else: + self.input_text = _safe_json_dumps(inp) + + def append_input_delta(self, partial: str) -> None: + if partial: + self.input_text += partial + + def render(self, ctx: "RenderCtx") -> str: + raw = self.input_text or "" + if ctx.tool_input_tail_max is not None and len(raw) > ctx.tool_input_tail_max: + raw = "..." + raw[-(ctx.tool_input_tail_max - 3) :] + inner = ctx.escape_code(raw) + name = ctx.code_inline(self.name) + return f"🛠 {ctx.bold('Tool call:')} {name}\n```\n{inner}\n```" + + +@dataclass +class ToolResultSegment(Segment): + tool_use_id: str + name: Optional[str] + content_text: str + is_error: bool = False + + def __init__( + self, + tool_use_id: str, + content: Any, + *, + name: Optional[str] = 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 SubagentHeaderSegment(Segment): + description: str + + def __init__(self, description: str) -> None: + super().__init__(kind="subagent") + self.description = str(description or "Subagent") + + def render(self, ctx: "RenderCtx") -> str: + return f"🤖 {ctx.bold('Subagent:')} {ctx.code_inline(self.description)}" + + +@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: Optional[int] = 1000 + tool_input_tail_max: Optional[int] = 1200 + tool_output_tail_max: Optional[int] = 1600 + text_tail_max: Optional[int] = 2000 + + +class TranscriptBuffer: + """Maintains an ordered, truncatable transcript of events.""" + + def __init__(self, *, show_tool_results: bool = True) -> 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] = [] + + def _in_subagent(self) -> bool: + return bool(self._subagent_stack) + + 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)) + 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)) + 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)) + tool_id = str(ev.get("id", "") or "") + name = str(ev.get("name", "") or "tool") + seg = ToolCallSegment(tool_id, name) + seg.set_initial_input(ev.get("input")) + self._segments.append(seg) + if idx >= 0: + self._open_tools_by_index[idx] = seg + if tool_id: + self._tool_name_by_id[tool_id] = name + + # Task tool indicates subagent. + if name == "Task": + desc = "" + inp = ev.get("input") + if isinstance(inp, dict): + desc = str(inp.get("description", "") or "") + if not desc: + desc = "Subagent" + self._segments.append(SubagentHeaderSegment(desc)) + if tool_id: + self._subagent_stack.append(tool_id) + return + + if et == "tool_use_delta": + idx = int(ev.get("index", -1)) + partial = str(ev.get("partial_json", "") or "") + seg = self._open_tools_by_index.get(idx) + if seg is not None: + seg.append_input_delta(partial) + 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 "") + name = str(ev.get("name", "") or "tool") + seg = ToolCallSegment(tool_id, name) + seg.set_initial_input(ev.get("input")) + seg.closed = True + self._segments.append(seg) + if tool_id: + self._tool_name_by_id[tool_id] = name + + if name == "Task": + desc = "" + inp = ev.get("input") + if isinstance(inp, dict): + desc = str(inp.get("description", "") or "") + if not desc: + desc = "Subagent" + self._segments.append(SubagentHeaderSegment(desc)) + if tool_id: + self._subagent_stack.append(tool_id) + return + + if et == "tool_result": + tool_id = str(ev.get("tool_use_id", "") or "") + name = self._tool_name_by_id.get(tool_id) + + # If this was the Task tool result, close subagent context. + if tool_id and self._subagent_stack and self._subagent_stack[-1] == tool_id: + self._subagent_stack.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: Optional[str]) -> 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: List[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. + parts = list(rendered) + dropped = False + while parts: + candidate = _join(parts, add_marker=True) + if len(candidate) <= limit_chars: + return candidate + parts.pop(0) + dropped = True + + # Nothing fits; return status only with marker if possible. + if dropped: + minimal = prefix_marker + status_text.lstrip("\n") + if len(minimal) <= limit_chars: + return minimal + return status or "" diff --git a/tests/test_cli.py b/tests/test_cli.py index 30dd866c..7456ef5d 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -20,7 +20,7 @@ class TestCLIParser: } result = parse_cli_event(event) assert len(result) == 1 - assert result[0]["type"] == "content" + assert result[0]["type"] == "text_chunk" assert result[0]["text"] == "Hello, world!" def test_parse_thinking_content(self): @@ -33,7 +33,7 @@ class TestCLIParser: } result = parse_cli_event(event) assert len(result) == 1 - assert result[0]["type"] == "thinking" + assert result[0]["type"] == "thinking_chunk" assert ( result[0]["text"] == "Let me think...\n" or result[0]["text"] == "Let me think..." @@ -52,9 +52,9 @@ class TestCLIParser: } result = parse_cli_event(event) assert len(result) == 2 - assert result[0]["type"] == "thinking" + assert result[0]["type"] == "thinking_chunk" assert result[0]["text"] == "Thinking..." - assert result[1]["type"] == "tool_start" + assert result[1]["type"] == "tool_use" def test_parse_tool_use(self): """Test parsing tool use content.""" @@ -72,30 +72,32 @@ class TestCLIParser: } result = parse_cli_event(event) assert len(result) == 1 - assert result[0]["type"] == "tool_start" - assert len(result[0]["tools"]) == 1 - assert result[0]["tools"][0]["name"] == "read_file" + assert result[0]["type"] == "tool_use" + assert result[0]["name"] == "read_file" + assert result[0]["input"] == {"path": "/test"} def test_parse_text_delta(self): """Test parsing streaming text delta.""" event = { "type": "content_block_delta", + "index": 0, "delta": {"type": "text_delta", "text": "streaming text"}, } result = parse_cli_event(event) assert len(result) == 1 - assert result[0]["type"] == "content" + assert result[0]["type"] == "text_delta" assert result[0]["text"] == "streaming text" def test_parse_thinking_delta(self): """Test parsing streaming thinking delta.""" event = { "type": "content_block_delta", + "index": 1, "delta": {"type": "thinking_delta", "thinking": "thinking..."}, } result = parse_cli_event(event) assert len(result) == 1 - assert result[0]["type"] == "thinking" + assert result[0]["type"] == "thinking_delta" assert result[0]["text"] == "thinking..." def test_parse_error(self): diff --git a/tests/test_event_parser.py b/tests/test_event_parser.py index 0ebae82b..aa7d76d2 100644 --- a/tests/test_event_parser.py +++ b/tests/test_event_parser.py @@ -13,8 +13,8 @@ def test_parse_cli_event_assistant_content(): } results = parse_cli_event(event) assert len(results) == 2 - assert results[0] == {"type": "thinking", "text": "Internal thought"} - assert results[1] == {"type": "content", "text": "Hello user"} + assert results[0] == {"type": "thinking_chunk", "text": "Internal thought"} + assert results[1] == {"type": "text_chunk", "text": "Hello user"} def test_parse_cli_event_assistant_tools(): @@ -26,8 +26,9 @@ def test_parse_cli_event_assistant_tools(): } results = parse_cli_event(event) assert len(results) == 1 - assert results[0]["type"] == "tool_start" - assert results[0]["tools"][0]["name"] == "ls" + assert results[0]["type"] == "tool_use" + assert results[0]["name"] == "ls" + assert results[0]["input"] == {"path": "."} def test_parse_cli_event_assistant_subagent(): @@ -45,31 +46,37 @@ def test_parse_cli_event_assistant_subagent(): } results = parse_cli_event(event) assert len(results) == 1 - assert results[0]["type"] == "subagent_start" - assert results[0]["tasks"] == ["Fix bug"] + assert results[0]["type"] == "tool_use" + assert results[0]["name"] == "Task" + assert results[0]["input"] == {"description": "Fix bug"} def test_parse_cli_event_content_block_delta(): # Text delta event_text = { "type": "content_block_delta", + "index": 0, "delta": {"type": "text_delta", "text": " more"}, } results_text = parse_cli_event(event_text) - assert results_text == [{"type": "content", "text": " more"}] + assert results_text == [{"type": "text_delta", "index": 0, "text": " more"}] # Thinking delta event_think = { "type": "content_block_delta", + "index": 1, "delta": {"type": "thinking_delta", "thinking": " more thought"}, } results_think = parse_cli_event(event_think) - assert results_think == [{"type": "thinking", "text": " more thought"}] + assert results_think == [ + {"type": "thinking_delta", "index": 1, "text": " more thought"} + ] def test_parse_cli_event_content_block_start(): event = { "type": "content_block_start", + "index": 2, "content_block": { "type": "tool_use", "name": "Task", @@ -77,7 +84,15 @@ def test_parse_cli_event_content_block_start(): }, } results = parse_cli_event(event) - assert results == [{"type": "subagent_start", "tasks": ["deploy"]}] + assert results == [ + { + "type": "tool_use_start", + "index": 2, + "id": "", + "name": "Task", + "input": {"description": "deploy"}, + } + ] def test_parse_cli_event_error(): @@ -86,6 +101,31 @@ def test_parse_cli_event_error(): assert results == [{"type": "error", "message": "something failed"}] +def test_parse_cli_event_user_tool_result(): + event = { + "type": "user", + "message": { + "content": [ + { + "type": "tool_result", + "tool_use_id": "tool_1", + "content": "ok", + "is_error": False, + } + ] + }, + } + results = parse_cli_event(event) + assert results == [ + { + "type": "tool_result", + "tool_use_id": "tool_1", + "content": "ok", + "is_error": False, + } + ] + + def test_parse_cli_event_exit_success(): event = {"type": "exit", "code": 0} results = parse_cli_event(event) diff --git a/tests/test_handler_format.py b/tests/test_handler_format.py index a7094a21..9c4d5830 100644 --- a/tests/test_handler_format.py +++ b/tests/test_handler_format.py @@ -1,6 +1,13 @@ import pytest from unittest.mock import MagicMock -from messaging.handler import ClaudeMessageHandler, escape_md_v2 +from messaging.handler import escape_md_v2 +from messaging.transcript import TranscriptBuffer, RenderCtx +from messaging.handler import ( + escape_md_v2_code, + mdv2_bold, + mdv2_code_inline, + render_markdown_to_mdv2, +) @pytest.fixture @@ -8,21 +15,49 @@ def handler(): platform = MagicMock() cli = MagicMock() store = MagicMock() - return ClaudeMessageHandler(platform, cli, store) + # Kept for backwards test structure; transcript rendering is now separate. + return (platform, cli, store) -def test_build_message_structure(handler): - """Verify the order of message components.""" - components = { - "thinking": ["Thinking process..."], - "tools": ["list_files", "read_file"], - "subagents": ["Searching codebase...", "Analyzing dependencies..."], - "content": ["Here is the file content."], - "errors": ["Some error happened"], - } +def _ctx() -> RenderCtx: + return RenderCtx( + bold=mdv2_bold, + code_inline=mdv2_code_inline, + escape_code=escape_md_v2_code, + escape_text=escape_md_v2, + render_markdown=render_markdown_to_mdv2, + ) + + +def test_transcript_structure_and_order(handler): + """Verify ordered transcript rendering (thinking/tool/subagent/text/error/status).""" status = "✅ *Complete*" + t = TranscriptBuffer() - msg = handler._build_message(components, status) + # Apply in a deliberate sequence. + t.apply({"type": "thinking_chunk", "text": "Thinking process..."}) + t.apply( + {"type": "tool_use", "id": "t1", "name": "list_files", "input": {"path": "."}} + ) + + # Subagent marker (Task tool). + t.apply( + { + "type": "tool_use", + "id": "task1", + "name": "Task", + "input": {"description": "Searching codebase..."}, + } + ) + t.apply( + {"type": "tool_use", "id": "t2", "name": "read_file", "input": {"path": "x.py"}} + ) + t.apply({"type": "tool_result", "tool_use_id": "task1", "content": "done"}) + + t.apply({"type": "text_chunk", "text": "Here is the file content."}) + t.apply({"type": "error", "message": "Some error happened"}) + + msg = t.render(_ctx(), limit_chars=3900, status=status) print(f"Generated Message:\n{msg}") @@ -35,37 +70,32 @@ def test_build_message_structure(handler): assert "Some error happened" in msg assert "✅ *Complete*" in msg - # Check headers - assert "💭 *Thinking:*" in msg - assert "🛠 *Tools:*" in msg + # Check headers/markers used in the transcript. + assert "💭 *Thinking*" in msg + assert "🛠 *Tool call:*" in msg assert "🤖 *Subagent:*" in msg assert "⚠️ *Error:*" in msg - # Check Order: Thinking -> Tools -> Subagents -> Content -> Errors -> Status + # Check Order: Thinking -> Tool call -> Subagent -> Content -> Errors -> Status p_thinking = msg.find("Thinking process...") - p_tools = msg.find("🛠 *Tools:*") - p_subagents = msg.find("🤖 *Subagent:*") + p_tool_call = msg.find("🛠 *Tool call:*") + p_subagent = msg.find("🤖 *Subagent:*") p_content = msg.find(escape_md_v2("Here is the file content.")) p_errors = msg.find("⚠️ *Error:*") p_status = msg.find("✅ *Complete*") - assert p_thinking < p_tools, "Thinking should come before Tools" - assert p_tools < p_subagents, "Tools should come before Subagents" - assert p_subagents < p_content, "Subagents should come before Content" + assert p_thinking < p_tool_call, "Thinking should come before tool calls" + assert p_tool_call < p_subagent, "Tool calls should come before subagent marker" + assert p_subagent < p_content, "Subagent should come before Content" assert p_content < p_errors, "Content should come before Errors" assert p_errors < p_status, "Errors should come before Status" -def test_build_message_simple(handler): - """Verify simple message with just content.""" - components = { - "thinking": [], - "tools": [], - "subagents": [], - "content": ["Simple message."], - "errors": [], - } - msg = handler._build_message(components, "Ready") +def test_transcript_simple(handler): + """Verify simple transcript with just text + status.""" + t = TranscriptBuffer() + t.apply({"type": "text_chunk", "text": "Simple message."}) + msg = t.render(_ctx(), limit_chars=3900, status="Ready") assert escape_md_v2("Simple message.") in msg assert "Ready" in msg @@ -74,15 +104,27 @@ def test_build_message_simple(handler): def test_subagents_formatting(handler): - """Verify subagents formatting.""" - components = { - "thinking": [], - "tools": [], - "subagents": ["Task 1", "Task 2"], - "content": [], - "errors": [], - } - msg = handler._build_message(components) + """Verify subagent formatting (Task tool).""" + t = TranscriptBuffer() + t.apply( + { + "type": "tool_use", + "id": "task_1", + "name": "Task", + "input": {"description": "Task 1"}, + } + ) + t.apply({"type": "tool_result", "tool_use_id": "task_1", "content": "done"}) + t.apply( + { + "type": "tool_use", + "id": "task_2", + "name": "Task", + "input": {"description": "Task 2"}, + } + ) + + msg = t.render(_ctx(), limit_chars=3900, status=None) assert "🤖 *Subagent:* `Task 1`" in msg assert "🤖 *Subagent:* `Task 2`" in msg diff --git a/tests/test_reliability.py b/tests/test_reliability.py index 5c351e35..35d0cf23 100644 --- a/tests/test_reliability.py +++ b/tests/test_reliability.py @@ -74,29 +74,37 @@ async def test_telegram_no_retry_on_bad_request(telegram_platform): def test_handler_build_message_hardening(): - handler = ClaudeMessageHandler(AsyncMock(), AsyncMock(), AsyncMock()) + # Formatting hardening now lives in TranscriptBuffer rendering. + from messaging.transcript import TranscriptBuffer, RenderCtx - # Case 1: Empty components - components = { - "thinking": [], - "tools": [], - "subagents": [], - "content": [], - "errors": [], - } - msg = handler._build_message(components) - assert msg == format_status("⏳", "Claude is working...") + from messaging.handler import ( + escape_md_v2, + escape_md_v2_code, + mdv2_bold, + mdv2_code_inline, + render_markdown_to_mdv2, + ) - # Case 2: Truncation with code block closing - long_thinking = "thought " * 200 # ~1400 chars - components["thinking"] = [long_thinking] - components["content"] = ["This is a very long message. " * 300] # ~ 8700 chars + ctx = RenderCtx( + bold=mdv2_bold, + code_inline=mdv2_code_inline, + escape_code=escape_md_v2_code, + escape_text=escape_md_v2, + render_markdown=render_markdown_to_mdv2, + ) - msg = handler._build_message(components, status="Finishing...") + # Case 1: Empty transcript + no status => empty string. + t = TranscriptBuffer() + msg = t.render(ctx, limit_chars=3900, status=None) + assert msg == "" + + # Case 2: Truncation with code block closing and status preserved. + t.apply({"type": "thinking_chunk", "text": ("thought " * 200)}) + t.apply({"type": "text_chunk", "text": ("This is a very long message. " * 300)}) + + msg = t.render(ctx, limit_chars=3900, status="Finishing...") assert len(msg) <= 4096 - assert "truncated" in msg assert "Finishing..." in msg - # If thinking contains backticks, they should be balanced if "```" in msg: assert msg.count("```") % 2 == 0 diff --git a/tests/test_robust_formatting.py b/tests/test_robust_formatting.py index 43d1dcb6..6fe9deac 100644 --- a/tests/test_robust_formatting.py +++ b/tests/test_robust_formatting.py @@ -1,6 +1,13 @@ import pytest from unittest.mock import MagicMock -from messaging.handler import ClaudeMessageHandler, escape_md_v2 +from messaging.handler import escape_md_v2 +from messaging.transcript import TranscriptBuffer, RenderCtx +from messaging.handler import ( + escape_md_v2_code, + mdv2_bold, + mdv2_code_inline, + render_markdown_to_mdv2, +) @pytest.fixture @@ -8,27 +15,37 @@ def handler(): platform = MagicMock() cli = MagicMock() store = MagicMock() - return ClaudeMessageHandler(platform, cli, store) + return (platform, cli, store) + + +def _ctx() -> RenderCtx: + return RenderCtx( + bold=mdv2_bold, + code_inline=mdv2_code_inline, + escape_code=escape_md_v2_code, + escape_text=escape_md_v2, + render_markdown=render_markdown_to_mdv2, + ) def test_truncation_closes_code_blocks(handler): """Verify that truncation correctly closes open code blocks.""" - components = { - "thinking": [ - "Starting some long thinking process that will definitely cause truncation later on..." - ], - "tools": [], - "subagents": [], - "content": [ - "```python\ndef very_long_function():\n # " + "A" * 4000 - ], # Long content - "errors": [], - } + t = TranscriptBuffer() + t.apply( + { + "type": "thinking_chunk", + "text": "Starting some long thinking process that will definitely cause truncation later on...", + } + ) + t.apply( + { + "type": "text_chunk", + "text": "```python\ndef very_long_function():\n # " + ("A" * 4000), + } + ) - msg = handler._build_message(components, "✅ *Complete*") + msg = t.render(_ctx(), limit_chars=3900, status="✅ *Complete*") - assert escape_md_v2("... (truncated)") in msg - # The limit is 3900. Our content + thinking is > 4000. # The backtick count must be even to be a valid block. assert msg.count("```") % 2 == 0 assert msg.endswith("```") or "✅ *Complete*" in msg.split("```")[-1] @@ -36,29 +53,18 @@ def test_truncation_closes_code_blocks(handler): def test_truncation_preserves_status(handler): """Verify that status is still appended after truncation.""" - components = { - "thinking": ["Thinking..."], - "tools": [], - "subagents": [], - "content": ["A" * 5000], - "errors": [], - } status = "READY_STATUS" - msg = handler._build_message(components, status) + t = TranscriptBuffer() + t.apply({"type": "thinking_chunk", "text": "Thinking..."}) + t.apply({"type": "text_chunk", "text": "A" * 5000}) + msg = t.render(_ctx(), limit_chars=3900, status=status) assert status in msg - assert escape_md_v2("... (truncated)") in msg def test_empty_components_with_status(handler): """Verify message building with just a status.""" - components = { - "thinking": [], - "tools": [], - "subagents": [], - "content": [], - "errors": [], - } status = "Simple Status" - msg = handler._build_message(components, status) + t = TranscriptBuffer() + msg = t.render(_ctx(), limit_chars=3900, status=status) assert msg == "\n\nSimple Status" diff --git a/tests/test_transcript.py b/tests/test_transcript.py new file mode 100644 index 00000000..c4c7e0ad --- /dev/null +++ b/tests/test_transcript.py @@ -0,0 +1,84 @@ +from messaging.transcript import TranscriptBuffer, RenderCtx + +from messaging.handler import ( + escape_md_v2, + escape_md_v2_code, + mdv2_bold, + mdv2_code_inline, + render_markdown_to_mdv2, +) + + +def _ctx() -> RenderCtx: + return RenderCtx( + bold=mdv2_bold, + code_inline=mdv2_code_inline, + escape_code=escape_md_v2_code, + escape_text=escape_md_v2, + render_markdown=render_markdown_to_mdv2, + thinking_tail_max=1000, + tool_input_tail_max=1200, + tool_output_tail_max=1600, + text_tail_max=2000, + ) + + +def test_transcript_order_thinking_tool_text(): + t = TranscriptBuffer() + t.apply({"type": "thinking_chunk", "text": "think1"}) + t.apply({"type": "tool_use", "id": "tool_1", "name": "ls", "input": {"path": "."}}) + t.apply({"type": "text_chunk", "text": "done"}) + + out = t.render(_ctx(), limit_chars=3900, status=None) + assert out.find("think1") < out.find("Tool call:") < out.find("done") + + +def test_transcript_subagent_suppresses_thinking_and_text_inside(): + t = TranscriptBuffer() + + # Enter subagent context (Task tool call). + t.apply( + { + "type": "tool_use", + "id": "task_1", + "name": "Task", + "input": {"description": "Fix bug"}, + } + ) + + # These should be suppressed while inside subagent context. + t.apply({"type": "thinking_delta", "index": -1, "text": "secret"}) + t.apply({"type": "text_chunk", "text": "visible?"}) + + # Tool activity should still show. + t.apply({"type": "tool_use", "id": "tool_2", "name": "ls", "input": {"path": "."}}) + t.apply({"type": "tool_result", "tool_use_id": "tool_2", "content": "x"}) + + # Close subagent context (Task tool result). + t.apply({"type": "tool_result", "tool_use_id": "task_1", "content": "done"}) + + # Now text should show again. + t.apply({"type": "text_chunk", "text": "after"}) + + out = t.render(_ctx(), limit_chars=3900, status=None) + assert "Subagent:" in out + assert "secret" not in out + assert "visible?" not in out + assert "after" in out + + +def test_transcript_truncates_by_dropping_oldest_segments(): + t = TranscriptBuffer() + + # Create many segments by opening/closing distinct text blocks. + for i in range(60): + t.apply({"type": "text_start", "index": i}) + t.apply( + {"type": "text_delta", "index": i, "text": f"segment_{i} " + ("x" * 120)} + ) + t.apply({"type": "block_stop", "index": i}) + + out = t.render(_ctx(), limit_chars=600, status="status") + assert escape_md_v2("... (truncated)") in out + assert escape_md_v2("segment_59") in out + assert escape_md_v2("segment_0") not in out