pi-context-manager/src/prune.test.ts

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 bulky tool results with distilled summaries inside older kept turns", () => {
  const policy = buildPolicy(3);
  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: "assistant", content: "observed turn 1" },
    { role: "user", content: "turn 2" },
    { role: "toolResult", toolName: "bash", content: bulkyFailure },
    { role: "assistant", content: "observed turn 2" },
    { role: "user", content: "turn 3" },
    { role: "assistant", content: "observed turn 3" },
    { role: "user", content: "turn 4" },
  ];

  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 drops older turns 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" },
    { role: "assistant", content: "after turn 4" },
    { role: "user", content: "turn 5" },
  ];

  const conservativePruned = pruneContextMessages(messages, conservative);
  const aggressivePruned = pruneContextMessages(messages, aggressive);

  assert.ok(conservativePruned.some((message) => message.role === "toolResult"));
  assert.ok(aggressivePruned.every((message) => message.role !== "toolResult"));
});

test("pruneContextMessages keeps newest-turn bulky tool results lossless", () => {
  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" },
  ];

  const pruned = pruneContextMessages(messages, policy);

  assert.deepEqual(pruned, messages);
});

test("pruneContextMessages drops 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.map((message) => message.content),
    ["turn 2", "observed turn 2", "turn 3"],
  );
});

test("pruneContextMessages keeps exactly-150-line tool results with a trailing newline inside a kept older turn", () => {
  const policy = buildPolicy(3);
  const messages = [
    { role: "user", content: "turn 1" },
    { role: "assistant", content: "after turn 1" },
    { role: "user", content: "turn 2" },
    { role: "toolResult", toolName: "read", content: thresholdWithTrailingNewline },
    { role: "assistant", content: "after threshold output" },
    { role: "user", content: "turn 3" },
    { role: "assistant", content: "after turn 3" },
    { role: "user", content: "turn 4" },
  ];

  const pruned = pruneContextMessages(messages, policy);

  assert.equal(pruned[1]?.content, thresholdWithTrailingNewline);
});

test("pruneContextMessages drops turns older than the kept suffix", () => {
  const policy = buildPolicy(2);
  const messages = [
    { role: "user", content: "turn 1" },
    { role: "assistant", content: "after turn 1" },
    { role: "user", content: "turn 2" },
    { role: "assistant", content: "after turn 2" },
    { role: "user", content: "turn 3" },
  ];

  const pruned = pruneContextMessages(messages, policy);

  assert.deepEqual(
    pruned.map((message) => message.content),
    ["turn 2", "after turn 2", "turn 3"],
  );
});

test("pruneContextMessages distills bulky tool results only inside older kept turns", () => {
  const policy = buildPolicy(2);
  const messages = [
    { role: "user", content: "turn 1" },
    { role: "assistant", content: "after turn 1" },
    { role: "user", content: "turn 2" },
    { role: "toolResult", toolName: "read", content: bulky },
    { role: "assistant", content: "after turn 2" },
    { role: "user", content: "turn 3" },
    { role: "toolResult", toolName: "read", content: boundaryBulky },
    { role: "assistant", content: "after turn 3" },
  ];

  const pruned = pruneContextMessages(messages, policy);

  assert.deepEqual(
    pruned.map((message) => message.role),
    ["user", "toolResult", "assistant", "user", "toolResult", "assistant"],
  );
  assert.match(pruned[1]?.content ?? "", /^\[distilled read output\]/);
  assert.equal(pruned[4]?.content, boundaryBulky);
});