|
| 1 | +"""Subagent tracing: correlate nested agent sessions into a parent-child tree. |
| 2 | +
|
| 3 | +Subagent sessions are linked via SessionMeta.parent_session_id and |
| 4 | +parent_event_id. This module provides: |
| 5 | +
|
| 6 | + - Tree building: reconstruct the full agent call tree from the store |
| 7 | + - Tree-aware replay: render the tree with inline subagent expansion |
| 8 | + - Aggregated stats: roll up tool calls, tokens, errors across the tree |
| 9 | +""" |
| 10 | + |
| 11 | +from __future__ import annotations |
| 12 | + |
| 13 | +import argparse |
| 14 | +import sys |
| 15 | +from dataclasses import dataclass, field |
| 16 | +from typing import TextIO |
| 17 | + |
| 18 | +from .models import EventType, SessionMeta, TraceEvent |
| 19 | +from .store import TraceStore |
| 20 | + |
| 21 | +MAX_DEPTH = 5 # configurable guard against runaway recursion |
| 22 | + |
| 23 | + |
| 24 | +# --------------------------------------------------------------------------- |
| 25 | +# Tree data structure |
| 26 | +# --------------------------------------------------------------------------- |
| 27 | + |
| 28 | +@dataclass |
| 29 | +class SessionNode: |
| 30 | + meta: SessionMeta |
| 31 | + events: list[TraceEvent] |
| 32 | + children: list[SessionNode] = field(default_factory=list) |
| 33 | + |
| 34 | + @property |
| 35 | + def depth(self) -> int: |
| 36 | + return self.meta.depth |
| 37 | + |
| 38 | + |
| 39 | +@dataclass |
| 40 | +class AggregatedStats: |
| 41 | + session_count: int = 0 |
| 42 | + tool_calls: int = 0 |
| 43 | + llm_requests: int = 0 |
| 44 | + errors: int = 0 |
| 45 | + total_tokens: int = 0 |
| 46 | + total_duration_ms: float = 0 |
| 47 | + |
| 48 | + |
| 49 | +# --------------------------------------------------------------------------- |
| 50 | +# Tree building |
| 51 | +# --------------------------------------------------------------------------- |
| 52 | + |
| 53 | +def build_tree(store: TraceStore, root_session_id: str) -> SessionNode: |
| 54 | + """Build a SessionNode tree rooted at *root_session_id*. |
| 55 | +
|
| 56 | + Discovers child sessions by scanning all sessions for ones whose |
| 57 | + parent_session_id matches a session already in the tree. |
| 58 | + Depth is bounded by MAX_DEPTH. |
| 59 | + """ |
| 60 | + all_meta = store.list_sessions() |
| 61 | + |
| 62 | + # Index by session_id for fast lookup |
| 63 | + meta_by_id: dict[str, SessionMeta] = {m.session_id: m for m in all_meta} |
| 64 | + |
| 65 | + # Index children by parent_session_id |
| 66 | + children_of: dict[str, list[SessionMeta]] = {} |
| 67 | + for m in all_meta: |
| 68 | + if m.parent_session_id: |
| 69 | + children_of.setdefault(m.parent_session_id, []).append(m) |
| 70 | + |
| 71 | + def _build(session_id: str, current_depth: int) -> SessionNode: |
| 72 | + if session_id not in meta_by_id: |
| 73 | + raise KeyError(f"Session not found in store: {session_id}") |
| 74 | + meta = meta_by_id[session_id] |
| 75 | + events = store.load_events(session_id) |
| 76 | + node = SessionNode(meta=meta, events=events) |
| 77 | + |
| 78 | + if current_depth < MAX_DEPTH: |
| 79 | + for child_meta in sorted( |
| 80 | + children_of.get(session_id, []), |
| 81 | + key=lambda m: m.started_at, |
| 82 | + ): |
| 83 | + node.children.append(_build(child_meta.session_id, current_depth + 1)) |
| 84 | + elif children_of.get(session_id): |
| 85 | + sys.stderr.write( |
| 86 | + f"agent-strace: subagent tree truncated at depth {MAX_DEPTH}" |
| 87 | + f" for session {session_id[:12]}\n" |
| 88 | + ) |
| 89 | + |
| 90 | + return node |
| 91 | + |
| 92 | + return _build(root_session_id, 0) |
| 93 | + |
| 94 | + |
| 95 | +def aggregate_stats(node: SessionNode) -> AggregatedStats: |
| 96 | + """Recursively aggregate stats across the full session tree.""" |
| 97 | + stats = AggregatedStats( |
| 98 | + session_count=1, |
| 99 | + tool_calls=node.meta.tool_calls, |
| 100 | + llm_requests=node.meta.llm_requests, |
| 101 | + errors=node.meta.errors, |
| 102 | + total_tokens=node.meta.total_tokens, |
| 103 | + # Duration is wall-clock: subagents run within parent time, so we take |
| 104 | + # the max of the root's own duration and each child subtree's duration. |
| 105 | + total_duration_ms=node.meta.total_duration_ms, |
| 106 | + ) |
| 107 | + for child in node.children: |
| 108 | + child_stats = aggregate_stats(child) |
| 109 | + stats.session_count += child_stats.session_count |
| 110 | + stats.tool_calls += child_stats.tool_calls |
| 111 | + stats.llm_requests += child_stats.llm_requests |
| 112 | + stats.errors += child_stats.errors |
| 113 | + stats.total_tokens += child_stats.total_tokens |
| 114 | + # Duration is wall-clock: subagents run within parent time, so take |
| 115 | + # the running max across all children (not the root's fixed value). |
| 116 | + stats.total_duration_ms = max( |
| 117 | + stats.total_duration_ms, child_stats.total_duration_ms |
| 118 | + ) |
| 119 | + return stats |
| 120 | + |
| 121 | + |
| 122 | +# --------------------------------------------------------------------------- |
| 123 | +# Tree-aware replay formatting |
| 124 | +# --------------------------------------------------------------------------- |
| 125 | + |
| 126 | +def _fmt_offset(base_ts: float, ts: float) -> str: |
| 127 | + offset = max(0.0, ts - base_ts) |
| 128 | + if offset < 60: |
| 129 | + return f"+{offset:5.2f}s" |
| 130 | + m = int(offset) // 60 |
| 131 | + s = offset % 60 |
| 132 | + return f"+{m}m{s:04.1f}s" |
| 133 | + |
| 134 | + |
| 135 | +def _indent(depth: int, last_child: bool = False) -> str: |
| 136 | + if depth == 0: |
| 137 | + return "" |
| 138 | + connector = "└─ " if last_child else "├─ " |
| 139 | + return "│ " * (depth - 1) + connector |
| 140 | + |
| 141 | + |
| 142 | +def format_tree( |
| 143 | + node: SessionNode, |
| 144 | + base_ts: float | None = None, |
| 145 | + out: TextIO = sys.stdout, |
| 146 | + expand: bool = True, |
| 147 | + last_child: bool = False, |
| 148 | +) -> None: |
| 149 | + """Render the session tree to *out*. |
| 150 | +
|
| 151 | + Parameters |
| 152 | + ---------- |
| 153 | + node : SessionNode |
| 154 | + Root of the tree to render. |
| 155 | + base_ts : float, optional |
| 156 | + Timestamp origin for relative offsets. Defaults to root session start. |
| 157 | + expand : bool |
| 158 | + If True, inline subagent events under their parent tool_call. |
| 159 | + last_child : bool |
| 160 | + Whether this node is the last child of its parent (affects tree chars). |
| 161 | + """ |
| 162 | + if base_ts is None: |
| 163 | + base_ts = node.meta.started_at |
| 164 | + |
| 165 | + indent = _indent(node.depth, last_child=last_child) |
| 166 | + w = out.write |
| 167 | + |
| 168 | + # Session header |
| 169 | + w(f"{indent}▶ session_start {node.meta.session_id[:12]}" |
| 170 | + f" agent={node.meta.agent_name or 'unknown'}" |
| 171 | + f" depth={node.depth}\n") |
| 172 | + |
| 173 | + # Build a lookup of child sessions by the parent_event_id that spawned them |
| 174 | + children_by_event: dict[str, SessionNode] = { |
| 175 | + c.meta.parent_event_id: c for c in node.children if c.meta.parent_event_id |
| 176 | + } |
| 177 | + |
| 178 | + for event in node.events: |
| 179 | + ts_str = _fmt_offset(base_ts, event.timestamp) |
| 180 | + etype = event.event_type.value |
| 181 | + |
| 182 | + if event.event_type == EventType.TOOL_CALL: |
| 183 | + tool_name = event.data.get("tool_name", "?") |
| 184 | + args = event.data.get("arguments", {}) |
| 185 | + detail = "" |
| 186 | + if tool_name.lower() == "bash": |
| 187 | + cmd = str(args.get("command", "")) |
| 188 | + detail = f" $ {cmd[:80]}{'...' if len(cmd) > 80 else ''}" |
| 189 | + elif tool_name.lower() in ("read", "write", "edit"): |
| 190 | + detail = f" {args.get('file_path', '')}" |
| 191 | + elif tool_name.lower() == "agent": |
| 192 | + prompt = str(args.get("prompt", "")) |
| 193 | + detail = f" \"{prompt[:60]}{'...' if len(prompt) > 60 else ''}\"" |
| 194 | + |
| 195 | + subagent_tag = "" |
| 196 | + if event.data.get("is_sidechain"): |
| 197 | + subagent_tag = " [sidechain]" |
| 198 | + if event.data.get("subagent_type"): |
| 199 | + subagent_tag += f" [{event.data['subagent_type']}]" |
| 200 | + |
| 201 | + w(f"{indent}{ts_str} → tool_call {tool_name}{subagent_tag}{detail}\n") |
| 202 | + |
| 203 | + # Inline expand subagent if this tool_call spawned one |
| 204 | + if expand and event.event_id in children_by_event: |
| 205 | + child = children_by_event[event.event_id] |
| 206 | + is_last = child == node.children[-1] if node.children else True |
| 207 | + format_tree(child, base_ts=base_ts, out=out, expand=expand, |
| 208 | + last_child=is_last) |
| 209 | + |
| 210 | + elif event.event_type == EventType.TOOL_RESULT: |
| 211 | + preview = (event.data.get("result", "") or |
| 212 | + event.data.get("content_preview", ""))[:80] |
| 213 | + w(f"{indent}{ts_str} ← tool_result" |
| 214 | + f"{' ' + preview if preview else ''}\n") |
| 215 | + |
| 216 | + elif event.event_type == EventType.ERROR: |
| 217 | + msg = (event.data.get("message", "") or |
| 218 | + event.data.get("error", ""))[:80] |
| 219 | + w(f"{indent}{ts_str} ✗ error {msg}\n") |
| 220 | + |
| 221 | + elif event.event_type == EventType.USER_PROMPT: |
| 222 | + prompt = event.data.get("prompt", "")[:80] |
| 223 | + w(f"{indent}{ts_str} 👤 \"{prompt}\"\n") |
| 224 | + |
| 225 | + elif event.event_type == EventType.ASSISTANT_RESPONSE: |
| 226 | + text = event.data.get("text", "")[:80] |
| 227 | + w(f"{indent}{ts_str} 🤖 \"{text}\"\n") |
| 228 | + |
| 229 | + elif event.event_type == EventType.SESSION_END: |
| 230 | + w(f"{indent}{ts_str} ■ session_end\n") |
| 231 | + |
| 232 | + w("\n") |
| 233 | + |
| 234 | + |
| 235 | +def format_tree_summary( |
| 236 | + node: SessionNode, |
| 237 | + out: TextIO = sys.stdout, |
| 238 | + last_child: bool = False, |
| 239 | +) -> None: |
| 240 | + """Print a compact tree structure showing session hierarchy.""" |
| 241 | + w = out.write |
| 242 | + indent = _indent(node.depth, last_child=last_child) |
| 243 | + |
| 244 | + duration = node.meta.total_duration_ms / 1000 if node.meta.total_duration_ms else 0 |
| 245 | + w(f"{indent}{node.meta.session_id[:12]}" |
| 246 | + f" {duration:.1f}s" |
| 247 | + f" {node.meta.tool_calls} tools" |
| 248 | + f" {node.meta.total_tokens:,} tokens" |
| 249 | + f"{' ✗ ' + str(node.meta.errors) + ' errors' if node.meta.errors else ''}\n") |
| 250 | + |
| 251 | + for i, child in enumerate(node.children): |
| 252 | + format_tree_summary(child, out=out, last_child=(i == len(node.children) - 1)) |
| 253 | + |
| 254 | + |
| 255 | +# --------------------------------------------------------------------------- |
| 256 | +# CLI handlers |
| 257 | +# --------------------------------------------------------------------------- |
| 258 | + |
| 259 | +def cmd_replay_tree(args: argparse.Namespace) -> int: |
| 260 | + store = TraceStore(args.trace_dir) |
| 261 | + |
| 262 | + session_id = args.session_id |
| 263 | + if not session_id: |
| 264 | + session_id = store.get_latest_session_id() |
| 265 | + if not session_id: |
| 266 | + sys.stderr.write("No sessions found.\n") |
| 267 | + return 1 |
| 268 | + full_id = store.find_session(session_id) |
| 269 | + if not full_id: |
| 270 | + sys.stderr.write(f"Session not found: {session_id}\n") |
| 271 | + return 1 |
| 272 | + |
| 273 | + tree = build_tree(store, full_id) |
| 274 | + expand = not getattr(args, "tree_only", False) |
| 275 | + |
| 276 | + if getattr(args, "tree", False): |
| 277 | + stats = aggregate_stats(tree) |
| 278 | + sys.stdout.write(f"\nSession tree for {full_id[:12]}\n\n") |
| 279 | + format_tree_summary(tree) |
| 280 | + sys.stdout.write( |
| 281 | + f"\nTotal: {stats.session_count} sessions, " |
| 282 | + f"{stats.tool_calls} tool calls, " |
| 283 | + f"{stats.llm_requests} LLM requests, " |
| 284 | + f"{stats.total_tokens:,} tokens" |
| 285 | + f"{', ' + str(stats.errors) + ' errors' if stats.errors else ''}\n\n" |
| 286 | + ) |
| 287 | + else: |
| 288 | + format_tree(tree, expand=expand) |
| 289 | + |
| 290 | + return 0 |
| 291 | + |
| 292 | + |
| 293 | +def cmd_stats_tree(args: argparse.Namespace) -> int: |
| 294 | + store = TraceStore(args.trace_dir) |
| 295 | + |
| 296 | + session_id = args.session_id |
| 297 | + if not session_id: |
| 298 | + session_id = store.get_latest_session_id() |
| 299 | + if not session_id: |
| 300 | + sys.stderr.write("No sessions found.\n") |
| 301 | + return 1 |
| 302 | + full_id = store.find_session(session_id) |
| 303 | + if not full_id: |
| 304 | + sys.stderr.write(f"Session not found: {session_id}\n") |
| 305 | + return 1 |
| 306 | + |
| 307 | + tree = build_tree(store, full_id) |
| 308 | + stats = aggregate_stats(tree) |
| 309 | + |
| 310 | + sys.stdout.write(f"\nAggregated stats for {full_id[:12]} (including subagents)\n\n") |
| 311 | + sys.stdout.write(f" Sessions: {stats.session_count}\n") |
| 312 | + sys.stdout.write(f" Tool calls: {stats.tool_calls}\n") |
| 313 | + sys.stdout.write(f" LLM requests: {stats.llm_requests}\n") |
| 314 | + sys.stdout.write(f" Total tokens: {stats.total_tokens:,}\n") |
| 315 | + sys.stdout.write(f" Errors: {stats.errors}\n") |
| 316 | + sys.stdout.write(f" Duration: {stats.total_duration_ms / 1000:.1f}s\n\n") |
| 317 | + |
| 318 | + if tree.children: |
| 319 | + sys.stdout.write("Session tree:\n\n") |
| 320 | + format_tree_summary(tree) |
| 321 | + sys.stdout.write("\n") |
| 322 | + |
| 323 | + return 0 |
0 commit comments