From 92f9e09a8e7f2f3e58ae3ea90a0984c497f35c17 Mon Sep 17 00:00:00 2001
From: Shakker <shakkerdroid@gmail.com>
Date: Sun, 12 Apr 2026 03:36:54 +0100
Subject: [PATCH] fix: preserve signed thinking turns during replay
 normalization

---
 .../pi-embedded-runner/run/attempt.test.ts    | 31 +++++++++++++++++++
 .../run/attempt.tool-call-normalization.ts    | 12 +++++++
 2 files changed, 43 insertions(+)

diff --git a/src/agents/pi-embedded-runner/run/attempt.test.ts b/src/agents/pi-embedded-runner/run/attempt.test.ts
index 648e1780618..9d073728912 100644
--- a/src/agents/pi-embedded-runner/run/attempt.test.ts
+++ b/src/agents/pi-embedded-runner/run/attempt.test.ts
@@ -946,6 +946,37 @@ describe("wrapStreamFnSanitizeMalformedToolCalls", () => {
     expect(seenContext.messages).toBe(messages);
   });
 
+  it("preserves signed thinking turns when replayed tool calls would otherwise be sanitized", async () => {
+    const messages = [
+      {
+        role: "assistant",
+        content: [
+          { type: "thinking", thinking: "internal", thinkingSignature: "sig_1" },
+          { type: "toolCall", id: "toolu_legacy", name: "gateway", arguments: {} },
+        ],
+      },
+      {
+        role: "user",
+        content: [{ type: "text", text: "retry" }],
+      },
+    ];
+    const baseFn = vi.fn((_model, _context) =>
+      createFakeStream({ events: [], resultMessage: { role: "assistant", content: [] } }),
+    );
+
+    const wrapped = wrapStreamFnSanitizeMalformedToolCalls(baseFn as never, new Set(["read"]));
+    const stream = wrapped(
+      { api: "anthropic-messages" } as never,
+      { messages } as never,
+      {} as never,
+    ) as FakeWrappedStream | Promise<FakeWrappedStream>;
+    await Promise.resolve(stream);
+
+    expect(baseFn).toHaveBeenCalledTimes(1);
+    const seenContext = baseFn.mock.calls[0]?.[1] as { messages: unknown[] };
+    expect(seenContext.messages).toBe(messages);
+  });
+
   it("preserves sessions_spawn attachment payloads on replay", async () => {
     const attachmentContent = "INLINE_ATTACHMENT_PAYLOAD";
     const messages = [
diff --git a/src/agents/pi-embedded-runner/run/attempt.tool-call-normalization.ts b/src/agents/pi-embedded-runner/run/attempt.tool-call-normalization.ts
index 883ae110907..00e1cc9f6ef 100644
--- a/src/agents/pi-embedded-runner/run/attempt.tool-call-normalization.ts
+++ b/src/agents/pi-embedded-runner/run/attempt.tool-call-normalization.ts
@@ -229,6 +229,14 @@ type AnthropicToolResultContentBlock = {
   toolUseId?: unknown;
 };
 
+function isThinkingLikeReplayBlock(block: unknown): boolean {
+  if (!block || typeof block !== "object") {
+    return false;
+  }
+  const type = (block as { type?: unknown }).type;
+  return type === "thinking" || type === "redacted_thinking";
+}
+
 function isReplayToolCallBlock(block: unknown): block is ReplayToolCallBlock {
   if (!block || typeof block !== "object") {
     return false;
@@ -283,6 +291,10 @@ function sanitizeReplayToolCallInputs(
       out.push(message);
       continue;
     }
+    if (message.content.some((block) => isThinkingLikeReplayBlock(block))) {
+      out.push(message);
+      continue;
+    }
 
     const nextContent: typeof message.content = [];
     let messageChanged = false;