feat(compaction): add staged helpers

This commit is contained in:
Peter Steinberger
2026-01-23 22:23:00 +00:00
parent 4ee70be690
commit ae0741a346
2 changed files with 411 additions and 0 deletions

View File

@@ -0,0 +1,70 @@
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import { describe, expect, it } from "vitest";
import {
estimateMessagesTokens,
pruneHistoryForContextShare,
splitMessagesByTokenShare,
} from "./compaction.js";
function makeMessage(id: number, size: number): AgentMessage {
return {
role: "user",
content: "x".repeat(size),
timestamp: id,
};
}
describe("splitMessagesByTokenShare", () => {
it("splits messages into two non-empty parts", () => {
const messages: AgentMessage[] = [
makeMessage(1, 4000),
makeMessage(2, 4000),
makeMessage(3, 4000),
makeMessage(4, 4000),
];
const parts = splitMessagesByTokenShare(messages, 2);
expect(parts.length).toBeGreaterThanOrEqual(2);
expect(parts[0]?.length).toBeGreaterThan(0);
expect(parts[1]?.length).toBeGreaterThan(0);
expect(parts.flat().length).toBe(messages.length);
});
});
describe("pruneHistoryForContextShare", () => {
it("drops older chunks until the history budget is met", () => {
const messages: AgentMessage[] = [
makeMessage(1, 4000),
makeMessage(2, 4000),
makeMessage(3, 4000),
makeMessage(4, 4000),
];
const maxContextTokens = 2000; // budget is 1000 tokens (50%)
const pruned = pruneHistoryForContextShare({
messages,
maxContextTokens,
maxHistoryShare: 0.5,
parts: 2,
});
expect(pruned.droppedChunks).toBeGreaterThan(0);
expect(pruned.keptTokens).toBeLessThanOrEqual(Math.floor(maxContextTokens * 0.5));
expect(pruned.messages.length).toBeGreaterThan(0);
});
it("keeps history when already within budget", () => {
const messages: AgentMessage[] = [makeMessage(1, 1000)];
const maxContextTokens = 2000;
const pruned = pruneHistoryForContextShare({
messages,
maxContextTokens,
maxHistoryShare: 0.5,
parts: 2,
});
expect(pruned.droppedChunks).toBe(0);
expect(pruned.messages.length).toBe(messages.length);
expect(pruned.keptTokens).toBe(estimateMessagesTokens(messages));
});
});