fix(session): refine sanitization logic and add tests
Addresses review feedback: parse valid JSON strings, add type guards, add logging, handle arguments alias.
This commit is contained in:
parent
5f4715acfc
commit
ccf00e10cb
@ -110,3 +110,80 @@ describe("sanitizeToolUseResultPairing", () => {
|
|||||||
expect(out.map((m) => m.role)).toEqual(["user", "assistant"]);
|
expect(out.map((m) => m.role)).toEqual(["user", "assistant"]);
|
||||||
});
|
});
|
||||||
});
|
});
|
||||||
|
|
||||||
|
import { sanitizeToolUseArgs } from "./session-transcript-repair.js";
|
||||||
|
|
||||||
|
describe("sanitizeToolUseArgs", () => {
|
||||||
|
it("preserves valid objects in input/arguments", () => {
|
||||||
|
const input = [
|
||||||
|
{
|
||||||
|
role: "assistant",
|
||||||
|
content: [
|
||||||
|
{ type: "toolUse", id: "1", name: "tool", input: { key: "value" } },
|
||||||
|
],
|
||||||
|
},
|
||||||
|
] as any;
|
||||||
|
const out = sanitizeToolUseArgs(input);
|
||||||
|
expect((out[0].content[0] as any).input).toEqual({ key: "value" });
|
||||||
|
expect(out).toBe(input); // No change, referentially equal
|
||||||
|
});
|
||||||
|
|
||||||
|
it("parses valid JSON strings in input", () => {
|
||||||
|
const input = [
|
||||||
|
{
|
||||||
|
role: "assistant",
|
||||||
|
content: [
|
||||||
|
{ type: "toolUse", id: "1", name: "tool", input: '{"key": "value"}' },
|
||||||
|
],
|
||||||
|
},
|
||||||
|
] as any;
|
||||||
|
const out = sanitizeToolUseArgs(input);
|
||||||
|
expect((out[0].content[0] as any).input).toEqual({ key: "value" });
|
||||||
|
expect(out).not.toBe(input); // Changed
|
||||||
|
});
|
||||||
|
|
||||||
|
it("sanitizes invalid JSON strings in input", () => {
|
||||||
|
const input = [
|
||||||
|
{
|
||||||
|
role: "assistant",
|
||||||
|
content: [
|
||||||
|
{ type: "toolUse", id: "1", name: "tool", input: '{ bad json }' },
|
||||||
|
],
|
||||||
|
},
|
||||||
|
] as any;
|
||||||
|
const out = sanitizeToolUseArgs(input);
|
||||||
|
const block = out[0].content[0] as any;
|
||||||
|
expect(block.input).toEqual({});
|
||||||
|
expect(block._sanitized).toBe(true);
|
||||||
|
expect(block._originalInput).toBe('{ bad json }');
|
||||||
|
});
|
||||||
|
|
||||||
|
it("handles 'arguments' alias", () => {
|
||||||
|
const input = [
|
||||||
|
{
|
||||||
|
role: "assistant",
|
||||||
|
content: [
|
||||||
|
{ type: "toolCall", id: "1", name: "tool", arguments: '{"key": "val"}' },
|
||||||
|
],
|
||||||
|
},
|
||||||
|
] as any;
|
||||||
|
const out = sanitizeToolUseArgs(input);
|
||||||
|
const block = out[0].content[0] as any;
|
||||||
|
expect(block.arguments).toEqual({ key: "val" });
|
||||||
|
});
|
||||||
|
|
||||||
|
it("sanitizes invalid JSON in 'arguments' alias", () => {
|
||||||
|
const input = [
|
||||||
|
{
|
||||||
|
role: "assistant",
|
||||||
|
content: [
|
||||||
|
{ type: "toolCall", id: "1", name: "tool", arguments: 'bad' },
|
||||||
|
],
|
||||||
|
},
|
||||||
|
] as any;
|
||||||
|
const out = sanitizeToolUseArgs(input);
|
||||||
|
const block = out[0].content[0] as any;
|
||||||
|
expect(block.arguments).toEqual({});
|
||||||
|
expect(block._sanitized).toBe(true);
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|||||||
@ -56,6 +56,78 @@ function makeMissingToolResult(params: {
|
|||||||
|
|
||||||
export { makeMissingToolResult };
|
export { makeMissingToolResult };
|
||||||
|
|
||||||
|
export function sanitizeToolUseArgs(messages: AgentMessage[]): AgentMessage[] {
|
||||||
|
// 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;
|
||||||
|
|
||||||
|
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.
|
||||||
|
console.warn(
|
||||||
|
`[SessionRepair] Sanitized malformed JSON in tool use '${toolBlock.name || "unknown"}'. Original: ${toolBlock[inputField]}`
|
||||||
|
);
|
||||||
|
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 changed ? out : messages;
|
||||||
|
}
|
||||||
export function sanitizeToolUseResultPairing(messages: AgentMessage[]): AgentMessage[] {
|
export function sanitizeToolUseResultPairing(messages: AgentMessage[]): AgentMessage[] {
|
||||||
return repairToolUseResultPairing(messages).messages;
|
return repairToolUseResultPairing(messages).messages;
|
||||||
}
|
}
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user