Files
dotfiles/.pi/agent/extensions/context-manager/src/prune.test.ts
2026-04-09 23:14:57 +01:00

132 lines
4.9 KiB
TypeScript

import test from "node:test";
import assert from "node:assert/strict";
import { resolvePolicy } from "./config.ts";
import { pruneContextMessages } from "./prune.ts";
const bulky = "line\n".repeat(300);
const boundaryBulky = "boundary\n".repeat(300);
const thresholdWithTrailingNewline = "threshold\n".repeat(150);
function buildPolicy(recentUserTurns = 4) {
return {
...resolvePolicy({ mode: "balanced", contextWindow: 200_000 }),
recentUserTurns,
};
}
test("pruneContextMessages replaces old bulky tool results with distilled summaries instead of deleting them", () => {
const policy = buildPolicy(2);
const bulkyFailure = [
"Build failed while compiling focus parser",
"Error: missing export createFocusMatcher from ./summary-focus.ts",
...Array.from({ length: 220 }, () => "stack frame"),
].join("\n");
const messages = [
{ role: "user", content: "turn 1" },
{ role: "toolResult", toolName: "bash", content: bulkyFailure },
{ role: "assistant", content: "observed turn 1" },
{ role: "user", content: "turn 2" },
{ role: "assistant", content: "observed turn 2" },
{ role: "user", content: "turn 3" },
];
const pruned = pruneContextMessages(messages, policy);
const distilled = pruned.find((message) => message.role === "toolResult");
assert.ok(distilled);
assert.match(distilled!.content, /missing export createFocusMatcher/);
assert.doesNotMatch(distilled!.content, /stack frame\nstack frame\nstack frame/);
});
test("aggressive mode distills an older bulky tool result sooner than conservative mode", () => {
const conservative = resolvePolicy({ mode: "conservative", contextWindow: 200_000 });
const aggressive = resolvePolicy({ mode: "aggressive", contextWindow: 200_000 });
const messages = [
{ role: "user", content: "turn 1" },
{ role: "toolResult", toolName: "read", content: bulky },
{ role: "assistant", content: "after turn 1" },
{ role: "user", content: "turn 2" },
{ role: "assistant", content: "after turn 2" },
{ role: "user", content: "turn 3" },
{ role: "assistant", content: "after turn 3" },
{ role: "user", content: "turn 4" },
];
const conservativePruned = pruneContextMessages(messages, conservative);
const aggressivePruned = pruneContextMessages(messages, aggressive);
assert.equal(conservativePruned[1]?.content, bulky);
assert.notEqual(aggressivePruned[1]?.content, bulky);
assert.match(aggressivePruned[1]?.content ?? "", /^\[distilled read output\]/);
});
test("pruneContextMessages keeps recent bulky tool results inside the recent-turn window", () => {
const policy = buildPolicy(2);
const messages = [
{ role: "user", content: "turn 1" },
{ role: "assistant", content: "observed turn 1" },
{ role: "user", content: "turn 2" },
{ role: "toolResult", toolName: "read", content: bulky },
{ role: "assistant", content: "observed turn 2" },
{ role: "user", content: "turn 3" },
];
const pruned = pruneContextMessages(messages, policy);
assert.deepEqual(pruned, messages);
});
test("pruneContextMessages keeps old non-bulky tool results outside the recent-turn window", () => {
const policy = buildPolicy(2);
const messages = [
{ role: "user", content: "turn 1" },
{ role: "toolResult", toolName: "read", content: "short output" },
{ role: "assistant", content: "observed turn 1" },
{ role: "user", content: "turn 2" },
{ role: "assistant", content: "observed turn 2" },
{ role: "user", content: "turn 3" },
];
const pruned = pruneContextMessages(messages, policy);
assert.deepEqual(pruned, messages);
});
test("pruneContextMessages keeps exactly-150-line tool results with a trailing newline", () => {
const policy = buildPolicy(2);
const messages = [
{ role: "user", content: "turn 1" },
{ role: "toolResult", toolName: "read", content: thresholdWithTrailingNewline },
{ role: "assistant", content: "after threshold output" },
{ role: "user", content: "turn 2" },
{ role: "assistant", content: "after turn 2" },
{ role: "user", content: "turn 3" },
];
const pruned = pruneContextMessages(messages, policy);
assert.deepEqual(pruned, messages);
});
test("pruneContextMessages honors the recent-user-turn boundary", () => {
const policy = buildPolicy(2);
const messages = [
{ role: "user", content: "turn 1" },
{ role: "toolResult", toolName: "read", content: bulky },
{ role: "assistant", content: "after turn 1" },
{ role: "user", content: "turn 2" },
{ role: "toolResult", toolName: "read", content: boundaryBulky },
{ role: "assistant", content: "after turn 2" },
{ role: "user", content: "turn 3" },
];
const pruned = pruneContextMessages(messages, policy);
assert.equal(pruned[1]?.role, "toolResult");
assert.match(pruned[1]?.content ?? "", /^\[distilled read output\]/);
assert.deepEqual(
pruned.map((message) => message.content),
["turn 1", pruned[1]!.content, "after turn 1", "turn 2", boundaryBulky, "after turn 2", "turn 3"]
);
});