fix(pruner): use dropAllThinkingBlocks for estimate when policy strips all thinking

SymbolStar · SymbolStar · commit e97c7bf188b3 · 2026-05-26T14:20:59.000+08:00
Reviewer feedback on PR #81534: P2-1: pruner was conflating dropThinkingBlocks (keeps latest assistant thinking block) and dropAllThinkingBlocks (strips every assistant thinking block) into a single boolean on the context-pruning runtime. For Copilot Claude proxy paths transcriptPolicy.dropAllThinkingBlocks is true, but the pruner only called dropThinkingBlocks() when estimating context size, so big non-latest thinking signatures still counted toward the token estimate and the pruner could trim tool results that would not actually have been sent. Thread a separate dropAllThinkingBlocksForEstimate signal through ContextPruningRuntimeValue -> extension.ts -> pruneContextMessages so the estimate path matches what the wire format will look like. P2-2: attempt.spawn-workspace.test-support.ts mocked ../thinking.js with only dropReasoningFromHistory + dropThinkingBlocks. attempt.ts also imports dropAllThinkingBlocks now, so module evaluation would fail. Add the missing field to the mock. Tests added: - pruner.test.ts: dropAllThinkingBlocksForEstimate strips even the latest assistant thinking block from the estimate, allowing the oversized tool result through unchanged. - context-pruning.test.ts: updated setContextPruningRuntime callers for new mandatory dropAllThinkingBlocks field.
diff --git a/src/agents/pi-embedded-runner/extensions.ts b/src/agents/pi-embedded-runner/extensions.ts
@@ -128,7 +128,8 @@ function buildContextPruningFactory(params: {
     settings,
     contextWindowTokens: resolveContextWindowTokens(params),
     isToolPrunable: makeToolPrunablePredicate(settings.tools),
-    dropThinkingBlocks: transcriptPolicy.dropThinkingBlocks || transcriptPolicy.dropAllThinkingBlocks,
+    dropThinkingBlocks: transcriptPolicy.dropThinkingBlocks === true,
+    dropAllThinkingBlocks: transcriptPolicy.dropAllThinkingBlocks === true,
     lastCacheTouchAt: readLastCacheTtlTimestamp(params.sessionManager, {
       provider: params.provider,
       modelId: params.modelId,
diff --git a/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.test-support.ts b/src/agents/pi-embedded-runner/run/attempt.spawn-workspace.test-support.ts
@@ -752,6 +752,7 @@ vi.mock("../sandbox-info.js", () => ({
 vi.mock("../thinking.js", () => ({
   dropReasoningFromHistory: <T>(messages: T) => messages,
   dropThinkingBlocks: <T>(messages: T) => messages,
+  dropAllThinkingBlocks: <T>(messages: T) => messages,
 }));
 
 vi.mock("../tool-name-allowlist.js", async (importOriginal) => {
diff --git a/src/agents/pi-hooks/context-pruning.test.ts b/src/agents/pi-hooks/context-pruning.test.ts
@@ -318,6 +318,7 @@ describe("context-pruning", () => {
       contextWindowTokens: 1000,
       isToolPrunable: () => true,
       dropThinkingBlocks: false,
+      dropAllThinkingBlocks: false,
       lastCacheTouchAt: Date.now() - DEFAULT_CONTEXT_PRUNING_SETTINGS.ttlMs - 1000,
     });
 
@@ -341,6 +342,7 @@ describe("context-pruning", () => {
       contextWindowTokens: 1000,
       isToolPrunable: () => true,
       dropThinkingBlocks: false,
+      dropAllThinkingBlocks: false,
       lastCacheTouchAt: lastTouch,
     });
 
diff --git a/src/agents/pi-hooks/context-pruning/extension.ts b/src/agents/pi-hooks/context-pruning/extension.ts
@@ -27,6 +27,7 @@ export default function contextPruningExtension(api: ExtensionAPI): void {
       isToolPrunable: runtime.isToolPrunable,
       contextWindowTokensOverride: runtime.contextWindowTokens ?? undefined,
       dropThinkingBlocksForEstimate: runtime.dropThinkingBlocks,
+      dropAllThinkingBlocksForEstimate: runtime.dropAllThinkingBlocks,
     });
 
     if (next === event.messages) {
diff --git a/src/agents/pi-hooks/context-pruning/pruner.test.ts b/src/agents/pi-hooks/context-pruning/pruner.test.ts
@@ -64,6 +64,7 @@ function makeToolResult(
 function pruneWithOversizedAssistantThinking(params: {
   assistantBlock: AssistantContentBlock;
   dropThinkingBlocksForEstimate?: boolean;
+  dropAllThinkingBlocksForEstimate?: boolean;
 }) {
   return pruneContextMessages({
     messages: [
@@ -77,6 +78,9 @@ function pruneWithOversizedAssistantThinking(params: {
     ctx: CONTEXT_WINDOW_5K,
     isToolPrunable: () => true,
     ...(params.dropThinkingBlocksForEstimate ? { dropThinkingBlocksForEstimate: true } : {}),
+    ...(params.dropAllThinkingBlocksForEstimate
+      ? { dropAllThinkingBlocksForEstimate: true }
+      : {}),
   });
 }
 
@@ -356,6 +360,23 @@ describe("pruneContextMessages", () => {
     expect(result).toBe(messages);
   });
 
+  it("ignores latest-turn thinking signatures when dropAllThinkingBlocksForEstimate is true", () => {
+    const result = pruneWithOversizedAssistantThinking({
+      assistantBlock: {
+        type: "thinking",
+        thinking: "internal",
+        thinkingSignature: "S".repeat(40_000),
+      } as unknown as AssistantContentBlock,
+      dropAllThinkingBlocksForEstimate: true,
+    });
+    const toolResult = result.find((m) => m.role === "toolResult") as Extract<
+      AgentMessage,
+      { role: "toolResult" }
+    >;
+    const textBlock = toolResult.content[0] as { type: "text"; text: string };
+    expect(textBlock.text).toBe("X".repeat(2_000));
+  });
+
   it("soft-trims image-containing tool results by replacing image blocks with placeholders", () => {
     const messages: AgentMessage[] = [
       makeUser("summarize this"),
diff --git a/src/agents/pi-hooks/context-pruning/pruner.ts b/src/agents/pi-hooks/context-pruning/pruner.ts
@@ -2,7 +2,7 @@ import type { AgentMessage } from "@earendil-works/pi-agent-core";
 import type { ImageContent, TextContent, ToolResultMessage } from "@earendil-works/pi-ai";
 import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
 import { CHARS_PER_TOKEN_ESTIMATE, estimateStringChars } from "../../../utils/cjk-chars.js";
-import { dropThinkingBlocks } from "../../pi-embedded-runner/thinking.js";
+import { dropAllThinkingBlocks, dropThinkingBlocks } from "../../pi-embedded-runner/thinking.js";
 import type { EffectiveContextPruningSettings } from "./settings.js";
 import { makeToolPrunablePredicate } from "./tools.js";
 
@@ -291,6 +291,7 @@ export function pruneContextMessages(params: {
   isToolPrunable?: (toolName: string) => boolean;
   contextWindowTokensOverride?: number;
   dropThinkingBlocksForEstimate?: boolean;
+  dropAllThinkingBlocksForEstimate?: boolean;
 }): AgentMessage[] {
   const { messages, settings, ctx } = params;
   const contextWindowTokens =
@@ -320,9 +321,11 @@ export function pruneContextMessages(params: {
   const pruneStartIndex = firstUserIndex === null ? messages.length : firstUserIndex;
 
   const isToolPrunable = params.isToolPrunable ?? makeToolPrunablePredicate(settings.tools);
-  const estimatedMessages = params.dropThinkingBlocksForEstimate
-    ? dropThinkingBlocks(messages)
-    : messages;
+  const estimatedMessages = params.dropAllThinkingBlocksForEstimate
+    ? dropAllThinkingBlocks(messages)
+    : params.dropThinkingBlocksForEstimate
+      ? dropThinkingBlocks(messages)
+      : messages;
 
   const totalCharsBefore = estimateContextChars(estimatedMessages);
   let totalChars = totalCharsBefore;
diff --git a/src/agents/pi-hooks/context-pruning/runtime.ts b/src/agents/pi-hooks/context-pruning/runtime.ts
@@ -6,6 +6,7 @@ export type ContextPruningRuntimeValue = {
   contextWindowTokens?: number | null;
   isToolPrunable: (toolName: string) => boolean;
   dropThinkingBlocks: boolean;
+  dropAllThinkingBlocks: boolean;
   lastCacheTouchAt?: number | null;
 };
 
