From 44afe200e6c8728e401d560f905bd0d59f332413 Mon Sep 17 00:00:00 2001
From: Shakker <shakkerdroid@gmail.com>
Date: Sun, 12 Apr 2026 07:11:21 +0100
Subject: [PATCH] fix: scope embedded replay tool-result stripping

---
 .../run/attempt.tool-call-normalization.ts    | 27 ++++++++++++++++---
 1 file changed, 23 insertions(+), 4 deletions(-)

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 94549026813..128f3c8c516 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
@@ -439,15 +439,30 @@ function extractAnthropicReplayToolResultIds(block: AnthropicToolResultContentBl
   return ids;
 }
 
+function isSignedThinkingReplayAssistantSpan(message: AgentMessage | undefined): boolean {
+  if (!message || typeof message !== "object" || message.role !== "assistant") {
+    return false;
+  }
+  const content = (message as { content?: unknown }).content;
+  if (!Array.isArray(content)) {
+    return false;
+  }
+  return (
+    content.some((block) => isThinkingLikeReplayBlock(block)) &&
+    content.some((block) => isReplayToolCallBlock(block))
+  );
+}
+
 function sanitizeAnthropicReplayToolResults(
   messages: AgentMessage[],
   options?: {
-    allowEmbeddedUserToolResults?: boolean;
+    disallowEmbeddedUserToolResultsForSignedThinkingReplay?: boolean;
   },
 ): AgentMessage[] {
   let changed = false;
   const out: AgentMessage[] = [];
-  const allowEmbeddedUserToolResults = options?.allowEmbeddedUserToolResults !== false;
+  const disallowEmbeddedUserToolResultsForSignedThinkingReplay =
+    options?.disallowEmbeddedUserToolResultsForSignedThinkingReplay === true;
 
   for (let index = 0; index < messages.length; index += 1) {
     const message = messages[index];
@@ -461,6 +476,9 @@ function sanitizeAnthropicReplayToolResults(
     }
 
     const previous = messages[index - 1];
+    const shouldStripEmbeddedToolResults =
+      disallowEmbeddedUserToolResultsForSignedThinkingReplay &&
+      isSignedThinkingReplayAssistantSpan(previous);
     const validToolUseIds = new Set<string>();
     if (previous && typeof previous === "object" && previous.role === "assistant") {
       const previousContent = (previous as { content?: unknown }).content;
@@ -489,7 +507,7 @@ function sanitizeAnthropicReplayToolResults(
       if (typedBlock.type !== "toolResult" && typedBlock.type !== "tool") {
         return true;
       }
-      if (!allowEmbeddedUserToolResults) {
+      if (shouldStripEmbeddedToolResults) {
         changed = true;
         return false;
       }
@@ -702,7 +720,8 @@ export function wrapStreamFnSanitizeMalformedToolCalls(
       : sanitized.messages;
     if (transcriptPolicy?.validateAnthropicTurns) {
       nextMessages = sanitizeAnthropicReplayToolResults(nextMessages, {
-        allowEmbeddedUserToolResults: !allowProviderOwnedThinkingReplay,
+        disallowEmbeddedUserToolResultsForSignedThinkingReplay:
+          allowProviderOwnedThinkingReplay,
       });
     }
     if (nextMessages === messages) {