From ccf00e10cb886138408320585c3cef3788191869 Mon Sep 17 00:00:00 2001
From: Nathan Hangen <hangen.2@wright.edu>
Date: Wed, 28 Jan 2026 19:18:14 -0500
Subject: [PATCH] fix(session): refine sanitization logic and add tests

Addresses review feedback: parse valid JSON strings, add type guards, add logging, handle arguments alias.
---
 src/agents/session-transcript-repair.test.ts | 77 ++++++++++++++++++++
 src/agents/session-transcript-repair.ts      | 72 ++++++++++++++++++
 2 files changed, 149 insertions(+)

diff --git a/src/agents/session-transcript-repair.test.ts b/src/agents/session-transcript-repair.test.ts
index ccc63ec7f..96d42ae1a 100644
--- a/src/agents/session-transcript-repair.test.ts
+++ b/src/agents/session-transcript-repair.test.ts
@@ -110,3 +110,80 @@ describe("sanitizeToolUseResultPairing", () => {
     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);
+  });
+});
diff --git a/src/agents/session-transcript-repair.ts b/src/agents/session-transcript-repair.ts
index d680beb4d..e5e4786ff 100644
--- a/src/agents/session-transcript-repair.ts
+++ b/src/agents/session-transcript-repair.ts
@@ -56,6 +56,78 @@ function makeMissingToolResult(params: {
 
 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[] {
   return repairToolUseResultPairing(messages).messages;
 }