This commit is contained in:
Nathan Hangen 2026-01-30 11:21:25 -05:00 committed by GitHub
commit a3c551c780
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 239 additions and 14 deletions

View File

@ -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);
});
});

View File

@ -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 = {