Compaction: sanitize token split accounting (#24058)

* Compaction: sanitize token split accounting

* Tests/Compaction: type sanitize token estimate callback
This commit is contained in:
Tak Hoffman
2026-02-22 20:13:21 -06:00
committed by GitHub
parent 259d863353
commit 50c5f75904
2 changed files with 58 additions and 2 deletions

View File

@@ -0,0 +1,52 @@
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import { describe, expect, it, vi } from "vitest";
const piCodingAgentMocks = vi.hoisted(() => ({
estimateTokens: vi.fn((_message: unknown) => 1),
generateSummary: vi.fn(async () => "summary"),
}));
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,
generateSummary: piCodingAgentMocks.generateSummary,
};
});
import { chunkMessagesByMaxTokens, splitMessagesByTokenShare } from "./compaction.js";
describe("compaction token accounting sanitization", () => {
it("does not pass toolResult.details into per-message token estimates", () => {
const messages: AgentMessage[] = [
{
role: "toolResult",
toolCallId: "call_1",
toolName: "browser",
isError: false,
content: [{ type: "text", text: "ok" }],
details: { raw: "x".repeat(50_000) },
timestamp: 1,
// oxlint-disable-next-line typescript/no-explicit-any
} as any,
{
role: "user",
content: "next",
timestamp: 2,
},
];
splitMessagesByTokenShare(messages, 2);
chunkMessagesByMaxTokens(messages, 16);
const calledWithDetails = piCodingAgentMocks.estimateTokens.mock.calls.some((call) => {
const message = call[0] as { details?: unknown } | undefined;
return Boolean(message?.details);
});
expect(calledWithDetails).toBe(false);
});
});

View File

@@ -23,6 +23,10 @@ export function estimateMessagesTokens(messages: AgentMessage[]): number {
return safe.reduce((sum, message) => sum + estimateTokens(message), 0);
}
function estimateCompactionMessageTokens(message: AgentMessage): number {
return estimateMessagesTokens([message]);
}
function normalizeParts(parts: number, messageCount: number): number {
if (!Number.isFinite(parts) || parts <= 1) {
return 1;
@@ -49,7 +53,7 @@ export function splitMessagesByTokenShare(
let currentTokens = 0;
for (const message of messages) {
const messageTokens = estimateTokens(message);
const messageTokens = estimateCompactionMessageTokens(message);
if (
chunks.length < normalizedParts - 1 &&
current.length > 0 &&
@@ -93,7 +97,7 @@ export function chunkMessagesByMaxTokens(
let currentTokens = 0;
for (const message of messages) {
const messageTokens = estimateTokens(message);
const messageTokens = estimateCompactionMessageTokens(message);
if (currentChunk.length > 0 && currentTokens + messageTokens > effectiveMax) {
chunks.push(currentChunk);
currentChunk = [];