From 033ff87861ea7e63c02a4ee5ff98f7c155dd1c43 Mon Sep 17 00:00:00 2001
From: Jiang Bohan <bhjiang@outlook.com>
Date: Thu, 26 Feb 2026 17:33:58 +0800
Subject: [PATCH] fix(agent): strip toolCalls from aborted/error assistant
 messages in transcript repair
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

When a streaming request is aborted mid-toolCall, the session persists an
assistant message with stopReason "aborted" containing partial toolCall blocks.
Our sanitizeToolUseResultPairing then inserts synthetic toolResults for these
toolCalls. However, pi-ai's transformMessages drops the entire aborted assistant
message downstream, leaving orphaned toolResults that reference non-existent
tool_use_ids — causing persistent 400 errors that block all subsequent
conversations in the session.

Fix: in repairToolCallInputs, strip toolCall blocks from assistant messages
with stopReason "aborted" or "error" before the result-pairing sanitizer runs.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 .../session/session-transcript-repair.test.ts | 80 +++++++++++++++++++
 .../session/session-transcript-repair.ts      | 19 +++++
 2 files changed, 99 insertions(+)

diff --git a/packages/core/src/agent/session/session-transcript-repair.test.ts b/packages/core/src/agent/session/session-transcript-repair.test.ts
index a4d9e79c..beb1ba15 100644
--- a/packages/core/src/agent/session/session-transcript-repair.test.ts
+++ b/packages/core/src/agent/session/session-transcript-repair.test.ts
@@ -257,4 +257,84 @@ describe("sanitizeToolCallInputs", () => {
     const out = sanitizeToolCallInputs(input);
     expect(out.map((m) => m.role)).toEqual(["user"]);
   });
+
+  it("strips toolCalls from aborted assistant but keeps text", () => {
+    const input = [
+      {
+        role: "assistant",
+        stopReason: "aborted",
+        content: [
+          { type: "text", text: "Let me try" },
+          { type: "toolCall", id: "call_1", name: "write", arguments: { path: "/tmp/x" } },
+        ],
+      },
+      { role: "user", content: "hello" },
+    ] as AgentMessage[];
+
+    const out = sanitizeToolCallInputs(input);
+    expect(out).toHaveLength(2);
+    expect(out[0]?.role).toBe("assistant");
+    const assistant = out[0] as Extract<AgentMessage, { role: "assistant" }>;
+    const types = Array.isArray(assistant.content)
+      ? assistant.content.map((b) => (b as { type?: unknown }).type)
+      : [];
+    expect(types).toEqual(["text"]);
+    expect(out[1]?.role).toBe("user");
+  });
+
+  it("drops aborted assistant entirely when only toolCalls remain", () => {
+    const input = [
+      {
+        role: "assistant",
+        stopReason: "aborted",
+        content: [
+          { type: "toolCall", id: "call_1", name: "write", arguments: { path: "/tmp/x" } },
+        ],
+      },
+      { role: "user", content: "hello" },
+    ] as AgentMessage[];
+
+    const out = sanitizeToolCallInputs(input);
+    expect(out.map((m) => m.role)).toEqual(["user"]);
+  });
+
+  it("strips toolCalls from error assistant messages", () => {
+    const input = [
+      {
+        role: "assistant",
+        stopReason: "error",
+        content: [
+          { type: "toolCall", id: "call_1", name: "read", arguments: { path: "a" } },
+        ],
+      },
+      { role: "user", content: "retry" },
+    ] as AgentMessage[];
+
+    const out = sanitizeToolCallInputs(input);
+    expect(out.map((m) => m.role)).toEqual(["user"]);
+  });
+
+  it("prevents orphan toolResults when aborted assistant is followed by user message", () => {
+    // Full scenario: aborted assistant with toolCall → sanitizeToolCallInputs strips toolCall
+    // → sanitizeToolUseResultPairing should NOT insert synthetic toolResult
+    const input = [
+      {
+        role: "assistant",
+        stopReason: "aborted",
+        content: [
+          { type: "toolCall", id: "call_1", name: "write", arguments: { path: "/tmp/x" } },
+        ],
+      },
+      { role: "user", content: "continue" },
+      { role: "assistant", content: [] },
+    ] as AgentMessage[];
+
+    // Run both sanitizers in sequence (same as transformContext)
+    const step1 = sanitizeToolCallInputs(input);
+    const step2 = sanitizeToolUseResultPairing(step1);
+
+    // No orphan toolResults should exist
+    const toolResults = step2.filter((m) => m.role === "toolResult");
+    expect(toolResults).toHaveLength(0);
+  });
 });
diff --git a/packages/core/src/agent/session/session-transcript-repair.ts b/packages/core/src/agent/session/session-transcript-repair.ts
index 9a082cd1..1143ea6e 100644
--- a/packages/core/src/agent/session/session-transcript-repair.ts
+++ b/packages/core/src/agent/session/session-transcript-repair.ts
@@ -118,6 +118,25 @@ export function repairToolCallInputs(messages: AgentMessage[]): ToolCallInputRep
       continue;
     }
 
+    // Drop toolCalls from aborted/error assistant messages.
+    // pi-ai's transformMessages drops these messages entirely, so any
+    // synthetic toolResults we'd insert would become orphaned.
+    const stopReason = (msg as { stopReason?: unknown }).stopReason;
+    if (stopReason === "aborted" || stopReason === "error") {
+      const filtered = msg.content.filter((block: unknown) => !isToolCallBlock(block));
+      if (filtered.length === 0) {
+        droppedAssistantMessages += 1;
+        changed = true;
+        continue;
+      }
+      if (filtered.length !== msg.content.length) {
+        droppedToolCalls += msg.content.length - filtered.length;
+        changed = true;
+        out.push({ ...msg, content: filtered });
+        continue;
+      }
+    }
+
     const nextContent = [];
     let droppedInMessage = 0;