fix: sanitize incomplete tool calls with partialJson
When a request is terminated mid-stream, tool call blocks may have partialJson but incomplete/missing arguments. These incomplete tool calls cause API rejections because their matching tool_result blocks become orphaned. This fix: - Adds sanitizePartialToolCalls() to remove incomplete tool calls - Updates extractToolCallsFromAssistant() to skip incomplete calls - Integrates with existing sanitizeToolUseResultPairing() flow Fixes: Tool call terminated mid-execution causes 400 error with 'unexpected tool_use_id found in tool_result blocks'
This commit is contained in:
Zedit42
parent
ded366d9ab
commit
1c82b096b8
@ -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<AgentMessage, { role: "assistant" }>;
|
||||
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<AgentMessage, { role: "assistant" }>;
|
||||
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<AgentMessage, { role: "assistant" }>;
|
||||
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
|
||||
});
|
||||
});
|
||||
|
||||
@ -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<AgentMessage, { role: "assistant" }>;
|
||||
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<AgentMessage, { role: "assistant" }>,
|
||||
): 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 = {
|
||||
|
||||
Loading…
Reference in New Issue
Block a user