diff --git a/src/agents/session-transcript-repair.test.ts b/src/agents/session-transcript-repair.test.ts index ccc63ec7f..0b79b2c3d 100644 --- a/src/agents/session-transcript-repair.test.ts +++ b/src/agents/session-transcript-repair.test.ts @@ -1,26 +1,33 @@ import type { AgentMessage } from "@mariozechner/pi-agent-core"; import { describe, expect, it } from "vitest"; -import { sanitizeToolUseResultPairing } from "./session-transcript-repair.js"; +import { sanitizeToolUseArgs, sanitizeToolUseResultPairing } from "./session-transcript-repair.js"; + +const now = Date.now(); describe("sanitizeToolUseResultPairing", () => { it("moves tool results directly after tool calls and inserts missing results", () => { - const input = [ + const input: AgentMessage[] = [ { role: "assistant", content: [ { type: "toolCall", id: "call_1", name: "read", arguments: {} }, { type: "toolCall", id: "call_2", name: "exec", arguments: {} }, ], + timestamp: now, + api: "openai", + provider: "openai", + model: "gpt-4", }, - { role: "user", content: "user message that should come after tool use" }, + { role: "user", content: "user message that should come after tool use", timestamp: now }, { role: "toolResult", toolCallId: "call_2", toolName: "exec", content: [{ type: "text", text: "ok" }], isError: false, + timestamp: now, }, - ] satisfies AgentMessage[]; + ]; const out = sanitizeToolUseResultPairing(input); expect(out[0]?.role).toBe("assistant"); @@ -32,10 +39,14 @@ describe("sanitizeToolUseResultPairing", () => { }); it("drops duplicate tool results for the same id within a span", () => { - const input = [ + const input: AgentMessage[] = [ { role: "assistant", content: [{ type: "toolCall", id: "call_1", name: "read", arguments: {} }], + timestamp: now, + api: "openai", + provider: "openai", + model: "gpt-4", }, { role: "toolResult", @@ -43,6 +54,7 @@ describe("sanitizeToolUseResultPairing", () => { toolName: "read", content: [{ type: "text", text: "first" }], isError: false, + timestamp: now, }, { role: "toolResult", @@ -50,19 +62,24 @@ describe("sanitizeToolUseResultPairing", () => { toolName: "read", content: [{ type: "text", text: "second" }], isError: false, + timestamp: now, }, - { role: "user", content: "ok" }, - ] satisfies AgentMessage[]; + { role: "user", content: "ok", timestamp: now }, + ]; 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 = [ + const input: AgentMessage[] = [ { role: "assistant", content: [{ type: "toolCall", id: "call_1", name: "read", arguments: {} }], + timestamp: now, + api: "openai", + provider: "openai", + model: "gpt-4", }, { role: "toolResult", @@ -70,16 +87,25 @@ describe("sanitizeToolUseResultPairing", () => { toolName: "read", content: [{ type: "text", text: "first" }], isError: false, + timestamp: now, + }, + { + role: "assistant", + content: [{ type: "text", text: "ok" }], + timestamp: now, + api: "openai", + provider: "openai", + model: "gpt-4", }, - { role: "assistant", content: [{ type: "text", text: "ok" }] }, { role: "toolResult", toolCallId: "call_1", toolName: "read", content: [{ type: "text", text: "second (duplicate)" }], isError: false, + timestamp: now, }, - ] satisfies AgentMessage[]; + ]; const out = sanitizeToolUseResultPairing(input); const results = out.filter((m) => m.role === "toolResult") as Array<{ @@ -90,23 +116,131 @@ describe("sanitizeToolUseResultPairing", () => { }); it("drops orphan tool results that do not match any tool call", () => { - const input = [ - { role: "user", content: "hello" }, + const input: AgentMessage[] = [ + { role: "user", content: "hello", timestamp: now }, { role: "toolResult", toolCallId: "call_orphan", toolName: "read", content: [{ type: "text", text: "orphan" }], isError: false, + timestamp: now, }, { role: "assistant", content: [{ type: "text", text: "ok" }], + timestamp: now, + api: "openai", + provider: "openai", + model: "gpt-4", }, - ] satisfies AgentMessage[]; + ]; const out = sanitizeToolUseResultPairing(input); expect(out.some((m) => m.role === "toolResult")).toBe(false); expect(out.map((m) => m.role)).toEqual(["user", "assistant"]); }); }); + +describe("sanitizeToolUseArgs", () => { + it("preserves valid JSON strings in input fields", () => { + const input: AgentMessage[] = [ + { + role: "assistant", + content: [ + { type: "toolCall", id: "call_1", name: "read", input: '{"path":"foo.txt"}' } as any, + ], + timestamp: now, + api: "openai", + provider: "openai", + model: "gpt-4", + }, + ]; + + const result = sanitizeToolUseArgs(input); + expect(result.changed).toBe(true); + const tool = (result.messages[0] as any).content[0]; + expect(tool.input).toEqual({ path: "foo.txt" }); + expect(result.sanitizedCount).toBe(0); + }); + + it("replaces invalid JSON strings with {} and sets metadata", () => { + const input: AgentMessage[] = [ + { + role: "assistant", + content: [{ type: "toolCall", id: "call_1", name: "read", input: "{ invalid json" } as any], + timestamp: now, + api: "openai", + provider: "openai", + model: "gpt-4", + }, + ]; + + const result = sanitizeToolUseArgs(input); + expect(result.changed).toBe(true); + expect(result.sanitizedCount).toBe(1); + const tool = (result.messages[0] as any).content[0]; + expect(tool.input).toEqual({}); + expect(tool._sanitized).toBe(true); + expect(tool._originalInput).toBe("{ invalid json"); + }); + + it("preserves already-parsed object values in input fields", () => { + const input: AgentMessage[] = [ + { + role: "assistant", + content: [ + { type: "toolCall", id: "call_1", name: "read", input: { path: "bar.txt" } } as any, + ], + timestamp: now, + api: "openai", + provider: "openai", + model: "gpt-4", + }, + ]; + + const result = sanitizeToolUseArgs(input); + expect(result.changed).toBe(false); + expect(result.messages).toBe(input); + const tool = (result.messages[0] as any).content[0]; + expect(tool.input).toEqual({ path: "bar.txt" }); + }); + + it("handles the 'arguments' alias used by some providers", () => { + const input: AgentMessage[] = [ + { + role: "assistant", + content: [ + { type: "toolCall", id: "call_1", name: "read", arguments: '{"path":"baz.txt"}' } as any, + ], + timestamp: now, + api: "openai", + provider: "openai", + model: "gpt-4", + }, + ]; + + const result = sanitizeToolUseArgs(input); + expect(result.changed).toBe(true); + const tool = (result.messages[0] as any).content[0]; + expect(tool.arguments).toEqual({ path: "baz.txt" }); + }); + + it("leaves messages without tool blocks unchanged", () => { + const input: AgentMessage[] = [ + { role: "user", content: "hello", timestamp: now }, + { + role: "assistant", + content: [{ type: "text", text: "hi" }], + timestamp: now, + api: "openai", + provider: "openai", + model: "gpt-4", + }, + ]; + + const result = sanitizeToolUseArgs(input); + expect(result.changed).toBe(false); + expect(result.messages).toBe(input); + }); +}); diff --git a/src/agents/session-transcript-repair.ts b/src/agents/session-transcript-repair.ts index 2728bb89c..477841310 100644 --- a/src/agents/session-transcript-repair.ts +++ b/src/agents/session-transcript-repair.ts @@ -56,8 +56,99 @@ function makeMissingToolResult(params: { export { makeMissingToolResult }; +export type ToolUseSanitizationReport = { + messages: AgentMessage[]; + sanitizedCount: number; + changed: boolean; +}; + +export function sanitizeToolUseArgs(messages: AgentMessage[]): ToolUseSanitizationReport { + // Creates new message objects only when sanitization is needed; otherwise + // returns the original messages to avoid unnecessary copying, while guarding + // against corrupt JSON in tool arguments that could break the session. + const out: AgentMessage[] = []; + let changed = false; + let sanitizedCount = 0; + + for (const msg of messages) { + if (msg.role !== "assistant" || !Array.isArray(msg.content)) { + out.push(msg); + continue; + } + + const content = msg.content; + const nextContent: any[] = []; + let msgChanged = false; + + for (const block of content) { + const anyBlock = block as any; + if ( + anyBlock && + typeof anyBlock === "object" && + (anyBlock.type === "toolUse" || + anyBlock.type === "toolCall" || + anyBlock.type === "functionCall") + ) { + const toolBlock = block as any; + // Handle both 'input' and 'arguments' fields (some providers use arguments) + const inputField = + "input" in toolBlock ? "input" : "arguments" in toolBlock ? "arguments" : null; + + if (inputField && typeof toolBlock[inputField] === "string") { + try { + // Consistency: Always parse valid JSON strings into objects + const parsed = JSON.parse(toolBlock[inputField]); + nextContent.push({ + ...toolBlock, + [inputField]: parsed, + }); + msgChanged = true; + } catch { + // Invalid JSON found in tool args. + // Replace with empty object to prevent downstream crashes. + sanitizedCount += 1; + const original = String(toolBlock[inputField]); + const sample = original.length > 100 ? `${original.slice(0, 100)}...` : original; + console.warn( + `[SessionRepair] Sanitized malformed JSON in tool use '${toolBlock.name || "unknown"}'. Original: ${sample}`, + ); + nextContent.push({ + ...toolBlock, + [inputField]: {}, + _sanitized: true, + _originalInput: toolBlock[inputField], + }); + msgChanged = true; + } + } else { + nextContent.push(block); + } + } else { + nextContent.push(block); + } + } + + if (msgChanged) { + out.push({ + ...msg, + content: nextContent, + } as AgentMessage); + changed = true; + } else { + out.push(msg); + } + } + + return { + messages: changed ? out : messages, + sanitizedCount, + changed, + }; +} + export function sanitizeToolUseResultPairing(messages: AgentMessage[]): AgentMessage[] { - return repairToolUseResultPairing(messages).messages; + const sanitized = sanitizeToolUseArgs(messages); + return repairToolUseResultPairing(sanitized.messages).messages; } export type ToolUseRepairReport = {