test(agents): isolate compaction token estimator mocks

This commit is contained in:
Peter Steinberger
2026-04-17 16:47:48 +01:00
parent c0a16650d5
commit 97f713f459
2 changed files with 89 additions and 12 deletions

View File

@@ -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",

View File

@@ -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;