From 08cc8f2281f76e251348807fe3444dadc8eb0131 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Sat, 10 Jan 2026 22:03:42 +0000 Subject: [PATCH] refactor(agents): extract transcript repair module --- src/agents/pi-embedded-helpers.test.ts | 94 ---------- src/agents/pi-embedded-helpers.ts | 172 ------------------ src/agents/pi-embedded-runner.ts | 2 +- src/agents/session-transcript-repair.test.ts | 95 ++++++++++ src/agents/session-transcript-repair.ts | 173 +++++++++++++++++++ 5 files changed, 269 insertions(+), 267 deletions(-) create mode 100644 src/agents/session-transcript-repair.test.ts create mode 100644 src/agents/session-transcript-repair.ts diff --git a/src/agents/pi-embedded-helpers.test.ts b/src/agents/pi-embedded-helpers.test.ts index 7a0aaa8a9..0a4d5069f 100644 --- a/src/agents/pi-embedded-helpers.test.ts +++ b/src/agents/pi-embedded-helpers.test.ts @@ -14,7 +14,6 @@ import { sanitizeGoogleTurnOrdering, sanitizeSessionMessagesImages, sanitizeToolCallId, - sanitizeToolUseResultPairing, validateGeminiTurns, } from "./pi-embedded-helpers.js"; import { @@ -585,99 +584,6 @@ describe("sanitizeSessionMessagesImages", () => { }); }); -describe("sanitizeToolUseResultPairing", () => { - it("moves tool results directly after tool calls and inserts missing results", () => { - const input = [ - { - role: "assistant", - content: [ - { type: "toolCall", id: "call_1", name: "read", arguments: {} }, - { type: "toolCall", id: "call_2", name: "bash", arguments: {} }, - ], - }, - { role: "user", content: "user message that should come after tool use" }, - { - role: "toolResult", - toolCallId: "call_2", - toolName: "bash", - content: [{ type: "text", text: "ok" }], - isError: false, - }, - ] satisfies AgentMessage[]; - - const out = sanitizeToolUseResultPairing(input); - expect(out[0]?.role).toBe("assistant"); - expect(out[1]?.role).toBe("toolResult"); - expect((out[1] as { toolCallId?: string }).toolCallId).toBe("call_1"); - expect(out[2]?.role).toBe("toolResult"); - expect((out[2] as { toolCallId?: string }).toolCallId).toBe("call_2"); - expect(out[3]?.role).toBe("user"); - }); - - it("drops duplicate tool results for the same id within a span", () => { - const input = [ - { - role: "assistant", - content: [ - { type: "toolCall", id: "call_1", name: "read", arguments: {} }, - ], - }, - { - role: "toolResult", - toolCallId: "call_1", - toolName: "read", - content: [{ type: "text", text: "first" }], - isError: false, - }, - { - role: "toolResult", - toolCallId: "call_1", - toolName: "read", - content: [{ type: "text", text: "second" }], - isError: false, - }, - { role: "user", content: "ok" }, - ] satisfies AgentMessage[]; - - const out = sanitizeToolUseResultPairing(input); - expect(out.filter((m) => m.role === "toolResult")).toHaveLength(1); - }); - - it("drops duplicate tool results for the same id across the transcript", () => { - const input = [ - { - role: "assistant", - content: [ - { type: "toolCall", id: "call_1", name: "read", arguments: {} }, - ], - }, - { - role: "toolResult", - toolCallId: "call_1", - toolName: "read", - content: [{ type: "text", text: "first" }], - isError: false, - }, - { role: "assistant", content: [{ type: "text", text: "ok" }] }, - { - role: "toolResult", - toolCallId: "call_1", - toolName: "read", - content: [{ type: "text", text: "second (duplicate)" }], - isError: false, - }, - ] satisfies AgentMessage[]; - - const out = sanitizeToolUseResultPairing(input); - const results = out.filter((m) => m.role === "toolResult") as Array<{ - toolCallId?: string; - content?: unknown; - }>; - expect(results).toHaveLength(1); - expect(results[0]?.toolCallId).toBe("call_1"); - }); -}); - describe("normalizeTextForComparison", () => { it("lowercases text", () => { expect(normalizeTextForComparison("Hello World")).toBe("hello world"); diff --git a/src/agents/pi-embedded-helpers.ts b/src/agents/pi-embedded-helpers.ts index 173fd5137..579e3f18e 100644 --- a/src/agents/pi-embedded-helpers.ts +++ b/src/agents/pi-embedded-helpers.ts @@ -218,178 +218,6 @@ export async function sanitizeSessionMessagesImages( return out; } -type ToolCallLike = { - id: string; - name?: string; -}; - -function extractToolCallsFromAssistant( - msg: Extract, -): ToolCallLike[] { - const content = msg.content; - if (!Array.isArray(content)) return []; - - const toolCalls: ToolCallLike[] = []; - for (const block of content) { - if (!block || typeof block !== "object") continue; - const rec = block as { type?: unknown; id?: unknown; name?: unknown }; - if (typeof rec.id !== "string" || !rec.id) continue; - - if ( - rec.type === "toolCall" || - rec.type === "toolUse" || - rec.type === "functionCall" - ) { - toolCalls.push({ - id: rec.id, - name: typeof rec.name === "string" ? rec.name : undefined, - }); - } - } - return toolCalls; -} - -function extractToolResultId( - msg: Extract, -): string | null { - const toolCallId = (msg as { toolCallId?: unknown }).toolCallId; - if (typeof toolCallId === "string" && toolCallId) return toolCallId; - const toolUseId = (msg as { toolUseId?: unknown }).toolUseId; - if (typeof toolUseId === "string" && toolUseId) return toolUseId; - return null; -} - -function makeMissingToolResult(params: { - toolCallId: string; - toolName?: string; -}): Extract { - return { - role: "toolResult", - toolCallId: params.toolCallId, - toolName: params.toolName ?? "unknown", - content: [ - { - type: "text", - text: "[clawdbot] missing tool result in session history; inserted synthetic error result for transcript repair.", - }, - ], - isError: true, - timestamp: Date.now(), - } as Extract; -} - -export function sanitizeToolUseResultPairing( - messages: AgentMessage[], -): AgentMessage[] { - // Anthropic (and Cloud Code Assist) reject transcripts where assistant tool calls are not - // immediately followed by matching tool results. Session files can end up with results - // displaced (e.g. after user turns) or duplicated. Repair by: - // - moving matching toolResult messages directly after their assistant toolCall turn - // - inserting synthetic error toolResults for missing ids - // - dropping duplicate toolResults for the same id (anywhere in the transcript) - const out: AgentMessage[] = []; - const seenToolResultIds = new Set(); - - const pushToolResult = ( - msg: Extract, - ) => { - const id = extractToolResultId(msg); - if (id && seenToolResultIds.has(id)) return; - if (id) seenToolResultIds.add(id); - out.push(msg); - }; - - for (let i = 0; i < messages.length; i += 1) { - const msg = messages[i] as AgentMessage; - if (!msg || typeof msg !== "object") { - out.push(msg); - continue; - } - - const role = (msg as { role?: unknown }).role; - if (role !== "assistant") { - if (role === "toolResult") { - pushToolResult(msg as Extract); - } else { - out.push(msg); - } - continue; - } - - const assistant = msg as Extract; - const toolCalls = extractToolCallsFromAssistant(assistant); - if (toolCalls.length === 0) { - out.push(msg); - continue; - } - - const toolCallIds = new Set(toolCalls.map((t) => t.id)); - - const spanResultsById = new Map< - string, - Extract - >(); - const remainder: AgentMessage[] = []; - - let j = i + 1; - for (; j < messages.length; j += 1) { - const next = messages[j] as AgentMessage; - if (!next || typeof next !== "object") { - remainder.push(next); - continue; - } - - const nextRole = (next as { role?: unknown }).role; - if (nextRole === "assistant") break; - - if (nextRole === "toolResult") { - const toolResult = next as Extract< - AgentMessage, - { role: "toolResult" } - >; - const id = extractToolResultId(toolResult); - if (id && toolCallIds.has(id)) { - if (seenToolResultIds.has(id)) { - continue; - } - if (!spanResultsById.has(id)) { - spanResultsById.set(id, toolResult); - } - continue; - } - } - - remainder.push(next); - } - - out.push(msg); - - for (const call of toolCalls) { - const existing = spanResultsById.get(call.id); - pushToolResult( - existing ?? - makeMissingToolResult({ toolCallId: call.id, toolName: call.name }), - ); - } - - for (const rem of remainder) { - if (!rem || typeof rem !== "object") { - out.push(rem); - continue; - } - const remRole = (rem as { role?: unknown }).role; - if (remRole === "toolResult") { - pushToolResult(rem as Extract); - continue; - } - out.push(rem); - } - i = j - 1; - } - - return out; -} - const GOOGLE_TURN_ORDER_BOOTSTRAP_TEXT = "(session bootstrap)"; export function isGoogleModelApi(api?: string | null): boolean { diff --git a/src/agents/pi-embedded-runner.ts b/src/agents/pi-embedded-runner.ts index a9e70e4bd..35af105c9 100644 --- a/src/agents/pi-embedded-runner.ts +++ b/src/agents/pi-embedded-runner.ts @@ -88,7 +88,6 @@ import { pickFallbackThinkingLevel, sanitizeGoogleTurnOrdering, sanitizeSessionMessagesImages, - sanitizeToolUseResultPairing, validateGeminiTurns, } from "./pi-embedded-helpers.js"; import { @@ -106,6 +105,7 @@ import { makeToolPrunablePredicate } from "./pi-extensions/context-pruning/tools import { toToolDefinitions } from "./pi-tool-definition-adapter.js"; import { createClawdbotCodingTools } from "./pi-tools.js"; import { resolveSandboxContext } from "./sandbox.js"; +import { sanitizeToolUseResultPairing } from "./session-transcript-repair.js"; import { applySkillEnvOverrides, applySkillEnvOverridesFromSnapshot, diff --git a/src/agents/session-transcript-repair.test.ts b/src/agents/session-transcript-repair.test.ts new file mode 100644 index 000000000..acb92fa31 --- /dev/null +++ b/src/agents/session-transcript-repair.test.ts @@ -0,0 +1,95 @@ +import type { AgentMessage } from "@mariozechner/pi-agent-core"; +import { describe, expect, it } from "vitest"; +import { sanitizeToolUseResultPairing } from "./session-transcript-repair.js"; + +describe("sanitizeToolUseResultPairing", () => { + it("moves tool results directly after tool calls and inserts missing results", () => { + const input = [ + { + role: "assistant", + content: [ + { type: "toolCall", id: "call_1", name: "read", arguments: {} }, + { type: "toolCall", id: "call_2", name: "bash", arguments: {} }, + ], + }, + { role: "user", content: "user message that should come after tool use" }, + { + role: "toolResult", + toolCallId: "call_2", + toolName: "bash", + content: [{ type: "text", text: "ok" }], + isError: false, + }, + ] satisfies AgentMessage[]; + + const out = sanitizeToolUseResultPairing(input); + expect(out[0]?.role).toBe("assistant"); + expect(out[1]?.role).toBe("toolResult"); + expect((out[1] as { toolCallId?: string }).toolCallId).toBe("call_1"); + expect(out[2]?.role).toBe("toolResult"); + expect((out[2] as { toolCallId?: string }).toolCallId).toBe("call_2"); + expect(out[3]?.role).toBe("user"); + }); + + it("drops duplicate tool results for the same id within a span", () => { + const input = [ + { + role: "assistant", + content: [ + { type: "toolCall", id: "call_1", name: "read", arguments: {} }, + ], + }, + { + role: "toolResult", + toolCallId: "call_1", + toolName: "read", + content: [{ type: "text", text: "first" }], + isError: false, + }, + { + role: "toolResult", + toolCallId: "call_1", + toolName: "read", + content: [{ type: "text", text: "second" }], + isError: false, + }, + { role: "user", content: "ok" }, + ] satisfies AgentMessage[]; + + const out = sanitizeToolUseResultPairing(input); + expect(out.filter((m) => m.role === "toolResult")).toHaveLength(1); + }); + + it("drops duplicate tool results for the same id across the transcript", () => { + const input = [ + { + role: "assistant", + content: [ + { type: "toolCall", id: "call_1", name: "read", arguments: {} }, + ], + }, + { + role: "toolResult", + toolCallId: "call_1", + toolName: "read", + content: [{ type: "text", text: "first" }], + isError: false, + }, + { role: "assistant", content: [{ type: "text", text: "ok" }] }, + { + role: "toolResult", + toolCallId: "call_1", + toolName: "read", + content: [{ type: "text", text: "second (duplicate)" }], + isError: false, + }, + ] satisfies AgentMessage[]; + + const out = sanitizeToolUseResultPairing(input); + const results = out.filter((m) => m.role === "toolResult") as Array<{ + toolCallId?: string; + }>; + expect(results).toHaveLength(1); + expect(results[0]?.toolCallId).toBe("call_1"); + }); +}); diff --git a/src/agents/session-transcript-repair.ts b/src/agents/session-transcript-repair.ts new file mode 100644 index 000000000..1c03c1add --- /dev/null +++ b/src/agents/session-transcript-repair.ts @@ -0,0 +1,173 @@ +import type { AgentMessage } from "@mariozechner/pi-agent-core"; + +type ToolCallLike = { + id: string; + name?: string; +}; + +function extractToolCallsFromAssistant( + msg: Extract, +): ToolCallLike[] { + const content = msg.content; + if (!Array.isArray(content)) return []; + + const toolCalls: ToolCallLike[] = []; + for (const block of content) { + if (!block || typeof block !== "object") continue; + const rec = block as { type?: unknown; id?: unknown; name?: unknown }; + if (typeof rec.id !== "string" || !rec.id) continue; + + if ( + rec.type === "toolCall" || + rec.type === "toolUse" || + rec.type === "functionCall" + ) { + toolCalls.push({ + id: rec.id, + name: typeof rec.name === "string" ? rec.name : undefined, + }); + } + } + return toolCalls; +} + +function extractToolResultId( + msg: Extract, +): string | null { + const toolCallId = (msg as { toolCallId?: unknown }).toolCallId; + if (typeof toolCallId === "string" && toolCallId) return toolCallId; + const toolUseId = (msg as { toolUseId?: unknown }).toolUseId; + if (typeof toolUseId === "string" && toolUseId) return toolUseId; + return null; +} + +function makeMissingToolResult(params: { + toolCallId: string; + toolName?: string; +}): Extract { + return { + role: "toolResult", + toolCallId: params.toolCallId, + toolName: params.toolName ?? "unknown", + content: [ + { + type: "text", + text: "[clawdbot] missing tool result in session history; inserted synthetic error result for transcript repair.", + }, + ], + isError: true, + timestamp: Date.now(), + } as Extract; +} + +export function sanitizeToolUseResultPairing( + messages: AgentMessage[], +): AgentMessage[] { + // Anthropic (and Cloud Code Assist) reject transcripts where assistant tool calls are not + // immediately followed by matching tool results. Session files can end up with results + // displaced (e.g. after user turns) or duplicated. Repair by: + // - moving matching toolResult messages directly after their assistant toolCall turn + // - inserting synthetic error toolResults for missing ids + // - dropping duplicate toolResults for the same id (anywhere in the transcript) + const out: AgentMessage[] = []; + const seenToolResultIds = new Set(); + + const pushToolResult = ( + msg: Extract, + ) => { + const id = extractToolResultId(msg); + if (id && seenToolResultIds.has(id)) return; + if (id) seenToolResultIds.add(id); + out.push(msg); + }; + + for (let i = 0; i < messages.length; i += 1) { + const msg = messages[i] as AgentMessage; + if (!msg || typeof msg !== "object") { + out.push(msg); + continue; + } + + const role = (msg as { role?: unknown }).role; + if (role !== "assistant") { + if (role === "toolResult") { + pushToolResult(msg as Extract); + } else { + out.push(msg); + } + continue; + } + + const assistant = msg as Extract; + const toolCalls = extractToolCallsFromAssistant(assistant); + if (toolCalls.length === 0) { + out.push(msg); + continue; + } + + const toolCallIds = new Set(toolCalls.map((t) => t.id)); + + const spanResultsById = new Map< + string, + Extract + >(); + const remainder: AgentMessage[] = []; + + let j = i + 1; + for (; j < messages.length; j += 1) { + const next = messages[j] as AgentMessage; + if (!next || typeof next !== "object") { + remainder.push(next); + continue; + } + + const nextRole = (next as { role?: unknown }).role; + if (nextRole === "assistant") break; + + if (nextRole === "toolResult") { + const toolResult = next as Extract< + AgentMessage, + { role: "toolResult" } + >; + const id = extractToolResultId(toolResult); + if (id && toolCallIds.has(id)) { + if (seenToolResultIds.has(id)) { + continue; + } + if (!spanResultsById.has(id)) { + spanResultsById.set(id, toolResult); + } + continue; + } + } + + remainder.push(next); + } + + out.push(msg); + + for (const call of toolCalls) { + const existing = spanResultsById.get(call.id); + pushToolResult( + existing ?? + makeMissingToolResult({ toolCallId: call.id, toolName: call.name }), + ); + } + + for (const rem of remainder) { + if (!rem || typeof rem !== "object") { + out.push(rem); + continue; + } + const remRole = (rem as { role?: unknown }).role; + if (remRole === "toolResult") { + pushToolResult(rem as Extract); + continue; + } + out.push(rem); + } + i = j - 1; + } + + return out; +}