fix: harden tool-result overflow recovery (#61651)

2026-05-04 10:00:21 +00:00 · 2026-04-06 13:48:40 +01:00
parent 4917009ac7
commit a42ee69ad4
9 changed files with 295 additions and 187 deletions
--- a/src/agents/pi-embedded-runner/run.overflow-compaction.loop.test.ts
+++ b/src/agents/pi-embedded-runner/run.overflow-compaction.loop.test.ts
@@ -18,6 +18,7 @@ import {
  mockedTruncateOversizedToolResultsInSession,
  overflowBaseRunParams as baseParams,
 } from "./run.overflow-compaction.harness.js";
+import type { EmbeddedRunAttemptResult } from "./run/types.js";

 let runEmbeddedPiAgent: typeof import("./run.js").runEmbeddedPiAgent;

@@ -278,7 +279,9 @@ describe("overflow compaction in run loop", () => {
    expect(mockedTruncateOversizedToolResultsInSession).not.toHaveBeenCalled();
    expect(mockedRunEmbeddedAttempt).toHaveBeenCalledTimes(2);
    expect(mockedLog.warn).toHaveBeenCalledWith(
-      expect.stringContaining("context overflow detected (attempt 1/3); attempting auto-compaction"),
+      expect.stringContaining(
+        "context overflow detected (attempt 1/3); attempting auto-compaction",
+      ),
    );
    expect(result.meta.error).toBeUndefined();
  });
--- a/src/agents/pi-embedded-runner/run/attempt.ts
+++ b/src/agents/pi-embedded-runner/run/attempt.ts
@@ -1769,17 +1769,18 @@ export async function runEmbeddedAttempt(
          }

          const reserveTokens = settingsManager.getCompactionReserveTokens();
+          const contextTokenBudget = params.contextTokenBudget ?? DEFAULT_CONTEXT_TOKENS;
          const preemptiveCompaction = shouldPreemptivelyCompactBeforePrompt({
            messages: activeSession.messages,
            systemPrompt: systemPromptText,
            prompt: effectivePrompt,
-            contextTokenBudget: params.contextTokenBudget,
+            contextTokenBudget,
            reserveTokens,
          });
          if (preemptiveCompaction.route === "truncate_tool_results_only") {
            const truncationResult = truncateOversizedToolResultsInSessionManager({
              sessionManager,
-              contextWindowTokens: params.contextTokenBudget,
+              contextWindowTokens: contextTokenBudget,
              sessionFile: params.sessionFile,
              sessionId: params.sessionId,
              sessionKey: params.sessionKey,
--- a/src/agents/pi-embedded-runner/run/preemptive-compaction.test.ts
+++ b/src/agents/pi-embedded-runner/run/preemptive-compaction.test.ts
@@ -1,3 +1,4 @@
+import type { AgentMessage } from "@mariozechner/pi-agent-core";
 import { describe, expect, it } from "vitest";
 import { estimateToolResultReductionPotential } from "../tool-result-truncation.js";
 import {
@@ -6,6 +7,27 @@ import {
  shouldPreemptivelyCompactBeforePrompt,
 } from "./preemptive-compaction.js";

+let timestamp = 1;
+
+function makeAssistantHistory(text: string): AgentMessage {
+  return {
+    role: "assistant",
+    content: [{ type: "text", text }],
+    timestamp: timestamp++,
+  } as AgentMessage;
+}
+
+function makeToolResultMessage(...texts: string[]): AgentMessage {
+  return {
+    role: "toolResult",
+    toolCallId: `call_${timestamp}`,
+    toolName: "read",
+    content: texts.map((text) => ({ type: "text", text })),
+    isError: false,
+    timestamp: timestamp++,
+  } as AgentMessage;
+}
+
 describe("preemptive-compaction", () => {
  const verboseHistory =
    "alpha beta gamma delta epsilon zeta eta theta iota kappa lambda mu ".repeat(40);
@@ -21,12 +43,12 @@ describe("preemptive-compaction", () => {

  it("raises the estimate as prompt-side content grows", () => {
    const smaller = estimatePrePromptTokens({
-      messages: [{ role: "assistant", content: verboseHistory }],
+      messages: [makeAssistantHistory(verboseHistory)],
      systemPrompt: "sys",
      prompt: "hello",
    });
    const larger = estimatePrePromptTokens({
-      messages: [{ role: "assistant", content: verboseHistory }],
+      messages: [makeAssistantHistory(verboseHistory)],
      systemPrompt: verboseSystem,
      prompt: verbosePrompt,
    });
@@ -36,7 +58,7 @@ describe("preemptive-compaction", () => {

  it("requests preemptive compaction when the reserve-based prompt budget would be exceeded", () => {
    const result = shouldPreemptivelyCompactBeforePrompt({
-      messages: [{ role: "assistant", content: verboseHistory }],
+      messages: [makeAssistantHistory(verboseHistory)],
      systemPrompt: verboseSystem,
      prompt: verbosePrompt,
      contextTokenBudget: 500,
@@ -50,7 +72,7 @@ describe("preemptive-compaction", () => {

  it("does not request preemptive compaction when the reserve-based prompt budget still fits", () => {
    const result = shouldPreemptivelyCompactBeforePrompt({
-      messages: [{ role: "assistant", content: "short history" }],
+      messages: [makeAssistantHistory("short history")],
      systemPrompt: "sys",
      prompt: "hello",
      contextTokenBudget: 10_000,
@@ -64,17 +86,9 @@ describe("preemptive-compaction", () => {

  it("routes to direct tool-result truncation when recent tool tails can clearly absorb the overflow", () => {
    const medium = "alpha beta gamma delta epsilon ".repeat(2200);
-    const messages = [
-      { role: "assistant", content: "short history" },
-      {
-        role: "toolResult",
-        content: [
-          { type: "text", text: medium },
-          { type: "text", text: medium },
-          { type: "text", text: medium },
-          { type: "text", text: medium },
-        ],
-      } as never,
+    const messages: AgentMessage[] = [
+      makeAssistantHistory("short history"),
+      makeToolResultMessage(medium, medium, medium, medium),
    ];
    const reserveTokens = 2_000;
    const contextTokenBudget = 26_000;
@@ -106,10 +120,10 @@ describe("preemptive-compaction", () => {
      5000,
    );
    const messages = [
-      { role: "assistant", content: longHistory },
-      { role: "toolResult", content: [{ type: "text", text: medium }] } as never,
-      { role: "toolResult", content: [{ type: "text", text: medium }] } as never,
-      { role: "toolResult", content: [{ type: "text", text: medium }] } as never,
+      makeAssistantHistory(longHistory),
+      makeToolResultMessage(medium),
+      makeToolResultMessage(medium),
+      makeToolResultMessage(medium),
    ];
    const reserveTokens = 500;
    const baseContextTokenBudget = 3_500;
@@ -119,7 +133,7 @@ describe("preemptive-compaction", () => {
      prompt: verbosePrompt,
    });
    const toolResultPotential = estimateToolResultReductionPotential({
-      messages: messages as never,
+      messages,
      contextWindowTokens: baseContextTokenBudget,
    });
    const desiredOverflowTokens = Math.ceil((toolResultPotential.maxReducibleChars + 4_096) / 4);
@@ -139,4 +153,40 @@ describe("preemptive-compaction", () => {
    expect(result.overflowTokens).toBeGreaterThan(0);
    expect(result.toolResultReducibleChars).toBeGreaterThan(0);
  });
+
+  it("treats mixed oversized-plus-aggregate tool tails as cumulative recovery potential", () => {
+    const oversized = "x".repeat(45_000);
+    const medium = "alpha beta gamma delta epsilon ".repeat(500);
+    const messages: AgentMessage[] = [
+      makeAssistantHistory("short history"),
+      makeToolResultMessage(oversized),
+      makeToolResultMessage(medium),
+      makeToolResultMessage(medium),
+    ];
+    const reserveTokens = 2_000;
+    const estimatedPromptTokens = estimatePrePromptTokens({
+      messages,
+      systemPrompt: "sys",
+      prompt: "hello",
+    });
+    const potential = estimateToolResultReductionPotential({
+      messages,
+      contextWindowTokens: 128_000,
+    });
+    const desiredOverflowTokens = 2_000;
+    const result = shouldPreemptivelyCompactBeforePrompt({
+      messages,
+      systemPrompt: "sys",
+      prompt: "hello",
+      contextTokenBudget: estimatedPromptTokens - desiredOverflowTokens + reserveTokens,
+      reserveTokens,
+    });
+
+    expect(potential.oversizedReducibleChars).toBeGreaterThan(0);
+    expect(potential.aggregateReducibleChars).toBeGreaterThan(0);
+    expect(potential.oversizedReducibleChars).toBeLessThan(desiredOverflowTokens * 4);
+    expect(potential.maxReducibleChars).toBeGreaterThan(desiredOverflowTokens * 4);
+    expect(result.route).toBe("truncate_tool_results_only");
+    expect(result.shouldCompact).toBe(false);
+  });
 });
--- a/src/agents/pi-embedded-runner/run/preemptive-compaction.ts
+++ b/src/agents/pi-embedded-runner/run/preemptive-compaction.ts
@@ -23,9 +23,13 @@ export function estimatePrePromptTokens(params: {
  const { messages, systemPrompt, prompt } = params;
  const syntheticMessages: AgentMessage[] = [];
  if (typeof systemPrompt === "string" && systemPrompt.trim().length > 0) {
-    syntheticMessages.push({ role: "system", content: systemPrompt } as AgentMessage);
+    syntheticMessages.push({
+      role: "system",
+      content: systemPrompt,
+      timestamp: 0,
+    } as unknown as AgentMessage);
  }
-  syntheticMessages.push({ role: "user", content: prompt } as AgentMessage);
+  syntheticMessages.push({ role: "user", content: prompt, timestamp: 0 } as AgentMessage);

  const estimated =
    estimateMessagesTokens(messages) +
--- a/src/agents/pi-embedded-runner/tool-result-context-guard.test.ts
+++ b/src/agents/pi-embedded-runner/tool-result-context-guard.test.ts
@@ -5,6 +5,7 @@ import {
  CONTEXT_LIMIT_TRUNCATION_NOTICE,
  formatContextLimitTruncationNotice,
  installToolResultContextGuard,
+  PREEMPTIVE_CONTEXT_OVERFLOW_MESSAGE,
 } from "./tool-result-context-guard.js";

 function makeUser(text: string): AgentMessage {
@@ -118,7 +119,7 @@ describe("installToolResultContextGuard", () => {

  it("does not preemptively overflow large non-tool context that is still under the high-water mark", async () => {
    const agent = makeGuardableAgent();
-    const contextForNextCall = [makeUser("u".repeat(50_000))];
+    const contextForNextCall = [makeUser("u".repeat(3_200))];

    const transformed = await applyGuardToContext(agent, contextForNextCall);

@@ -180,19 +181,20 @@ describe("installToolResultContextGuard", () => {
    expect((contextForNextCall[0] as { details?: unknown }).details).toBeDefined();
  });

-  it("does not preemptively overflow when total context remains large after one-time truncation", async () => {
+  it("throws a preemptive overflow when total context still exceeds the high-water mark", async () => {
    const agent = makeGuardableAgent();
    const contextForNextCall = [
      makeUser("u".repeat(50_000)),
-      makeToolResult("call_ok", "x".repeat(500)),
+      makeToolResult("call_big", "x".repeat(5_000)),
    ];

-    const transformed = await applyGuardToContext(agent, contextForNextCall);
-
-    expect(transformed).toBe(contextForNextCall);
+    await expect(applyGuardToContext(agent, contextForNextCall)).rejects.toThrow(
+      PREEMPTIVE_CONTEXT_OVERFLOW_MESSAGE,
+    );
+    expect(getToolResultText(contextForNextCall[1])).toBe("x".repeat(5_000));
  });

-  it("does not rewrite older tool results under aggregate pressure", async () => {
+  it("throws instead of rewriting older tool results under aggregate pressure", async () => {
    const agent = makeGuardableAgent();
    const contextForNextCall = [
      makeUser("u".repeat(50_000)),
@@ -201,15 +203,15 @@ describe("installToolResultContextGuard", () => {
      makeToolResult("call_3", "c".repeat(500)),
    ];

-    const transformed = await applyGuardToContext(agent, contextForNextCall);
-
-    expect(transformed).toBe(contextForNextCall);
+    await expect(applyGuardToContext(agent, contextForNextCall)).rejects.toThrow(
+      PREEMPTIVE_CONTEXT_OVERFLOW_MESSAGE,
+    );
    expect(getToolResultText(contextForNextCall[1])).toBe("a".repeat(500));
    expect(getToolResultText(contextForNextCall[2])).toBe("b".repeat(500));
    expect(getToolResultText(contextForNextCall[3])).toBe("c".repeat(500));
  });

-  it("does not special-case the latest read result under aggregate pressure", async () => {
+  it("does not special-case the latest read result before throwing under aggregate pressure", async () => {
    const agent = makeGuardableAgent();
    const contextForNextCall = [
      makeUser("u".repeat(50_000)),
@@ -217,9 +219,9 @@ describe("installToolResultContextGuard", () => {
      makeReadToolResult("call_new", "y".repeat(500)),
    ];

-    const transformed = await applyGuardToContext(agent, contextForNextCall);
-
-    expect(transformed).toBe(contextForNextCall);
+    await expect(applyGuardToContext(agent, contextForNextCall)).rejects.toThrow(
+      PREEMPTIVE_CONTEXT_OVERFLOW_MESSAGE,
+    );
    expect(getToolResultText(contextForNextCall[1])).toBe("x".repeat(400));
    expect(getToolResultText(contextForNextCall[2])).toBe("y".repeat(500));
  });
--- a/src/agents/pi-embedded-runner/tool-result-context-guard.ts
+++ b/src/agents/pi-embedded-runner/tool-result-context-guard.ts
@@ -1,8 +1,10 @@
 import type { AgentMessage } from "@mariozechner/pi-agent-core";
 import {
+  CHARS_PER_TOKEN_ESTIMATE,
  TOOL_RESULT_CHARS_PER_TOKEN_ESTIMATE,
  type MessageCharEstimateCache,
  createMessageCharEstimateCache,
+  estimateContextChars,
  estimateMessageCharsCached,
  getToolResultText,
  invalidateMessageCharsCacheEntry,
@@ -10,10 +12,12 @@ import {
 } from "./tool-result-char-estimator.js";

 const SINGLE_TOOL_RESULT_CONTEXT_SHARE = 0.5;
+const PREEMPTIVE_OVERFLOW_RATIO = 0.9;

 export const CONTEXT_LIMIT_TRUNCATION_NOTICE = "more characters truncated";
-const TOOL_RESULT_ESTIMATE_TO_TEXT_RATIO =
-  4 / TOOL_RESULT_CHARS_PER_TOKEN_ESTIMATE;
+export const PREEMPTIVE_CONTEXT_OVERFLOW_MESSAGE =
+  "Context overflow: estimated context size exceeds safe threshold during tool loop.";
+const TOOL_RESULT_ESTIMATE_TO_TEXT_RATIO = 4 / TOOL_RESULT_CHARS_PER_TOKEN_ESTIMATE;

 type GuardableTransformContext = (
  messages: AgentMessage[],
@@ -133,6 +137,14 @@ function toolResultsNeedTruncation(params: {
  return false;
 }

+function exceedsPreemptiveOverflowThreshold(params: {
+  messages: AgentMessage[];
+  maxContextChars: number;
+}): boolean {
+  const estimateCache = createMessageCharEstimateCache();
+  return estimateContextChars(params.messages, estimateCache) > params.maxContextChars;
+}
+
 function applyMessageMutationInPlace(
  target: AgentMessage,
  source: AgentMessage,
@@ -176,6 +188,10 @@ export function installToolResultContextGuard(params: {
  contextWindowTokens: number;
 }): () => void {
  const contextWindowTokens = Math.max(1, Math.floor(params.contextWindowTokens));
+  const maxContextChars = Math.max(
+    1_024,
+    Math.floor(contextWindowTokens * CHARS_PER_TOKEN_ESTIMATE * PREEMPTIVE_OVERFLOW_RATIO),
+  );
  const maxSingleToolResultChars = Math.max(
    1_024,
    Math.floor(
@@ -206,6 +222,14 @@ export function installToolResultContextGuard(params: {
        maxSingleToolResultChars,
      });
    }
+    if (
+      exceedsPreemptiveOverflowThreshold({
+        messages: contextMessages,
+        maxContextChars,
+      })
+    ) {
+      throw new Error(PREEMPTIVE_CONTEXT_OVERFLOW_MESSAGE);
+    }

    return contextMessages;
  }) as GuardableTransformContext;
--- a/src/agents/pi-embedded-runner/tool-result-truncation.test.ts
+++ b/src/agents/pi-embedded-runner/tool-result-truncation.test.ts
@@ -490,51 +490,36 @@ describe("truncateOversizedToolResultsInSession", () => {
    expect(text.length).toBeLessThan(2_000);
    expect(text).toContain("truncated");
  });
-
-  it("applies aggregate recovery after oversized truncation for mixed tool-result tails", async () => {
+  it("combines oversized and aggregate recovery truncation in the same session rewrite", async () => {
    const dir = await createTmpDir();
    const sm = SessionManager.create(dir, dir);
    sm.appendMessage(makeUserMessage("hello"));
    sm.appendMessage(makeAssistantMessage("calling tools"));
-    const oversized = "x".repeat(500_000);
+    sm.appendMessage(makeToolResult("x".repeat(500_000), "call_1"));
    const medium = "alpha beta gamma delta epsilon ".repeat(800);
-    sm.appendMessage(makeToolResult(oversized, "call_1"));
    sm.appendMessage(makeToolResult(medium, "call_2"));
    sm.appendMessage(makeToolResult(medium, "call_3"));
    const sessionFile = sm.getSessionFile()!;

-    const beforeBranch = SessionManager.open(sessionFile).getBranch();
-    const beforeToolResults = beforeBranch.filter(
-      (entry) => entry.type === "message" && entry.message.role === "toolResult",
-    );
-    const beforeLengths = beforeToolResults.map((entry) =>
-      entry.type === "message" ? getToolResultTextLength(entry.message) : 0,
-    );
-
    const result = await truncateOversizedToolResultsInSession({
      sessionFile,
-      contextWindowTokens: 128_000,
+      contextWindowTokens: 100,
    });

    expect(result.truncated).toBe(true);
-    expect(result.truncatedCount).toBeGreaterThan(1);
+    expect(result.truncatedCount).toBe(3);

    const afterBranch = SessionManager.open(sessionFile).getBranch();
-    const afterToolResults = afterBranch.filter(
+    const toolResults = afterBranch.filter(
      (entry) => entry.type === "message" && entry.message.role === "toolResult",
    );
-    const afterLengths = afterToolResults.map((entry) =>
-      entry.type === "message" ? getToolResultTextLength(entry.message) : 0,
+    const toolTexts = toolResults.map((entry) =>
+      entry.type === "message" ? getFirstToolResultText(entry.message) : "",
    );

-    expect(afterLengths[0]).toBeLessThan(beforeLengths[0] ?? Infinity);
-    expect(
-      (afterLengths[1] ?? Infinity) < (beforeLengths[1] ?? Infinity) ||
-        (afterLengths[2] ?? Infinity) < (beforeLengths[2] ?? Infinity),
-    ).toBe(true);
-    expect(afterLengths.reduce((sum, value) => sum + value, 0)).toBeLessThan(
-      beforeLengths.reduce((sum, value) => sum + value, 0),
-    );
+    expect(toolTexts[0]).toContain("truncated");
+    expect(toolTexts[1]).toContain("truncated");
+    expect(toolTexts[2].length).toBeGreaterThan(0);
  });
 });

--- a/src/agents/pi-embedded-runner/tool-result-truncation.ts
+++ b/src/agents/pi-embedded-runner/tool-result-truncation.ts
@@ -46,6 +46,18 @@ const DEFAULT_SUFFIX = (truncatedChars: number) =>
 export const MIN_TRUNCATED_TEXT_CHARS = MIN_KEEP_CHARS + DEFAULT_SUFFIX(1).length;
 const RECOVERY_MIN_TRUNCATED_TEXT_CHARS = RECOVERY_MIN_KEEP_CHARS + DEFAULT_SUFFIX(1).length;

+function resolveSuffixFactory(
+  suffix: ToolResultTruncationOptions["suffix"],
+): (truncatedChars: number) => string {
+  if (typeof suffix === "function") {
+    return suffix;
+  }
+  if (typeof suffix === "string") {
+    return () => suffix;
+  }
+  return DEFAULT_SUFFIX;
+}
+
 /**
 * Marker inserted between head and tail when using head+tail truncation.
 */
@@ -81,10 +93,7 @@ export function truncateToolResultText(
  maxChars: number,
  options: ToolResultTruncationOptions = {},
 ): string {
-  const suffixFactory: (truncatedChars: number) => string =
-    typeof options.suffix === "function"
-      ? options.suffix
-      : () => options.suffix ?? DEFAULT_SUFFIX(1);
+  const suffixFactory = resolveSuffixFactory(options.suffix);
  const minKeepChars = options.minKeepChars ?? MIN_KEEP_CHARS;
  if (text.length <= maxChars) {
    return text;
@@ -174,10 +183,7 @@ export function truncateToolResultMessage(
  maxChars: number,
  options: ToolResultTruncationOptions = {},
 ): AgentMessage {
-  const suffixFactory: (truncatedChars: number) => string =
-    typeof options.suffix === "function"
-      ? options.suffix
-      : () => options.suffix ?? DEFAULT_SUFFIX(1);
+  const suffixFactory = resolveSuffixFactory(options.suffix);
  const minKeepChars = options.minKeepChars ?? MIN_KEEP_CHARS;
  const content = (msg as { content?: unknown }).content;
  if (!Array.isArray(content)) {
@@ -267,11 +273,22 @@ export type ToolResultReductionPotential = {
  maxReducibleChars: number;
 };

+type ToolResultBranchEntry = {
+  id: string;
+  type: string;
+  message?: AgentMessage;
+};
+
+type ToolResultReplacement = {
+  entryId: string;
+  message: AgentMessage;
+};
+
 function buildAggregateToolResultReplacements(params: {
-  branch: Array<{ id: string; type: string; message?: AgentMessage }>;
+  branch: ToolResultBranchEntry[];
  aggregateBudgetChars: number;
  minKeepChars?: number;
-}): Array<{ entryId: string; message: AgentMessage }> {
+}): ToolResultReplacement[] {
  const minKeepChars = params.minKeepChars ?? MIN_KEEP_CHARS;
  const minTruncatedTextChars = minKeepChars + DEFAULT_SUFFIX(1).length;
  const candidates = params.branch
@@ -339,30 +356,126 @@ function buildAggregateToolResultReplacements(params: {
  return replacements;
 }

-function applyToolResultReplacementsToBranch(params: {
-  branch: Array<{ id: string; type: string; message?: AgentMessage }>;
-  replacements: Array<{ entryId: string; message: AgentMessage }>;
-}): Array<{ id: string; type: string; message?: AgentMessage }> {
-  if (params.replacements.length === 0) {
-    return params.branch;
+function buildOversizedToolResultReplacements(params: {
+  branch: ToolResultBranchEntry[];
+  maxChars: number;
+  minKeepChars?: number;
+}): ToolResultReplacement[] {
+  const minKeepChars = params.minKeepChars ?? MIN_KEEP_CHARS;
+  const replacements: ToolResultReplacement[] = [];
+
+  for (const entry of params.branch) {
+    if (entry.type !== "message" || !entry.message) {
+      continue;
+    }
+    const msg = entry.message;
+    if ((msg as { role?: string }).role !== "toolResult") {
+      continue;
+    }
+    if (getToolResultTextLength(msg) <= params.maxChars) {
+      continue;
+    }
+    replacements.push({
+      entryId: entry.id,
+      message: truncateToolResultMessage(msg, params.maxChars, {
+        minKeepChars,
+      }),
+    });
  }

-  const replacementsById = new Map(
-    params.replacements.map((replacement) => [replacement.entryId, replacement.message]),
-  );
+  return replacements;
+}

-  return params.branch.map((entry) => {
-    if (entry.type !== "message") {
-      return entry;
+function calculateReplacementReduction(
+  branch: ToolResultBranchEntry[],
+  replacements: ToolResultReplacement[],
+): number {
+  if (replacements.length === 0) {
+    return 0;
+  }
+  const branchById = new Map(branch.map((entry) => [entry.id, entry]));
+  let reduction = 0;
+
+  for (const replacement of replacements) {
+    const entry = branchById.get(replacement.entryId);
+    if (!entry?.message) {
+      continue;
    }
+    reduction += Math.max(
+      0,
+      getToolResultTextLength(entry.message) - getToolResultTextLength(replacement.message),
+    );
+  }
+
+  return reduction;
+}
+
+function applyToolResultReplacementsToBranch(
+  branch: ToolResultBranchEntry[],
+  replacements: ToolResultReplacement[],
+): ToolResultBranchEntry[] {
+  if (replacements.length === 0) {
+    return branch;
+  }
+  const replacementsById = new Map(
+    replacements.map((replacement) => [replacement.entryId, replacement]),
+  );
+  return branch.map((entry) => {
    const replacement = replacementsById.get(entry.id);
-    if (!replacement) {
+    if (!replacement || entry.type !== "message") {
      return entry;
    }
-    return { ...entry, message: replacement };
+    return {
+      ...entry,
+      message: replacement.message,
+    };
  });
 }

+function buildToolResultReplacementPlan(params: {
+  branch: ToolResultBranchEntry[];
+  maxChars: number;
+  aggregateBudgetChars: number;
+  minKeepChars?: number;
+}): {
+  replacements: ToolResultReplacement[];
+  oversizedReplacementCount: number;
+  aggregateReplacementCount: number;
+  oversizedReducibleChars: number;
+  aggregateReducibleChars: number;
+} {
+  const minKeepChars = params.minKeepChars ?? MIN_KEEP_CHARS;
+  const oversizedReplacements = buildOversizedToolResultReplacements({
+    branch: params.branch,
+    maxChars: params.maxChars,
+    minKeepChars,
+  });
+  const oversizedReducibleChars = calculateReplacementReduction(
+    params.branch,
+    oversizedReplacements,
+  );
+  const oversizedTrimmedBranch = applyToolResultReplacementsToBranch(
+    params.branch,
+    oversizedReplacements,
+  );
+  const aggregateReplacements = buildAggregateToolResultReplacements({
+    branch: oversizedTrimmedBranch,
+    aggregateBudgetChars: params.aggregateBudgetChars,
+    minKeepChars,
+  });
+  const aggregateReducibleChars = calculateReplacementReduction(
+    oversizedTrimmedBranch,
+    aggregateReplacements,
+  );
+
+  return {
+    replacements: [...oversizedReplacements, ...aggregateReplacements],
+    oversizedReplacementCount: oversizedReplacements.length,
+    aggregateReplacementCount: aggregateReplacements.length,
+    oversizedReducibleChars,
+    aggregateReducibleChars,
+  };
+}
 export function estimateToolResultReductionPotential(params: {
  messages: AgentMessage[];
  contextWindowTokens: number;
@@ -370,15 +483,15 @@ export function estimateToolResultReductionPotential(params: {
  const { messages, contextWindowTokens } = params;
  const maxChars = calculateMaxToolResultChars(contextWindowTokens);
  const aggregateBudgetChars = calculateRecoveryAggregateToolResultChars(contextWindowTokens);
+  const branch = messages.map((message, index) => ({
+    id: `message-${index}`,
+    type: "message",
+    message,
+  }));

  let toolResultCount = 0;
  let totalToolResultChars = 0;
-  let oversizedCount = 0;
-  let oversizedReducibleChars = 0;
-  const individuallyTrimmedMessages = messages.slice();
-
-  for (let index = 0; index < messages.length; index += 1) {
-    const msg = messages[index];
+  for (const msg of messages) {
    if ((msg as { role?: string }).role !== "toolResult") {
      continue;
    }
@@ -388,50 +501,23 @@ export function estimateToolResultReductionPotential(params: {
    }
    toolResultCount += 1;
    totalToolResultChars += textLength;
-    if (textLength <= maxChars) {
-      continue;
-    }
-    oversizedCount += 1;
-    const truncatedMessage = truncateToolResultMessage(msg, maxChars, {
-      minKeepChars: RECOVERY_MIN_KEEP_CHARS,
-    });
-    individuallyTrimmedMessages[index] = truncatedMessage;
-    oversizedReducibleChars += Math.max(0, textLength - getToolResultTextLength(truncatedMessage));
  }
-
-  const aggregateReplacements = buildAggregateToolResultReplacements({
-    branch: individuallyTrimmedMessages.map((message, index) => ({
-      id: `message-${index}`,
-      type: "message",
-      message,
-    })),
+  const plan = buildToolResultReplacementPlan({
+    branch,
+    maxChars,
    aggregateBudgetChars,
    minKeepChars: RECOVERY_MIN_KEEP_CHARS,
  });
-  const individuallyTrimmedBranch = individuallyTrimmedMessages.map((message, index) => ({
-    id: `message-${index}`,
-    type: "message",
-    message,
-  }));
-  const aggregateReducibleChars = aggregateReplacements.reduce((sum, replacement) => {
-    const match = individuallyTrimmedBranch.find((entry) => entry.id === replacement.entryId);
-    const originalLength =
-      match && match.message
-        ? getToolResultTextLength(match.message)
-        : getToolResultTextLength(replacement.message);
-    const newLength = getToolResultTextLength(replacement.message);
-    return sum + Math.max(0, originalLength - newLength);
-  }, 0);
-  const maxReducibleChars = oversizedReducibleChars + aggregateReducibleChars;
+  const maxReducibleChars = plan.oversizedReducibleChars + plan.aggregateReducibleChars;

  return {
    maxChars,
    aggregateBudgetChars,
    toolResultCount,
    totalToolResultChars,
-    oversizedCount,
-    oversizedReducibleChars,
-    aggregateReducibleChars,
+    oversizedCount: plan.oversizedReplacementCount,
+    oversizedReducibleChars: plan.oversizedReducibleChars,
+    aggregateReducibleChars: plan.aggregateReducibleChars,
    maxReducibleChars,
  };
 }
@@ -446,84 +532,37 @@ function truncateOversizedToolResultsInExistingSessionManager(params: {
  const { sessionManager, contextWindowTokens } = params;
  const maxChars = calculateMaxToolResultChars(contextWindowTokens);
  const aggregateBudgetChars = calculateRecoveryAggregateToolResultChars(contextWindowTokens);
-  const branch = sessionManager.getBranch();
+  const branch = sessionManager.getBranch() as ToolResultBranchEntry[];

  if (branch.length === 0) {
    return { truncated: false, truncatedCount: 0, reason: "empty session" };
  }

-  const oversizedIndices: number[] = [];
-  for (let i = 0; i < branch.length; i += 1) {
-    const entry = branch[i];
-    if (entry.type !== "message") {
-      continue;
-    }
-    const msg = entry.message;
-    if ((msg as { role?: string }).role !== "toolResult") {
-      continue;
-    }
-    if (getToolResultTextLength(msg) > maxChars) {
-      oversizedIndices.push(i);
-    }
-  }
-
-  const oversizedReplacements = oversizedIndices.flatMap((index) => {
-    const entry = branch[index];
-    if (!entry || entry.type !== "message") {
-      return [];
-    }
-    return [
-      {
-        entryId: entry.id,
-        message: truncateToolResultMessage(entry.message, maxChars, {
-          minKeepChars: RECOVERY_MIN_KEEP_CHARS,
-        }),
-      },
-    ];
-  });
-
-  const oversizedTrimmedBranch = applyToolResultReplacementsToBranch({
-    branch: branch as Array<{ id: string; type: string; message?: AgentMessage }>,
-    replacements: oversizedReplacements,
-  });
-  const aggregateReplacements = buildAggregateToolResultReplacements({
-    branch: oversizedTrimmedBranch,
+  const plan = buildToolResultReplacementPlan({
+    branch,
+    maxChars,
    aggregateBudgetChars,
    minKeepChars: RECOVERY_MIN_KEEP_CHARS,
  });
-  const replacements = [...oversizedReplacements];
-  const replacementIndexById = new Map(replacements.map((replacement) => [replacement.entryId, replacement]));
-  for (const replacement of aggregateReplacements) {
-    replacementIndexById.set(replacement.entryId, replacement);
-  }
-  const mergedReplacements = Array.from(replacementIndexById.values());
-
-  if (mergedReplacements.length === 0) {
+  if (plan.replacements.length === 0) {
    return {
      truncated: false,
      truncatedCount: 0,
      reason: "no oversized or aggregate tool results",
    };
  }
-
  const rewriteResult = rewriteTranscriptEntriesInSessionManager({
    sessionManager,
-    replacements: mergedReplacements,
+    replacements: plan.replacements,
  });
  if (rewriteResult.changed && params.sessionFile) {
    emitSessionTranscriptUpdate(params.sessionFile);
  }

-  const truncatedKinds = [
-    oversizedReplacements.length > 0 ? "oversized" : "",
-    aggregateReplacements.length > 0 ? "aggregate" : "",
-  ]
-    .filter(Boolean)
-    .join("+");
-
  log.info(
-    `[tool-result-truncation] Truncated ${rewriteResult.rewrittenEntries} ${truncatedKinds || "tool"} tool result(s) in session ` +
-      `(contextWindow=${contextWindowTokens} maxChars=${maxChars} aggregateBudgetChars=${aggregateBudgetChars}) ` +
+    `[tool-result-truncation] Truncated ${rewriteResult.rewrittenEntries} tool result(s) in session ` +
+      `(contextWindow=${contextWindowTokens} maxChars=${maxChars} aggregateBudgetChars=${aggregateBudgetChars} ` +
+      `oversized=${plan.oversizedReplacementCount} aggregate=${plan.aggregateReplacementCount}) ` +
      `sessionKey=${params.sessionKey ?? params.sessionId ?? "unknown"}`,
  );