Merge f041a95006 into 09be5d45d5
This commit is contained in:
commit
a3c551c780
@ -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);
|
||||
});
|
||||
});
|
||||
|
||||
@ -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 = {
|
||||
|
||||
Loading…
Reference in New Issue
Block a user