@townco/debugger 0.1.66 → 0.1.68

package/src/lib/segmentation.ts

@@ -0,0 +1,391 @@
+import type { ConversationTrace, Log, Span } from "../types";
+
+// ============================================================================
+// Types
+// ============================================================================
+
+export interface ConversationItem {
+  id: string;
+  type: "user" | "agent" | "thinking" | "tool_call";
+  content: string;
+  timestamp: number;
+  spanId?: string;
+  toolName?: string;
+  toolInput?: unknown;
+  toolOutput?: unknown;
+  /** Label shown above the content (e.g., "Agent", "Thinking stream") */
+  label?: string;
+  /** If this item is from a subagent, the span_id of the parent Task span */
+  subagentSpanId?: string;
+  /** Short identifier for the subagent (last 4 chars of span_id) */
+  subagentId?: string;
+}
+
+export interface Segment {
+  id: string;
+  traceId: string;
+  startTime: number;
+  endTime: number;
+  conversationItems: ConversationItem[];
+  spans: Span[];
+  logs: Log[];
+}
+
+export interface TraceData {
+  trace: ConversationTrace;
+  spans: Span[];
+  logs: Log[];
+}
+
+// ============================================================================
+// Utility Functions
+// ============================================================================
+
+function parseAttributes(attrs: string | null): Record<string, unknown> {
+  if (!attrs) return {};
+  try {
+    return JSON.parse(attrs);
+  } catch {
+    return {};
+  }
+}
+
+function formatTimestamp(timestampNano: number): string {
+  const date = new Date(timestampNano / 1_000_000);
+  return date.toLocaleTimeString("en-US", {
+    hour: "2-digit",
+    minute: "2-digit",
+    second: "2-digit",
+    hour12: false,
+  });
+}
+
+/**
+ * Extract thinking/reasoning content from a chat span if available
+ */
+function extractThinkingContent(span: Span): string | null {
+  const attrs = parseAttributes(span.attributes);
+  const outputMessages = attrs["gen_ai.output.messages"];
+
+  if (!outputMessages) return null;
+
+  try {
+    const messages =
+      typeof outputMessages === "string"
+        ? JSON.parse(outputMessages)
+        : outputMessages;
+
+    if (!Array.isArray(messages)) return null;
+
+    // Look for thinking blocks in the content
+    for (const msg of messages) {
+      if (msg.role === "assistant" || msg.role === "ai") {
+        const content = msg.content;
+        if (Array.isArray(content)) {
+          for (const block of content) {
+            if (block.type === "thinking" && block.thinking) {
+              return block.thinking;
+            }
+          }
+        }
+      }
+    }
+  } catch {
+    // Ignore parse errors
+  }
+
+  return null;
+}
+
+/**
+ * Check if a span is a chat/LLM call span
+ */
+function isChatSpan(span: Span): boolean {
+  const attrs = parseAttributes(span.attributes);
+  return span.name.startsWith("chat") && "gen_ai.input.messages" in attrs;
+}
+
+/**
+ * Check if a span is a subagent (Task) span
+ *
+ * A span is considered a subagent if:
+ * - It's a tool call span (agent.tool_call)
+ * - AND the tool name matches one of these patterns:
+ *   - Exactly "Task" (the built-in subagent tool)
+ *   - Contains "subagent" (case-insensitive)
+ */
+function isSubagentSpan(span: Span): boolean {
+  if (span.name !== "agent.tool_call") return false;
+  const attrs = parseAttributes(span.attributes);
+  const toolName = attrs["tool.name"];
+
+  if (typeof toolName !== "string") return false;
+
+  // Check for Task or any tool name containing "subagent" (case-insensitive)
+  const isSubagent =
+    toolName === "Task" || toolName.toLowerCase().includes("subagent");
+
+  return isSubagent;
+}
+
+/**
+ * Get a short identifier from a span_id (last 4 characters)
+ */
+function getShortId(spanId: string): string {
+  return spanId.slice(-4);
+}
+
+interface SubagentInfo {
+  spanId: string;
+  shortId: string;
+}
+
+/**
+ * Build a map of span_id -> subagent info for all spans that are descendants of a Task span
+ */
+function buildSubagentMap(spans: Span[]): Map<string, SubagentInfo> {
+  const subagentMap = new Map<string, SubagentInfo>();
+  const spanById = new Map(spans.map((s) => [s.span_id, s]));
+
+  // Find all Task (subagent) spans
+  const taskSpans = spans.filter(isSubagentSpan);
+
+  console.log("[DEBUG] Found subagent tool call spans:", {
+    count: taskSpans.length,
+    spans: taskSpans.map((s) => {
+      const attrs = parseAttributes(s.attributes);
+      return {
+        spanId: s.span_id.slice(-4),
+        toolName: attrs["tool.name"],
+        spanName: s.name,
+      };
+    }),
+  });
+
+  // For each Task span, mark it and all its descendants
+  for (const taskSpan of taskSpans) {
+    const subagentInfo: SubagentInfo = {
+      spanId: taskSpan.span_id,
+      shortId: getShortId(taskSpan.span_id),
+    };
+
+    // Mark the Task span itself
+    subagentMap.set(taskSpan.span_id, subagentInfo);
+
+    // Track descendants for logging
+    const descendants: string[] = [];
+
+    // Find all descendants using BFS
+    const queue: string[] = [taskSpan.span_id];
+    while (queue.length > 0) {
+      const currentId = queue.shift();
+      if (!currentId) break;
+      // Find children of current span
+      for (const span of spans) {
+        if (span.parent_span_id === currentId) {
+          subagentMap.set(span.span_id, subagentInfo);
+          queue.push(span.span_id);
+          descendants.push(span.span_id);
+        }
+      }
+    }
+
+    console.log("[DEBUG] Subagent descendant tree:", {
+      rootSpanId: taskSpan.span_id.slice(-4),
+      shortId: subagentInfo.shortId,
+      descendantCount: descendants.length,
+      descendants: descendants.map((id) => ({
+        spanId: id.slice(-4),
+        spanName: spanById.get(id)?.name || "unknown",
+      })),
+    });
+  }
+
+  return subagentMap;
+}
+
+// ============================================================================
+// Segmentation Function
+// ============================================================================
+
+/**
+ * Per-Turn Segmentation: Each user message + full agent response = one segment
+ * Spans are grouped by their trace_id, preserving full hierarchy
+ */
+export function segmentByTurn(
+  traces: ConversationTrace[],
+  traceDataMap: Map<string, TraceData>,
+): Segment[] {
+  const segments: Segment[] = [];
+
+  for (const trace of traces) {
+    const data = traceDataMap.get(trace.trace_id);
+    if (!data) continue;
+
+    // Build subagent map for this trace's spans
+    const subagentMap = buildSubagentMap(data.spans);
+
+    // Debug: Log subagent map contents
+    if (subagentMap.size > 0) {
+      console.log("[DEBUG] Subagent map for trace:", trace.trace_id.slice(-4), {
+        subagentCount: subagentMap.size,
+        subagents: Array.from(subagentMap.entries()).map(([spanId, info]) => ({
+          spanId: spanId.slice(-4),
+          subagentSpanId: info.spanId.slice(-4),
+          shortId: info.shortId,
+        })),
+      });
+    }
+
+    const conversationItems: ConversationItem[] = [];
+
+    // Add user message if available
+    if (trace.userInput) {
+      conversationItems.push({
+        id: `${trace.trace_id}-user`,
+        type: "user",
+        content: trace.userInput,
+        timestamp: trace.start_time_unix_nano,
+      });
+    }
+
+    // Add agent messages (chat and tool calls)
+    for (const msg of trace.agentMessages) {
+      // Check if this message is from a subagent
+      const subagentInfo = msg.spanId ? subagentMap.get(msg.spanId) : undefined;
+
+      // Debug: Log message processing
+      if (msg.spanId) {
+        console.log("[DEBUG] Processing agent message:", {
+          type: msg.type,
+          spanId: msg.spanId.slice(-4),
+          hasSubagentInfo: !!subagentInfo,
+          subagentShortId: subagentInfo?.shortId,
+          contentPreview: msg.content.slice(0, 50),
+        });
+      }
+
+      if (msg.type === "chat") {
+        const item: ConversationItem = {
+          id: `${trace.trace_id}-agent-${msg.spanId}`,
+          type: "agent",
+          content: msg.content,
+          timestamp: msg.timestamp,
+          spanId: msg.spanId,
+          label: subagentInfo ? `Subagent .${subagentInfo.shortId}` : "Agent",
+        };
+        if (subagentInfo) {
+          item.subagentSpanId = subagentInfo.spanId;
+          item.subagentId = subagentInfo.shortId;
+        }
+        conversationItems.push(item);
+      } else if (msg.type === "tool_call" && msg.toolName) {
+        const item: ConversationItem = {
+          id: `${trace.trace_id}-tool-${msg.spanId}`,
+          type: "tool_call",
+          content: msg.toolName || msg.content,
+          timestamp: msg.timestamp,
+          spanId: msg.spanId,
+          toolName: msg.toolName,
+          toolInput: msg.toolInput,
+          toolOutput: msg.toolOutput,
+          label: "Tool Call",
+        };
+        if (subagentInfo) {
+          item.subagentSpanId = subagentInfo.spanId;
+          item.subagentId = subagentInfo.shortId;
+        }
+        conversationItems.push(item);
+      }
+    }
+
+    // Check for thinking content in chat spans
+    for (const span of data.spans) {
+      if (isChatSpan(span)) {
+        const thinking = extractThinkingContent(span);
+        if (thinking) {
+          // Insert thinking before the corresponding chat message
+          const chatItemIndex = conversationItems.findIndex(
+            (item) => item.spanId === span.span_id && item.type === "agent",
+          );
+          if (chatItemIndex > 0) {
+            const subagentInfo = subagentMap.get(span.span_id);
+            const item: ConversationItem = {
+              id: `${trace.trace_id}-thinking-${span.span_id}`,
+              type: "thinking",
+              content: thinking,
+              timestamp: span.start_time_unix_nano,
+              spanId: span.span_id,
+              label: "Thinking stream",
+            };
+            if (subagentInfo) {
+              item.subagentSpanId = subagentInfo.spanId;
+              item.subagentId = subagentInfo.shortId;
+            }
+            conversationItems.splice(chatItemIndex, 0, item);
+          }
+        }
+      }
+    }
+
+    // Sort by timestamp
+    conversationItems.sort((a, b) => a.timestamp - b.timestamp);
+
+    const startTime =
+      data.spans.length > 0
+        ? Math.min(...data.spans.map((s) => s.start_time_unix_nano))
+        : trace.start_time_unix_nano;
+    const endTime =
+      data.spans.length > 0
+        ? Math.max(...data.spans.map((s) => s.end_time_unix_nano))
+        : trace.start_time_unix_nano;
+
+    segments.push({
+      id: trace.trace_id,
+      traceId: trace.trace_id,
+      startTime,
+      endTime,
+      conversationItems,
+      spans: data.spans,
+      logs: data.logs,
+    });
+  }
+
+  return segments;
+}
+
+// ============================================================================
+// Span Utilities
+// ============================================================================
+
+/**
+ * Get token count for a span
+ */
+export function getSpanTokens(span: Span): number {
+  const attrs = parseAttributes(span.attributes);
+  const inputTokens = (attrs["gen_ai.usage.input_tokens"] as number) || 0;
+  const outputTokens = (attrs["gen_ai.usage.output_tokens"] as number) || 0;
+  return inputTokens + outputTokens;
+}
+
+/**
+ * Get duration in milliseconds for a span
+ */
+export function getSpanDurationMs(span: Span): number {
+  return (span.end_time_unix_nano - span.start_time_unix_nano) / 1_000_000;
+}
+
+/**
+ * Format duration for display
+ */
+export function formatDuration(durationMs: number): string {
+  if (durationMs < 1) {
+    return `${(durationMs * 1000).toFixed(0)}μs`;
+  }
+  if (durationMs < 1000) {
+    return `${durationMs.toFixed(0)}ms`;
+  }
+  return `${(durationMs / 1000).toFixed(1)}s`;
+}
+
+export { formatTimestamp };