@tracemarketplace/shared 0.0.6 → 0.0.9

package/src/extractors/codex.ts

@@ -5,6 +5,22 @@ import type {
   Turn,
   ContentBlock,
 } from "../types.js";
+import {
+  collectTraceMetrics,
+  createPassiveEnvState,
+  extractTextFragments,
+  pushUniqueTextBlock,
+} from "./common.js";
+
+function toNumber(value: unknown): number | null {
+  return typeof value === "number" && Number.isFinite(value) ? value : null;
+}
+
+function extractExitCode(output: string | undefined): number | null {
+  if (!output) return null;
+  const match = output.match(/Process exited with code (\d+)/);
+  return match ? Number.parseInt(match[1], 10) : null;
+}
 
 export async function extractCodex(
   rolloutFileBuffer: Buffer,
@@ -29,6 +45,10 @@ export async function extractCodex(
   let currentBlocks: ContentBlock[] = [];
   let taskTimestamp: string | null = null;
   let currentModel: string | null = null;
+  let currentTaskLastTimestamp: string | null = null;
+  let totalInputTokens: number | null = null;
+  let totalOutputTokens: number | null = null;
+  let totalCacheReadTokens: number | null = null;
 
   for (const event of events) {
     const type = event["type"] as string;
@@ -41,6 +61,7 @@ export async function extractCodex(
       currentUserMessage = null;
       currentBlocks = [];
       taskTimestamp = timestamp;
+      currentTaskLastTimestamp = timestamp;
       continue;
     }
 
@@ -61,10 +82,7 @@ export async function extractCodex(
 
         const lastMsg = payload["last_agent_message"] as string | undefined;
         if (lastMsg) {
-          const alreadyCaptured = currentBlocks.some(
-            b => b.type === "text" && (b as { type: "text"; text: string }).text === lastMsg
-          );
-          if (!alreadyCaptured) currentBlocks.push({ type: "text", text: lastMsg });
+          pushUniqueTextBlock(currentBlocks, "text", lastMsg);
         }
 
         if (currentBlocks.length > 0) {
@@ -72,7 +90,7 @@ export async function extractCodex(
             turn_id: currentTurnId ?? randomUUID(),
             parent_turn_id: currentUserMessage ? "user_" + currentTurnId : null,
             role: "assistant",
-            timestamp: taskTimestamp,
+            timestamp: currentTaskLastTimestamp ?? timestamp ?? taskTimestamp,
             content: currentBlocks,
             model: currentModel,
             usage: null,
@@ -86,11 +104,14 @@ export async function extractCodex(
       currentBlocks = [];
       taskTimestamp = null;
       currentModel = null;
+      currentTaskLastTimestamp = null;
       continue;
     }
 
     if (!inTask) continue;
 
+    if (timestamp) currentTaskLastTimestamp = timestamp;
+
     if (type === "turn_context") {
       currentModel = (payload["model"] as string | undefined) ?? null;
       continue;
@@ -102,14 +123,42 @@ export async function extractCodex(
           currentUserMessage = (payload["message"] as string | undefined) ?? "";
           break;
         case "agent_reasoning":
-          if (payload["text"]) currentBlocks.push({ type: "thinking", text: payload["text"] as string });
+          pushUniqueTextBlock(currentBlocks, "thinking", payload["text"] as string | undefined);
           break;
         case "agent_message":
-          if (payload["message"]) currentBlocks.push({ type: "text", text: payload["message"] as string });
+          pushUniqueTextBlock(currentBlocks, "text", payload["message"] as string | undefined);
+          break;
+        case "token_count": {
+          const info = (payload["info"] ?? {}) as Record<string, unknown>;
+          const usage = (info["total_token_usage"] ?? {}) as Record<string, unknown>;
+          totalInputTokens = toNumber(usage["input_tokens"]) ?? totalInputTokens;
+          totalOutputTokens = toNumber(usage["output_tokens"]) ?? totalOutputTokens;
+          totalCacheReadTokens =
+            toNumber(usage["cached_input_tokens"]) ?? totalCacheReadTokens;
           break;
+        }
       }
     } else if (type === "response_item") {
       switch (payload["type"] as string) {
+        case "message": {
+          const role = payload["role"] as string | undefined;
+          const texts = extractTextFragments(payload["content"]);
+          if (role === "assistant") {
+            for (const text of texts) {
+              pushUniqueTextBlock(currentBlocks, "text", text);
+            }
+          } else if (role === "user" && !currentUserMessage) {
+            currentUserMessage = texts.join("\n\n") || currentUserMessage;
+          }
+          break;
+        }
+        case "reasoning": {
+          const reasoningText =
+            extractTextFragments(payload["content"]).join("\n\n") ||
+            extractTextFragments(payload["summary"]).join("\n\n");
+          pushUniqueTextBlock(currentBlocks, "thinking", reasoningText);
+          break;
+        }
         case "function_call": {
           const callId = (payload["call_id"] as string | undefined) ?? randomUUID();
           let toolInput: Record<string, unknown> = {};
@@ -126,15 +175,13 @@ export async function extractCodex(
         }
         case "function_call_output": {
           const output = payload["output"] as string | undefined;
-          const isError = typeof output === "string"
-            && output.includes("Process exited with code")
-            && !output.includes("code 0");
+          const exitCode = extractExitCode(output);
           currentBlocks.push({
             type: "tool_result",
             tool_call_id: (payload["call_id"] as string | undefined) ?? "",
-            is_error: isError,
+            is_error: exitCode !== null ? exitCode !== 0 : false,
             result_content: output ?? null,
-            exit_code: null,
+            exit_code: exitCode,
           });
           break;
         }
@@ -152,15 +199,7 @@ export async function extractCodex(
     }
   }
 
-  const allBlocks = turns.flatMap(t => t.content);
-  const toolCallCount = allBlocks.filter(b => b.type === "tool_use").length;
-  const hasFileChanges = allBlocks.some(
-    b => b.type === "tool_use" && ["write_file", "file_change", "create_file"].includes((b as { type: "tool_use"; tool_name: string }).tool_name)
-  );
-  const hasShellCommands = allBlocks.some(
-    b => b.type === "tool_use" && ["exec_command", "bash", "shell"].includes((b as { type: "tool_use"; tool_name: string }).tool_name)
-  );
-  const hasThinking = allBlocks.some(b => b.type === "thinking");
+  const metrics = collectTraceMetrics(turns);
   const startedAt = (sessionMetaPayload["timestamp"] as string | undefined) ?? events[0]?.["timestamp"] as string ?? new Date().toISOString();
   const endedAt = events[events.length - 1]?.["timestamp"] as string ?? new Date().toISOString();
 
@@ -179,24 +218,16 @@ export async function extractCodex(
     ended_at: endedAt,
     turns,
     turn_count: turns.length,
-    tool_call_count: toolCallCount,
-    has_tool_calls: toolCallCount > 0,
-    has_thinking_blocks: hasThinking,
-    has_file_changes: hasFileChanges,
-    has_shell_commands: hasShellCommands,
-    total_input_tokens: null,
-    total_output_tokens: null,
-    total_cache_read_tokens: null,
+    tool_call_count: metrics.toolCallCount,
+    has_tool_calls: metrics.toolCallCount > 0,
+    has_thinking_blocks: metrics.hasThinkingBlocks,
+    has_file_changes: metrics.hasFileChanges,
+    has_shell_commands: metrics.hasShellCommands,
+    total_input_tokens: totalInputTokens,
+    total_output_tokens: totalOutputTokens,
+    total_cache_read_tokens: totalCacheReadTokens,
     content_fidelity: "full",
-    env_state: {
-      git_branch: null,
-      inferred_file_tree: null,
-      inferred_changed_files: null,
-      inferred_error_files: null,
-      shell_exit_codes: null,
-      open_files_in_editor: null,
-      extraction_method: "passive",
-    },
+    env_state: createPassiveEnvState(),
     score: null,
     raw_r2_key: "",
     normalized_r2_key: "",

package/src/extractors/common.ts

@@ -0,0 +1,139 @@
+import type { ContentBlock, EnvState, Turn } from "../types.js";
+
+const SHELL_TOOL_NAMES = ["exec_command", "bash", "shell", "write_stdin"];
+const FILE_MUTATION_TOOL_NAMES = [
+  "apply_patch",
+  "write_file",
+  "create_file",
+  "delete_file",
+  "rename_file",
+  "move_file",
+  "file_change",
+  "edit",
+  "edit_file",
+  "multiedit",
+  "write",
+];
+
+function normalizeToolName(toolName: string): string {
+  return toolName.trim().toLowerCase();
+}
+
+function toolNameMatches(toolName: string, candidate: string): boolean {
+  const normalized = normalizeToolName(toolName);
+  return normalized === candidate || normalized.endsWith(`.${candidate}`);
+}
+
+export function isShellToolName(toolName: string): boolean {
+  return SHELL_TOOL_NAMES.some((candidate) => toolNameMatches(toolName, candidate));
+}
+
+export function isFileMutationTool(
+  toolName: string,
+  toolInput: Record<string, unknown>,
+): boolean {
+  if (
+    FILE_MUTATION_TOOL_NAMES.some((candidate) => toolNameMatches(toolName, candidate))
+  ) {
+    return true;
+  }
+
+  if (!isShellToolName(toolName)) return false;
+
+  const command =
+    typeof toolInput.cmd === "string"
+      ? toolInput.cmd
+      : typeof toolInput.command === "string"
+        ? toolInput.command
+        : null;
+
+  return command !== null && /\bapply_patch\b|\bsed\s+-i\b|\bperl\s+-pi\b/.test(command);
+}
+
+export function collectTraceMetrics(turns: Turn[]) {
+  const allBlocks = turns.flatMap((turn) => turn.content);
+  const toolUses = allBlocks.filter(
+    (block): block is Extract<ContentBlock, { type: "tool_use" }> =>
+      block.type === "tool_use",
+  );
+
+  return {
+    toolCallCount: toolUses.length,
+    hasFileChanges: toolUses.some((block) =>
+      isFileMutationTool(block.tool_name, block.tool_input),
+    ),
+    hasShellCommands: toolUses.some((block) => isShellToolName(block.tool_name)),
+    hasThinkingBlocks: allBlocks.some((block) => block.type === "thinking"),
+  };
+}
+
+export function createPassiveEnvState(
+  overrides: Partial<EnvState> = {},
+): EnvState {
+  return {
+    git_branch: null,
+    inferred_file_tree: null,
+    inferred_changed_files: null,
+    inferred_error_files: null,
+    shell_exit_codes: null,
+    open_files_in_editor: null,
+    extraction_method: "passive",
+    ...overrides,
+  };
+}
+
+export function extractTextFragments(value: unknown): string[] {
+  if (typeof value === "string") {
+    return value.trim() ? [value] : [];
+  }
+
+  if (Array.isArray(value)) {
+    return value.flatMap((item) => extractTextFragments(item));
+  }
+
+  if (!value || typeof value !== "object") {
+    return [];
+  }
+
+  const record = value as Record<string, unknown>;
+
+  if (typeof record.text === "string" && record.text.trim()) {
+    return [record.text];
+  }
+
+  const nested: string[] = [];
+  for (const key of ["content", "summary", "thinking", "message"]) {
+    if (key in record) {
+      nested.push(...extractTextFragments(record[key]));
+    }
+  }
+
+  return nested;
+}
+
+export function pushUniqueTextBlock(
+  blocks: ContentBlock[],
+  type: "text" | "thinking",
+  text: string | null | undefined,
+): void {
+  const normalized = text?.trim();
+  if (!normalized) return;
+
+  const exists = blocks.some(
+    (block) => block.type === type && block.text === normalized,
+  );
+  if (!exists) {
+    blocks.push({ type, text: normalized });
+  }
+}
+
+export function normalizeTimestamp(value: unknown): string | null {
+  if (typeof value === "string" && value.trim()) return value;
+
+  if (typeof value === "number" && Number.isFinite(value)) {
+    const millis = value > 1_000_000_000_000 ? value : value * 1000;
+    return new Date(millis).toISOString();
+  }
+
+  return null;
+}