Yash030/claude-code-proxy
2
1"""Ordered transcript builder for messaging UIs (Telegram, etc.).2 3This module maintains an ordered list of "segments" that represent what the user4should see in the chat transcript: thinking, tool calls, tool results, subagent5headers, and assistant text. It is designed for in-place message editing where6the transcript grows over time and older content must be truncated.7"""8 9from __future__ import annotations10 11import json12from abc import ABC, abstractmethod13from collections import deque14from collections.abc import Callable, Iterable15from dataclasses import dataclass, field16from typing import Any17 18from loguru import logger19 20 21def _safe_json_dumps(obj: Any) -> str:22 try:23 return json.dumps(obj, indent=2, ensure_ascii=False, sort_keys=True)24 except Exception:25 return str(obj)26 27 28@dataclass29class Segment(ABC):30 kind: str31 32 @abstractmethod33 def render(self, ctx: RenderCtx) -> str: ...34 35 36@dataclass37class ThinkingSegment(Segment):38 def __init__(self) -> None:39 super().__init__(kind="thinking")40 self._parts: list[str] = []41 42 def append(self, t: str) -> None:43 if t:44 self._parts.append(t)45 46 @property47 def text(self) -> str:48 return "".join(self._parts)49 50 def render(self, ctx: RenderCtx) -> str:51 raw = self.text or ""52 if ctx.thinking_tail_max is not None and len(raw) > ctx.thinking_tail_max:53 raw = "..." + raw[-(ctx.thinking_tail_max - 3) :]54 inner = ctx.escape_code(raw)55 return f"💭 {ctx.bold('Thinking')}\n```\n{inner}\n```"56 57 58@dataclass59class TextSegment(Segment):60 def __init__(self) -> None:61 super().__init__(kind="text")62 self._parts: list[str] = []63 64 def append(self, t: str) -> None:65 if t:66 self._parts.append(t)67 68 @property69 def text(self) -> str:70 return "".join(self._parts)71 72 def render(self, ctx: RenderCtx) -> str:73 raw = self.text or ""74 if ctx.text_tail_max is not None and len(raw) > ctx.text_tail_max:75 raw = "..." + raw[-(ctx.text_tail_max - 3) :]76 return ctx.render_markdown(raw)77 78 79@dataclass80class ToolCallSegment(Segment):81 tool_use_id: str82 name: str83 closed: bool = False84 indent_level: int = 085 86 def __init__(self, tool_use_id: str, name: str, *, indent_level: int = 0) -> None:87 super().__init__(kind="tool_call")88 self.tool_use_id = str(tool_use_id or "")89 self.name = str(name or "tool")90 self.indent_level = max(0, int(indent_level))91 92 def render(self, ctx: RenderCtx) -> str:93 name = ctx.code_inline(self.name)94 # Per UX requirement: do not display tool args/results, only the tool call.95 prefix = " " * self.indent_level96 return f"{prefix}🛠 {ctx.bold('Tool call:')} {name}"97 98 99@dataclass100class ToolResultSegment(Segment):101 tool_use_id: str102 name: str | None103 content_text: str104 is_error: bool = False105 106 def __init__(107 self,108 tool_use_id: str,109 content: Any,110 *,111 name: str | None = None,112 is_error: bool = False,113 ) -> None:114 super().__init__(kind="tool_result")115 self.tool_use_id = str(tool_use_id or "")116 self.name = str(name) if name is not None else None117 self.is_error = bool(is_error)118 if isinstance(content, str):119 self.content_text = content120 else:121 self.content_text = _safe_json_dumps(content)122 123 def render(self, ctx: RenderCtx) -> str:124 raw = self.content_text or ""125 if ctx.tool_output_tail_max is not None and len(raw) > ctx.tool_output_tail_max:126 raw = "..." + raw[-(ctx.tool_output_tail_max - 3) :]127 inner = ctx.escape_code(raw)128 label = "Tool error:" if self.is_error else "Tool result:"129 maybe_name = f" {ctx.code_inline(self.name)}" if self.name else ""130 return f"📤 {ctx.bold(label)}{maybe_name}\n```\n{inner}\n```"131 132 133@dataclass134class SubagentSegment(Segment):135 description: str136 tool_calls: int = 0137 tools_used: set[str] = field(default_factory=set)138 current_tool: ToolCallSegment | None = None139 140 def __init__(self, description: str) -> None:141 super().__init__(kind="subagent")142 self.description = str(description or "Subagent")143 self.tool_calls = 0144 self.tools_used = set()145 self.current_tool = None146 147 def set_current_tool_call(self, tool_use_id: str, name: str) -> ToolCallSegment:148 tool_use_id = str(tool_use_id or "")149 name = str(name or "tool")150 self.tools_used.add(name)151 self.tool_calls += 1152 self.current_tool = ToolCallSegment(tool_use_id, name, indent_level=1)153 return self.current_tool154 155 def render(self, ctx: RenderCtx) -> str:156 inner_prefix = " "157 158 lines: list[str] = [159 f"🤖 {ctx.bold('Subagent:')} {ctx.code_inline(self.description)}"160 ]161 162 if self.current_tool is not None:163 try:164 rendered = self.current_tool.render(ctx)165 except Exception:166 rendered = ""167 if rendered:168 lines.append(rendered)169 170 tools_used = sorted(self.tools_used)171 tools_set_raw = "{{{}}}".format(", ".join(tools_used)) if tools_used else "{}"172 173 # Keep braces inside a code entity so MarkdownV2 doesn't require escaping them.174 lines.append(175 f"{inner_prefix}{ctx.bold('Tools used:')} {ctx.code_inline(tools_set_raw)}"176 )177 lines.append(178 f"{inner_prefix}{ctx.bold('Tool calls:')} {ctx.code_inline(str(self.tool_calls))}"179 )180 return "\n".join(lines)181 182 183@dataclass184class ErrorSegment(Segment):185 message: str186 187 def __init__(self, message: str) -> None:188 super().__init__(kind="error")189 self.message = str(message or "Unknown error")190 191 def render(self, ctx: RenderCtx) -> str:192 return f"⚠️ {ctx.bold('Error:')} {ctx.code_inline(self.message)}"193 194 195@dataclass196class RenderCtx:197 bold: Callable[[str], str]198 code_inline: Callable[[str], str]199 escape_code: Callable[[str], str]200 escape_text: Callable[[str], str]201 render_markdown: Callable[[str], str]202 203 thinking_tail_max: int | None = 1000204 tool_input_tail_max: int | None = 1200205 tool_output_tail_max: int | None = 1600206 text_tail_max: int | None = 2000207 208 209class TranscriptBuffer:210 """Maintains an ordered, truncatable transcript of events."""211 212 def __init__(213 self,214 *,215 show_tool_results: bool = True,216 debug_subagent_stack: bool = False,217 ) -> None:218 self._segments: list[Segment] = []219 self._open_thinking_by_index: dict[int, ThinkingSegment] = {}220 self._open_text_by_index: dict[int, TextSegment] = {}221 222 # content_block index -> tool call segment (for streaming tool args)223 self._open_tools_by_index: dict[int, ToolCallSegment] = {}224 225 # tool_use_id -> tool name (for tool_result labeling)226 self._tool_name_by_id: dict[str, str] = {}227 228 self._show_tool_results = bool(show_tool_results)229 230 # subagent context stack. Each entry is the Task tool_use_id we are waiting to close.231 self._subagent_stack: list[str] = []232 # Parallel stack of segments for rendering nested subagents.233 self._subagent_segments: list[SubagentSegment] = []234 self._debug_subagent_stack = debug_subagent_stack235 236 def _in_subagent(self) -> bool:237 return bool(self._subagent_stack)238 239 def _subagent_current(self) -> SubagentSegment | None:240 return self._subagent_segments[-1] if self._subagent_segments else None241 242 def _task_heading_from_input(self, inp: Any) -> str:243 # We never display full JSON args; only extract a short heading.244 if isinstance(inp, dict):245 desc = str(inp.get("description", "") or "").strip()246 if desc:247 return desc248 subagent_type = str(inp.get("subagent_type", "") or "").strip()249 if subagent_type:250 return subagent_type251 typ = str(inp.get("type", "") or "").strip()252 if typ:253 return typ254 return "Subagent"255 256 def _subagent_push(self, tool_id: str, seg: SubagentSegment) -> None:257 # Some providers can omit ids; still track depth for UI suppression.258 tool_id = (259 str(tool_id or "").strip() or f"__task_{len(self._subagent_stack) + 1}"260 )261 self._subagent_stack.append(tool_id)262 self._subagent_segments.append(seg)263 if self._debug_subagent_stack:264 logger.debug(265 "SUBAGENT_STACK: push id=%r depth=%d heading=%r",266 tool_id,267 len(self._subagent_stack),268 getattr(seg, "description", None),269 )270 271 def _subagent_pop(self, tool_id: str) -> bool:272 tool_id = str(tool_id or "").strip()273 if not self._subagent_stack:274 return False275 276 def _ids_roughly_match(stack_id: str, result_id: str) -> bool:277 if not stack_id or not result_id:278 return False279 if stack_id == result_id:280 return True281 # Some providers emit Task result ids with a suffix/prefix variant.282 # Treat those as the same logical Task invocation.283 return stack_id.startswith(result_id) or result_id.startswith(stack_id)284 285 if tool_id:286 # O(1) common case: LIFO - top of stack matches.287 if _ids_roughly_match(self._subagent_stack[-1], tool_id):288 self._subagent_stack.pop()289 if self._subagent_segments:290 self._subagent_segments.pop()291 if self._debug_subagent_stack:292 logger.debug(293 "SUBAGENT_STACK: pop id=%r depth=%d (LIFO)",294 tool_id,295 len(self._subagent_stack),296 )297 return True298 # Pop to the matching id (defensive against non-LIFO emissions).299 idx = -1300 for i in range(len(self._subagent_stack) - 1, -1, -1):301 if _ids_roughly_match(self._subagent_stack[i], tool_id):302 idx = i303 break304 if idx < 0:305 return False306 while len(self._subagent_stack) > idx:307 popped = self._subagent_stack.pop()308 if self._subagent_segments:309 self._subagent_segments.pop()310 if self._debug_subagent_stack:311 logger.debug(312 "SUBAGENT_STACK: pop id=%r depth=%d (matched=%r)",313 popped,314 len(self._subagent_stack),315 tool_id,316 )317 return True318 319 # No id in result; only close if we have a synthetic top marker.320 if self._subagent_stack and self._subagent_stack[-1].startswith("__task_"):321 popped = self._subagent_stack.pop()322 if self._subagent_segments:323 self._subagent_segments.pop()324 if self._debug_subagent_stack:325 logger.debug(326 "SUBAGENT_STACK: pop id=%r depth=%d (synthetic)",327 popped,328 len(self._subagent_stack),329 )330 return True331 return False332 333 def _ensure_thinking(self) -> ThinkingSegment:334 seg = ThinkingSegment()335 self._segments.append(seg)336 return seg337 338 def _ensure_text(self) -> TextSegment:339 seg = TextSegment()340 self._segments.append(seg)341 return seg342 343 def apply(self, ev: dict[str, Any]) -> None:344 """Apply a parsed event to the transcript."""345 et = ev.get("type")346 347 # Subagent rules: inside a Task/subagent, we only show tool calls/results.348 if self._in_subagent() and et in (349 "thinking_start",350 "thinking_delta",351 "thinking_chunk",352 "text_start",353 "text_delta",354 "text_chunk",355 ):356 return357 358 if et == "thinking_start":359 idx = int(ev.get("index", -1))360 if idx >= 0:361 # Defensive: if a provider reuses indices without emitting a stop,362 # close the previous open segment first.363 self.apply({"type": "block_stop", "index": idx})364 seg = self._ensure_thinking()365 if idx >= 0:366 self._open_thinking_by_index[idx] = seg367 return368 if et in ("thinking_delta", "thinking_chunk"):369 idx = int(ev.get("index", -1))370 seg = self._open_thinking_by_index.get(idx)371 if seg is None:372 seg = self._ensure_thinking()373 if idx >= 0:374 self._open_thinking_by_index[idx] = seg375 seg.append(str(ev.get("text", "")))376 return377 if et == "thinking_stop":378 idx = int(ev.get("index", -1))379 if idx >= 0:380 self._open_thinking_by_index.pop(idx, None)381 return382 383 if et == "text_start":384 idx = int(ev.get("index", -1))385 if idx >= 0:386 self.apply({"type": "block_stop", "index": idx})387 seg = self._ensure_text()388 if idx >= 0:389 self._open_text_by_index[idx] = seg390 return391 if et in ("text_delta", "text_chunk"):392 idx = int(ev.get("index", -1))393 seg = self._open_text_by_index.get(idx)394 if seg is None:395 seg = self._ensure_text()396 if idx >= 0:397 self._open_text_by_index[idx] = seg398 seg.append(str(ev.get("text", "")))399 return400 if et == "text_stop":401 idx = int(ev.get("index", -1))402 if idx >= 0:403 self._open_text_by_index.pop(idx, None)404 return405 406 if et == "tool_use_start":407 idx = int(ev.get("index", -1))408 if idx >= 0:409 self.apply({"type": "block_stop", "index": idx})410 tool_id = str(ev.get("id", "") or "").strip()411 name = str(ev.get("name", "") or "tool")412 if tool_id:413 self._tool_name_by_id[tool_id] = name414 415 # Task tool indicates subagent.416 if name == "Task":417 heading = self._task_heading_from_input(ev.get("input"))418 seg = SubagentSegment(heading)419 self._segments.append(seg)420 self._subagent_push(tool_id, seg)421 return422 423 # Normal tool call.424 if self._in_subagent():425 parent = self._subagent_current()426 if parent is not None:427 seg = parent.set_current_tool_call(tool_id, name)428 else:429 seg = ToolCallSegment(tool_id, name)430 self._segments.append(seg)431 else:432 seg = ToolCallSegment(tool_id, name)433 self._segments.append(seg)434 435 if idx >= 0:436 self._open_tools_by_index[idx] = seg437 return438 439 if et == "tool_use_delta":440 # Track open tool by index for tool_use_stop (closing state).441 return442 443 if et == "tool_use_stop":444 idx = int(ev.get("index", -1))445 seg = self._open_tools_by_index.pop(idx, None)446 if seg is not None:447 seg.closed = True448 return449 450 if et == "block_stop":451 idx = int(ev.get("index", -1))452 if idx in self._open_tools_by_index:453 self.apply({"type": "tool_use_stop", "index": idx})454 return455 if idx in self._open_thinking_by_index:456 self.apply({"type": "thinking_stop", "index": idx})457 return458 if idx in self._open_text_by_index:459 self.apply({"type": "text_stop", "index": idx})460 return461 return462 463 if et == "tool_use":464 tool_id = str(ev.get("id", "") or "").strip()465 name = str(ev.get("name", "") or "tool")466 if tool_id:467 self._tool_name_by_id[tool_id] = name468 469 if name == "Task":470 heading = self._task_heading_from_input(ev.get("input"))471 seg = SubagentSegment(heading)472 self._segments.append(seg)473 self._subagent_push(tool_id, seg)474 return475 476 if self._in_subagent():477 parent = self._subagent_current()478 if parent is not None:479 seg = parent.set_current_tool_call(tool_id, name)480 else:481 seg = ToolCallSegment(tool_id, name)482 self._segments.append(seg)483 else:484 seg = ToolCallSegment(tool_id, name)485 self._segments.append(seg)486 487 seg.closed = True488 return489 490 if et == "tool_result":491 tool_id = str(ev.get("tool_use_id", "") or "").strip()492 name = self._tool_name_by_id.get(tool_id)493 494 # If this was the Task tool result, close subagent context.495 if self._subagent_stack:496 popped = self._subagent_pop(tool_id)497 top = self._subagent_stack[-1] if self._subagent_stack else ""498 looks_like_task_id = "task" in tool_id.lower()499 # Some streams omit Task tool_use ids (synthetic stack ids), but include500 # a real Task id on tool_result (e.g. "functions.Task:0"). Reconcile that.501 if (502 not popped503 and tool_id504 and top.startswith("__task_")505 and (name in (None, "Task"))506 and looks_like_task_id507 ):508 self._subagent_pop("")509 510 if not self._show_tool_results:511 return512 513 seg = ToolResultSegment(514 tool_id,515 ev.get("content"),516 name=name,517 is_error=bool(ev.get("is_error", False)),518 )519 self._segments.append(seg)520 return521 522 if et == "error":523 self._segments.append(ErrorSegment(str(ev.get("message", ""))))524 return525 526 def render(self, ctx: RenderCtx, *, limit_chars: int, status: str | None) -> str:527 """Render transcript with truncation (drop oldest segments)."""528 # Filter out empty rendered segments.529 rendered: list[str] = []530 for seg in self._segments:531 try:532 out = seg.render(ctx)533 except Exception:534 continue535 if out:536 rendered.append(out)537 538 status_text = f"\n\n{status}" if status else ""539 prefix_marker = ctx.escape_text("... (truncated)\n")540 541 def _join(parts: Iterable[str], add_marker: bool) -> str:542 body = "\n".join(parts)543 if add_marker and body:544 body = prefix_marker + body545 return body + status_text if (body or status_text) else status_text546 547 # Fast path.548 candidate = _join(rendered, add_marker=False)549 if len(candidate) <= limit_chars:550 return candidate551 552 # Drop oldest segments until under limit (keep the tail).553 # Use deque for O(1) popleft; list.pop(0) would be O(n) per iteration.554 parts: deque[str] = deque(rendered)555 dropped = False556 last_part: str | None = None557 while parts:558 candidate = _join(parts, add_marker=True)559 if len(candidate) <= limit_chars:560 return candidate561 last_part = parts.popleft()562 dropped = True563 564 # Nothing fits - preserve tail of last segment instead of only marker+status.565 if dropped and last_part:566 budget = limit_chars - len(prefix_marker) - len(status_text)567 if budget > 20:568 if len(last_part) > budget:569 tail = "..." + last_part[-(budget - 3) :]570 else:571 tail = last_part572 candidate = prefix_marker + tail + status_text573 if len(candidate) <= limit_chars:574 return candidate575 576 # Fallback: marker + status only.577 if dropped:578 minimal = prefix_marker + status_text.lstrip("\n")579 if len(minimal) <= limit_chars:580 return minimal581 return status or ""582