fix: preserve anthropic thinking blocks on replay

obviyus · obviyus · commit af6c237e21bd · 2026-04-01T15:12:22.000+05:30
diff --git a/src/agents/pi-embedded-helpers/bootstrap.ts b/src/agents/pi-embedded-helpers/bootstrap.ts
@@ -65,6 +65,10 @@ export function stripThoughtSignatures<T>(
     if (!block || typeof block !== "object") {
       return block;
     }
+    const blockType = (block as { type?: unknown }).type;
+    if (blockType === "thinking" || blockType === "redacted_thinking") {
+      return block;
+    }
     const rec = block as ContentBlockWithSignature;
     const stripSnake = shouldStripSignature(rec.thought_signature);
     const stripCamel = includeCamelCase ? shouldStripSignature(rec.thoughtSignature) : false;
diff --git a/src/agents/pi-embedded-runner/proxy-stream-wrappers.ts b/src/agents/pi-embedded-runner/proxy-stream-wrappers.ts
@@ -71,17 +71,32 @@ export function createOpenRouterSystemCacheWrapper(baseStreamFn: StreamFn | unde
       const messages = payloadObj.messages;
       if (Array.isArray(messages)) {
         for (const msg of messages as Array<{ role?: string; content?: unknown }>) {
-          if (msg.role !== "system" && msg.role !== "developer") {
+          if (msg.role === "system" || msg.role === "developer") {
+            if (typeof msg.content === "string") {
+              msg.content = [
+                { type: "text", text: msg.content, cache_control: { type: "ephemeral" } },
+              ];
+            } else if (Array.isArray(msg.content) && msg.content.length > 0) {
+              const last = msg.content[msg.content.length - 1];
+              if (last && typeof last === "object") {
+                const record = last as Record<string, unknown>;
+                if (record.type !== "thinking" && record.type !== "redacted_thinking") {
+                  record.cache_control = { type: "ephemeral" };
+                }
+              }
+            }
             continue;
           }
-          if (typeof msg.content === "string") {
-            msg.content = [
-              { type: "text", text: msg.content, cache_control: { type: "ephemeral" } },
-            ];
-          } else if (Array.isArray(msg.content) && msg.content.length > 0) {
-            const last = msg.content[msg.content.length - 1];
-            if (last && typeof last === "object") {
-              (last as Record<string, unknown>).cache_control = { type: "ephemeral" };
+
+          if (msg.role === "assistant" && Array.isArray(msg.content)) {
+            for (const block of msg.content) {
+              if (!block || typeof block !== "object") {
+                continue;
+              }
+              const record = block as Record<string, unknown>;
+              if (record.type === "thinking" || record.type === "redacted_thinking") {
+                delete record.cache_control;
+              }
             }
           }
         }
diff --git a/src/agents/pi-embedded-runner/thinking.ts b/src/agents/pi-embedded-runner/thinking.ts
@@ -12,6 +12,15 @@ export function isAssistantMessageWithContent(message: AgentMessage): message is
   );
 }
 
+function isThinkingBlock(block: AssistantContentBlock): boolean {
+  return (
+    !!block &&
+    typeof block === "object" &&
+    ((block as { type?: unknown }).type === "thinking" ||
+      (block as { type?: unknown }).type === "redacted_thinking")
+  );
+}
+
 /**
  * Strip all `type: "thinking"` content blocks from assistant messages.
  *
@@ -23,17 +32,30 @@ export function isAssistantMessageWithContent(message: AgentMessage): message is
  * use reference equality to skip downstream work).
  */
 export function dropThinkingBlocks(messages: AgentMessage[]): AgentMessage[] {
+  let latestAssistantIndex = -1;
+  for (let i = messages.length - 1; i >= 0; i -= 1) {
+    if (isAssistantMessageWithContent(messages[i])) {
+      latestAssistantIndex = i;
+      break;
+    }
+  }
+
   let touched = false;
   const out: AgentMessage[] = [];
-  for (const msg of messages) {
+  for (let i = 0; i < messages.length; i += 1) {
+    const msg = messages[i];
     if (!isAssistantMessageWithContent(msg)) {
       out.push(msg);
       continue;
     }
+    if (i === latestAssistantIndex) {
+      out.push(msg);
+      continue;
+    }
     const nextContent: AssistantContentBlock[] = [];
     let changed = false;
     for (const block of msg.content) {
-      if (block && typeof block === "object" && (block as { type?: unknown }).type === "thinking") {
+      if (isThinkingBlock(block)) {
         touched = true;
         changed = true;
         continue;
diff --git a/src/agents/pi-hooks/context-pruning/pruner.ts b/src/agents/pi-hooks/context-pruning/pruner.ts
@@ -146,8 +146,16 @@ function estimateMessageChars(message: AgentMessage): number {
       if (b.type === "text" && typeof b.text === "string") {
         chars += estimateWeightedTextChars(b.text);
       }
-      if (b.type === "thinking" && typeof b.thinking === "string") {
-        chars += estimateWeightedTextChars(b.thinking);
+      const blockType = (b as { type?: unknown }).type;
+      if (blockType === "thinking" || blockType === "redacted_thinking") {
+        const thinking = (b as { thinking?: unknown }).thinking;
+        if (typeof thinking === "string") {
+          chars += estimateWeightedTextChars(thinking);
+        }
+        const signature = (b as { thinkingSignature?: unknown }).thinkingSignature;
+        if (typeof signature === "string") {
+          chars += estimateWeightedTextChars(signature);
+        }
       }
       if (b.type === "toolCall") {
         try {
