diff --git a/src/agents/session-transcript-repair.test.ts b/src/agents/session-transcript-repair.test.ts index ccc63ec7f..681d0d360 100644 --- a/src/agents/session-transcript-repair.test.ts +++ b/src/agents/session-transcript-repair.test.ts @@ -1,6 +1,9 @@ import type { AgentMessage } from "@mariozechner/pi-agent-core"; import { describe, expect, it } from "vitest"; -import { sanitizeToolUseResultPairing } from "./session-transcript-repair.js"; +import { + sanitizeToolUseResultPairing, + sanitizePartialToolCalls, +} from "./session-transcript-repair.js"; describe("sanitizeToolUseResultPairing", () => { it("moves tool results directly after tool calls and inserts missing results", () => { @@ -109,4 +112,133 @@ describe("sanitizeToolUseResultPairing", () => { expect(out.some((m) => m.role === "toolResult")).toBe(false); expect(out.map((m) => m.role)).toEqual(["user", "assistant"]); }); + + it("removes incomplete tool calls with partialJson and drops their orphaned results", () => { + // This simulates a terminated request where tool call was incomplete + // Note: arguments is undefined/missing when truly incomplete + const input = [ + { + role: "assistant", + content: [ + { type: "text", text: "Let me write that file:" }, + { + type: "toolCall", + id: "call_incomplete", + name: "write", + // arguments is missing - only partialJson exists + partialJson: '{"path": "/tmp/test.md", "content": "# Hello', + }, + ], + }, + { + role: "toolResult", + toolCallId: "call_incomplete", + toolName: "write", + content: [{ type: "text", text: "[clawdbot] synthetic error result" }], + isError: true, + }, + ] satisfies AgentMessage[]; + + const out = sanitizeToolUseResultPairing(input); + // The incomplete tool call should be removed from content + const assistant = out[0] as Extract; + expect(assistant.content).toHaveLength(1); + expect(assistant.content[0]).toEqual({ type: "text", text: "Let me write that file:" }); + // The orphaned tool result should also be dropped + expect(out.filter((m) => m.role === "toolResult")).toHaveLength(0); + }); + + it("keeps complete tool calls even if partialJson was captured", () => { + // If arguments is complete, the tool call should be kept + const input = [ + { + role: "assistant", + content: [ + { + type: "toolCall", + id: "call_complete", + name: "read", + arguments: { path: "/tmp/test.md" }, + partialJson: '{"path": "/tmp/test.md"}', // partialJson exists but args complete + }, + ], + }, + { + role: "toolResult", + toolCallId: "call_complete", + toolName: "read", + content: [{ type: "text", text: "file contents" }], + isError: false, + }, + ] satisfies AgentMessage[]; + + const out = sanitizeToolUseResultPairing(input); + expect(out).toHaveLength(2); + expect(out[0]?.role).toBe("assistant"); + expect(out[1]?.role).toBe("toolResult"); + }); +}); + +describe("sanitizePartialToolCalls", () => { + it("removes tool calls with partialJson but no arguments", () => { + const input = [ + { + role: "assistant", + content: [ + { type: "text", text: "Working on it" }, + { + type: "toolCall", + id: "call_1", + name: "write", + partialJson: '{"path": "/tmp/file.md", "content": "partial...', + }, + ], + }, + ] satisfies AgentMessage[]; + + const out = sanitizePartialToolCalls(input); + const assistant = out[0] as Extract; + expect(assistant.content).toHaveLength(1); + expect((assistant.content[0] as { type: string }).type).toBe("text"); + }); + + it("removes tool calls with partialJson and empty arguments", () => { + const input = [ + { + role: "assistant", + content: [ + { + type: "toolCall", + id: "call_1", + name: "write", + arguments: {}, + partialJson: '{"path": "/tmp', + }, + ], + }, + ] satisfies AgentMessage[]; + + const out = sanitizePartialToolCalls(input); + const assistant = out[0] as Extract; + expect(assistant.content).toHaveLength(0); + }); + + it("preserves complete tool calls without partialJson", () => { + const input = [ + { + role: "assistant", + content: [ + { + type: "toolCall", + id: "call_1", + name: "read", + arguments: { path: "/tmp/file.md" }, + }, + ], + }, + ] satisfies AgentMessage[]; + + const out = sanitizePartialToolCalls(input); + expect(out).toBe(input); // Should return same reference if unchanged + }); }); diff --git a/src/agents/session-transcript-repair.ts b/src/agents/session-transcript-repair.ts index d680beb4d..2547855f7 100644 --- a/src/agents/session-transcript-repair.ts +++ b/src/agents/session-transcript-repair.ts @@ -5,6 +5,82 @@ type ToolCallLike = { name?: string; }; +type ToolCallBlock = { + type?: unknown; + id?: unknown; + name?: unknown; + partialJson?: unknown; + arguments?: unknown; +}; + +/** + * Checks if a tool call block is incomplete (has partialJson but no complete arguments). + * Incomplete tool calls happen when a request is terminated mid-stream. + */ +function isIncompleteToolCall(block: ToolCallBlock): boolean { + if (!block || typeof block !== "object") return false; + // If partialJson exists and arguments is missing or incomplete, it's partial + if (typeof block.partialJson === "string" && block.partialJson) { + // Check if arguments is missing or empty + const args = block.arguments; + if (args === undefined || args === null) return true; + if (typeof args === "object" && Object.keys(args as object).length === 0) return true; + } + return false; +} + +/** + * Removes incomplete (partial) tool calls from assistant messages. + * These occur when a request is terminated mid-stream and can cause API rejections. + */ +export function sanitizePartialToolCalls(messages: AgentMessage[]): AgentMessage[] { + let changed = false; + const out: AgentMessage[] = []; + + for (const msg of messages) { + if (!msg || typeof msg !== "object") { + out.push(msg); + continue; + } + + const role = (msg as { role?: unknown }).role; + if (role !== "assistant") { + out.push(msg); + continue; + } + + const assistant = msg as Extract; + const content = assistant.content; + if (!Array.isArray(content)) { + out.push(msg); + continue; + } + + // Filter out incomplete tool calls + const filteredContent = content.filter((block) => { + if (!block || typeof block !== "object") return true; + const rec = block as ToolCallBlock; + if (rec.type !== "toolCall" && rec.type !== "toolUse" && rec.type !== "functionCall") { + return true; + } + if (isIncompleteToolCall(rec)) { + changed = true; + return false; + } + return true; + }); + + if (filteredContent.length !== content.length) { + // Content was modified, create new message + out.push({ ...assistant, content: filteredContent } as AgentMessage); + } else { + out.push(msg); + } + } + + return changed ? out : messages; +} + function extractToolCallsFromAssistant( msg: Extract, ): ToolCallLike[] { @@ -14,10 +90,13 @@ function extractToolCallsFromAssistant( const toolCalls: ToolCallLike[] = []; for (const block of content) { if (!block || typeof block !== "object") continue; - const rec = block as { type?: unknown; id?: unknown; name?: unknown }; + const rec = block as ToolCallBlock; if (typeof rec.id !== "string" || !rec.id) continue; if (rec.type === "toolCall" || rec.type === "toolUse" || rec.type === "functionCall") { + // Skip incomplete tool calls - they don't count as valid tool calls + if (isIncompleteToolCall(rec)) continue; + toolCalls.push({ id: rec.id, name: typeof rec.name === "string" ? rec.name : undefined, @@ -57,7 +136,10 @@ function makeMissingToolResult(params: { export { makeMissingToolResult }; export function sanitizeToolUseResultPairing(messages: AgentMessage[]): AgentMessage[] { - return repairToolUseResultPairing(messages).messages; + // First, sanitize partial tool calls that were terminated mid-stream + const sanitizedPartials = sanitizePartialToolCalls(messages); + // Then repair tool use/result pairing + return repairToolUseResultPairing(sanitizedPartials).messages; } export type ToolUseRepairReport = {