test: share token estimate mock

This commit is contained in:
Peter Steinberger
2026-04-18 22:29:56 +01:00
parent f0f4fa6978
commit 7481478303
3 changed files with 37 additions and 64 deletions

View File

@@ -2,38 +2,7 @@ import type { AgentMessage } from "@mariozechner/pi-agent-core";
import type { AssistantMessage, ToolResultMessage } from "@mariozechner/pi-ai";
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,
};
});
import "./test-helpers/pi-coding-agent-token-mock.js";
let estimateMessagesTokens: typeof import("./compaction.js").estimateMessagesTokens;
let pruneHistoryForContextShare: typeof import("./compaction.js").pruneHistoryForContextShare;

View File

@@ -1,39 +1,8 @@
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import { beforeAll, describe, expect, it, vi } from "vitest";
import "../../test-helpers/pi-coding-agent-token-mock.js";
import { estimateToolResultReductionPotential } from "../tool-result-truncation.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;

View File

@@ -0,0 +1,35 @@
import { vi } from "vitest";
const piCodingAgentTokenMocks = vi.hoisted(() => {
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));
}
return {
estimateTokens: vi.fn((message: unknown) => estimateTokenish(message)),
};
});
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: piCodingAgentTokenMocks.estimateTokens,
};
});