From 97f713f459774c1becf1ff20188d506faa2ab725 Mon Sep 17 00:00:00 2001
From: Peter Steinberger <steipete@gmail.com>
Date: Fri, 17 Apr 2026 16:47:48 +0100
Subject: [PATCH] test(agents): isolate compaction token estimator mocks

---
 src/agents/compaction.test.ts                 | 49 ++++++++++++++---
 .../run/preemptive-compaction.test.ts         | 52 ++++++++++++++++---
 2 files changed, 89 insertions(+), 12 deletions(-)

diff --git a/src/agents/compaction.test.ts b/src/agents/compaction.test.ts
index 94c09e240f4..0f5648b6109 100644
--- a/src/agents/compaction.test.ts
+++ b/src/agents/compaction.test.ts
@@ -1,13 +1,50 @@
 import type { AgentMessage } from "@mariozechner/pi-agent-core";
 import type { AssistantMessage, ToolResultMessage } from "@mariozechner/pi-ai";
-import { describe, expect, it } from "vitest";
-import {
-  estimateMessagesTokens,
-  pruneHistoryForContextShare,
-  splitMessagesByTokenShare,
-} from "./compaction.js";
+import { beforeAll, describe, expect, it, vi } from "vitest";
 import { makeAgentAssistantMessage } from "./test-helpers/agent-message-fixtures.js";
 
+const piCodingAgentMocks = vi.hoisted(() => ({
+  estimateTokens: vi.fn((message: unknown) => estimateTokenish(message)),
+}));
+
+function readText(value: unknown): string {
+  if (typeof value === "string") {
+    return value;
+  }
+  if (Array.isArray(value)) {
+    return value.map(readText).join("");
+  }
+  if (value && typeof value === "object") {
+    const record = value as { text?: unknown; content?: unknown; arguments?: unknown };
+    return `${readText(record.text)}${readText(record.content)}${readText(record.arguments)}`;
+  }
+  return "";
+}
+
+function estimateTokenish(message: unknown): number {
+  return Math.max(1, Math.ceil(readText(message).length / 4));
+}
+
+vi.mock("@mariozechner/pi-coding-agent", async () => {
+  const actual = await vi.importActual<typeof import("@mariozechner/pi-coding-agent")>(
+    "@mariozechner/pi-coding-agent",
+  );
+  return {
+    ...actual,
+    estimateTokens: piCodingAgentMocks.estimateTokens,
+  };
+});
+
+let estimateMessagesTokens: typeof import("./compaction.js").estimateMessagesTokens;
+let pruneHistoryForContextShare: typeof import("./compaction.js").pruneHistoryForContextShare;
+let splitMessagesByTokenShare: typeof import("./compaction.js").splitMessagesByTokenShare;
+
+beforeAll(async () => {
+  vi.resetModules();
+  ({ estimateMessagesTokens, pruneHistoryForContextShare, splitMessagesByTokenShare } =
+    await import("./compaction.js"));
+});
+
 function makeMessage(id: number, size: number): AgentMessage {
   return {
     role: "user",
diff --git a/src/agents/pi-embedded-runner/run/preemptive-compaction.test.ts b/src/agents/pi-embedded-runner/run/preemptive-compaction.test.ts
index cfd1aeb9797..5fcecfb97f7 100644
--- a/src/agents/pi-embedded-runner/run/preemptive-compaction.test.ts
+++ b/src/agents/pi-embedded-runner/run/preemptive-compaction.test.ts
@@ -1,11 +1,51 @@
 import type { AgentMessage } from "@mariozechner/pi-agent-core";
-import { describe, expect, it } from "vitest";
+import { beforeAll, describe, expect, it, vi } from "vitest";
 import { estimateToolResultReductionPotential } from "../tool-result-truncation.js";
-import {
-  PREEMPTIVE_OVERFLOW_ERROR_TEXT,
-  estimatePrePromptTokens,
-  shouldPreemptivelyCompactBeforePrompt,
-} from "./preemptive-compaction.js";
+
+const piCodingAgentMocks = vi.hoisted(() => ({
+  estimateTokens: vi.fn((message: unknown) => estimateTokenish(message)),
+}));
+
+function readText(value: unknown): string {
+  if (typeof value === "string") {
+    return value;
+  }
+  if (Array.isArray(value)) {
+    return value.map(readText).join("");
+  }
+  if (value && typeof value === "object") {
+    const record = value as { text?: unknown; content?: unknown };
+    return `${readText(record.text)}${readText(record.content)}`;
+  }
+  return "";
+}
+
+function estimateTokenish(message: unknown): number {
+  return Math.max(1, Math.ceil(readText(message).length / 4));
+}
+
+vi.mock("@mariozechner/pi-coding-agent", async () => {
+  const actual = await vi.importActual<typeof import("@mariozechner/pi-coding-agent")>(
+    "@mariozechner/pi-coding-agent",
+  );
+  return {
+    ...actual,
+    estimateTokens: piCodingAgentMocks.estimateTokens,
+  };
+});
+
+let PREEMPTIVE_OVERFLOW_ERROR_TEXT: typeof import("./preemptive-compaction.js").PREEMPTIVE_OVERFLOW_ERROR_TEXT;
+let estimatePrePromptTokens: typeof import("./preemptive-compaction.js").estimatePrePromptTokens;
+let shouldPreemptivelyCompactBeforePrompt: typeof import("./preemptive-compaction.js").shouldPreemptivelyCompactBeforePrompt;
+
+beforeAll(async () => {
+  vi.resetModules();
+  ({
+    PREEMPTIVE_OVERFLOW_ERROR_TEXT,
+    estimatePrePromptTokens,
+    shouldPreemptivelyCompactBeforePrompt,
+  } = await import("./preemptive-compaction.js"));
+});
 
 let timestamp = 1;