romtuck/openclaw
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(session): improve sanitization observability and add type-safe tests

Browse Source
This commit is contained in:
Nathan Hangen 2026-01-29 11:12:11 -05:00
parent dc2c40468d
commit f041a95006
2 changed files with 140 additions and 57 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

173
src/agents/session-transcript-repair.test.ts
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,20 +116,25 @@ 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);
@ -111,69 +142,105 @@ describe("sanitizeToolUseResultPairing", () => {
});
});
import { sanitizeToolUseArgs } from "./session-transcript-repair.js";
describe("sanitizeToolUseArgs", () => {
it("preserves valid objects in input/arguments", () => {
const input = [
it("preserves valid JSON strings in input fields", () => {
const input: AgentMessage[] = [
{
role: "assistant",
content: [{ type: "toolUse", id: "1", name: "tool", input: { key: "value" } }],
content: [
{ type: "toolCall", id: "call_1", name: "read", input: '{"path":"foo.txt"}' } as any,
],
timestamp: now,
api: "openai",
provider: "openai",
model: "gpt-4",
},
] 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
];
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("parses valid JSON strings in input", () => {
const input = [
it("replaces invalid JSON strings with {} and sets metadata", () => {
const input: AgentMessage[] = [
{
role: "assistant",
content: [{ type: "toolUse", id: "1", name: "tool", input: '{"key": "value"}' }],
content: [{ type: "toolCall", id: "call_1", name: "read", input: "{ invalid json" } as any],
timestamp: now,
api: "openai",
provider: "openai",
model: "gpt-4",
},
] as any;
const out = sanitizeToolUseArgs(input);
expect((out[0].content[0] as any).input).toEqual({ key: "value" });
expect(out).not.toBe(input); // Changed
];
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("sanitizes invalid JSON strings in input", () => {
const input = [
it("preserves already-parsed object values in input fields", () => {
const input: AgentMessage[] = [
{
role: "assistant",
content: [{ type: "toolUse", id: "1", name: "tool", input: "{ bad json }" }],
content: [
{ type: "toolCall", id: "call_1", name: "read", input: { path: "bar.txt" } } as any,
],
timestamp: now,
api: "openai",
provider: "openai",
model: "gpt-4",
},
] 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 }");
];
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 'arguments' alias", () => {
const input = [
it("handles the 'arguments' alias used by some providers", () => {
const input: AgentMessage[] = [
{
role: "assistant",
content: [{ type: "toolCall", id: "1", name: "tool", arguments: '{"key": "val"}' }],
content: [
{ type: "toolCall", id: "call_1", name: "read", arguments: '{"path":"baz.txt"}' } as any,
],
timestamp: now,
api: "openai",
provider: "openai",
model: "gpt-4",
},
] as any;
const out = sanitizeToolUseArgs(input);
const block = out[0].content[0] as any;
expect(block.arguments).toEqual({ key: "val" });
];
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("sanitizes invalid JSON in 'arguments' alias", () => {
const input = [
it("leaves messages without tool blocks unchanged", () => {
const input: AgentMessage[] = [
{ role: "user", content: "hello", timestamp: now },
{
role: "assistant",
content: [{ type: "toolCall", id: "1", name: "tool", arguments: "bad" }],
content: [{ type: "text", text: "hi" }],
timestamp: now,
api: "openai",
provider: "openai",
model: "gpt-4",
},
] as any;
const out = sanitizeToolUseArgs(input);
const block = out[0].content[0] as any;
expect(block.arguments).toEqual({});
expect(block._sanitized).toBe(true);
];
const result = sanitizeToolUseArgs(input);
expect(result.changed).toBe(false);
expect(result.messages).toBe(input);
});
});

24
src/agents/session-transcript-repair.ts
View File

@ -56,12 +56,19 @@ function makeMissingToolResult(params: {
export { makeMissingToolResult };
export function sanitizeToolUseArgs(messages: AgentMessage[]): AgentMessage[] {
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)) {
@ -99,8 +106,11 @@ export function sanitizeToolUseArgs(messages: AgentMessage[]): AgentMessage[] {
} 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: ${toolBlock[inputField]}`,
`[SessionRepair] Sanitized malformed JSON in tool use '${toolBlock.name || "unknown"}'. Original: ${sample}`,
);
nextContent.push({
...toolBlock,
@ -129,10 +139,16 @@ export function sanitizeToolUseArgs(messages: AgentMessage[]): AgentMessage[] {
}
}
return changed ? out : messages;
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 = {