Merge f041a95006 into 09be5d45d5

src/agents/session-transcript-repair.test.ts

@@ -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" },
+      { role: "user", content: "ok", timestamp: now },
-    ] satisfies AgentMessage[];
+    ];
 
     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 = [
+    const input: AgentMessage[] = [
-      { role: "user", content: "hello" },
+      { 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);
+  });
+});

src/agents/session-transcript-repair.ts

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