From aa415b25534bd3b861b44d3ba3f4d0fe524fde4f Mon Sep 17 00:00:00 2001
From: Vincent Koc <vincentkoc@ieee.org>
Date: Sun, 12 Apr 2026 05:19:41 +0100
Subject: [PATCH] test(agents): share context pruning trim fixtures

---
 .../pi-hooks/context-pruning/pruner.test.ts   | 125 ++++++++----------
 1 file changed, 57 insertions(+), 68 deletions(-)

diff --git a/src/agents/pi-hooks/context-pruning/pruner.test.ts b/src/agents/pi-hooks/context-pruning/pruner.test.ts
index 7343fd307e8..809dda94204 100644
--- a/src/agents/pi-hooks/context-pruning/pruner.test.ts
+++ b/src/agents/pi-hooks/context-pruning/pruner.test.ts
@@ -10,6 +10,9 @@ type AssistantContentBlock = AssistantMessage["content"][number];
 const CONTEXT_WINDOW_1M = {
   model: { contextWindow: 1_000_000 },
 } as unknown as ExtensionContext;
+const CONTEXT_WINDOW_5K = {
+  model: { contextWindow: 5_000 },
+} as unknown as ExtensionContext;
 
 function makeUser(text: string): AgentMessage {
   return {
@@ -58,6 +61,43 @@ function makeToolResult(
   } as AgentMessage;
 }
 
+function pruneWithOversizedAssistantThinking(params: {
+  assistantBlock: AssistantContentBlock;
+  dropThinkingBlocksForEstimate?: boolean;
+}) {
+  return pruneContextMessages({
+    messages: [
+      makeUser("hello"),
+      makeToolResult([{ type: "text", text: "X".repeat(2_000) }]),
+      makeAssistant([params.assistantBlock, { type: "text", text: "done" }]),
+    ],
+    settings: {
+      ...buildToolTrimSettings(),
+    },
+    ctx: CONTEXT_WINDOW_5K,
+    isToolPrunable: () => true,
+    ...(params.dropThinkingBlocksForEstimate ? { dropThinkingBlocksForEstimate: true } : {}),
+  });
+}
+
+function buildToolTrimSettings() {
+  return {
+    keepLastAssistants: 1,
+    softTrimRatio: 0.5,
+    softTrim: { maxChars: 200, headChars: 100, tailChars: 50 },
+    hardClear: { ...DEFAULT_CONTEXT_PRUNING_SETTINGS.hardClear, enabled: false },
+  };
+}
+
+function expectToolResultWasTrimmed(result: AgentMessage[]) {
+  const toolResult = result.find((message) => message.role === "toolResult") as Extract<
+    AgentMessage,
+    { role: "toolResult" }
+  >;
+  const textBlock = toolResult.content[0] as { type: "text"; text: string };
+  expect(textBlock.text).toContain("[Tool result trimmed:");
+}
+
 describe("pruneContextMessages", () => {
   it("does not crash on assistant message with malformed thinking block (missing thinking string)", () => {
     const messages: AgentMessage[] = [
@@ -124,74 +164,26 @@ describe("pruneContextMessages", () => {
   });
 
   it("counts thinkingSignature bytes when estimating assistant message size", () => {
-    const messages: AgentMessage[] = [
-      makeUser("hello"),
-      makeToolResult([{ type: "text", text: "X".repeat(2_000) }]),
-      makeAssistant([
-        {
-          type: "thinking",
-          thinking: "[redacted]",
-          thinkingSignature: "S".repeat(40_000),
-          redacted: true,
-        } as unknown as AssistantContentBlock,
-        { type: "text", text: "done" },
-      ]),
-    ];
-
-    const result = pruneContextMessages({
-      messages,
-      settings: {
-        ...DEFAULT_CONTEXT_PRUNING_SETTINGS,
-        keepLastAssistants: 1,
-        softTrimRatio: 0.5,
-        softTrim: { maxChars: 200, headChars: 100, tailChars: 50 },
-        hardClear: { ...DEFAULT_CONTEXT_PRUNING_SETTINGS.hardClear, enabled: false },
-      },
-      ctx: { model: { contextWindow: 5_000 } } as unknown as ExtensionContext,
-      isToolPrunable: () => true,
+    const result = pruneWithOversizedAssistantThinking({
+      assistantBlock: {
+        type: "thinking",
+        thinking: "[redacted]",
+        thinkingSignature: "S".repeat(40_000),
+        redacted: true,
+      } as unknown as AssistantContentBlock,
     });
-
-    const toolResult = result.find((message) => message.role === "toolResult") as Extract<
-      AgentMessage,
-      { role: "toolResult" }
-    >;
-    const textBlock = toolResult.content[0] as { type: "text"; text: string };
-    expect(textBlock.text).toContain("[Tool result trimmed:");
+    expectToolResultWasTrimmed(result);
   });
 
   it("counts redacted_thinking data bytes when estimating assistant message size", () => {
-    const messages: AgentMessage[] = [
-      makeUser("hello"),
-      makeToolResult([{ type: "text", text: "X".repeat(2_000) }]),
-      makeAssistant([
-        {
-          type: "redacted_thinking",
-          data: "D".repeat(40_000),
-          thinkingSignature: "sig",
-        } as unknown as AssistantContentBlock,
-        { type: "text", text: "done" },
-      ]),
-    ];
-
-    const result = pruneContextMessages({
-      messages,
-      settings: {
-        ...DEFAULT_CONTEXT_PRUNING_SETTINGS,
-        keepLastAssistants: 1,
-        softTrimRatio: 0.5,
-        softTrim: { maxChars: 200, headChars: 100, tailChars: 50 },
-        hardClear: { ...DEFAULT_CONTEXT_PRUNING_SETTINGS.hardClear, enabled: false },
-      },
-      ctx: { model: { contextWindow: 5_000 } } as unknown as ExtensionContext,
-      isToolPrunable: () => true,
+    const result = pruneWithOversizedAssistantThinking({
+      assistantBlock: {
+        type: "redacted_thinking",
+        data: "D".repeat(40_000),
+        thinkingSignature: "sig",
+      } as unknown as AssistantContentBlock,
     });
-
-    const toolResult = result.find((message) => message.role === "toolResult") as Extract<
-      AgentMessage,
-      { role: "toolResult" }
-    >;
-    const textBlock = toolResult.content[0] as { type: "text"; text: string };
-    expect(textBlock.text).toContain("[Tool result trimmed:");
+    expectToolResultWasTrimmed(result);
   });
 
   it("ignores non-latest thinking signatures that will be dropped before send", () => {
@@ -214,12 +206,9 @@ describe("pruneContextMessages", () => {
       messages,
       settings: {
         ...DEFAULT_CONTEXT_PRUNING_SETTINGS,
-        keepLastAssistants: 1,
-        softTrimRatio: 0.5,
-        softTrim: { maxChars: 200, headChars: 100, tailChars: 50 },
-        hardClear: { ...DEFAULT_CONTEXT_PRUNING_SETTINGS.hardClear, enabled: false },
+        ...buildToolTrimSettings(),
       },
-      ctx: { model: { contextWindow: 5_000 } } as unknown as ExtensionContext,
+      ctx: CONTEXT_WINDOW_5K,
       isToolPrunable: () => true,
       dropThinkingBlocksForEstimate: true,
     });