mirror of
https://github.com/Alishahryar1/free-claude-code.git
synced 2026-04-28 03:20:01 +00:00
Revamped telegram result display
This commit is contained in:
Alishahryar1
parent
85eed8a1bc
commit
b95f2ef9c4
9 changed files with 812 additions and 261 deletions
|
|
@ -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":
|
||||
|
|
|
|||
|
|
@ -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],
|
||||
|
|
|
|||
402
messaging/transcript.py
Normal file
402
messaging/transcript.py
Normal file
|
|
@ -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 ""
|
||||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
84
tests/test_transcript.py
Normal file
84
tests/test_transcript.py
Normal file
|
|
@ -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
|
||||
Loading…
Add table
Add a link
Reference in a new issue