Compaction: sanitize token split accounting (#24058)
* Compaction: sanitize token split accounting * Tests/Compaction: type sanitize token estimate callback
This commit is contained in:
52
src/agents/compaction.token-sanitize.test.ts
Normal file
52
src/agents/compaction.token-sanitize.test.ts
Normal 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);
|
||||
});
|
||||
});
|
||||
@@ -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 = [];
|
||||
|
||||
Reference in New Issue
Block a user