my-pi/pi-compaction/test/masking.test.ts

import test from "node:test";
import assert from "node:assert/strict";
import { MASK_PREFIX, planMasking, rebuildMaskedLedger } from "../src/masking.ts";
import { effectiveThreshold } from "../src/token-accounting.ts";
import type { MaskingConfig } from "../src/types.ts";

const MASKING: MaskingConfig = {
  enabled: true,
  triggerTokens: 1_000,
  triggerRatio: 0.9,
  targetTokens: 400,
  targetRatio: 0.35,
  minClearTokens: 200,
  keepRecentTurns: 2,
  minEntryTokens: 10,
  excludeTools: ["ask_user_question"],
};

const estimate = (message: any): number => Math.ceil(JSON.stringify(message).length / 4);

function assistantMessage(id: string, calls: Array<{ id: string; name: string; args?: unknown }>) {
  return {
    role: "assistant",
    content: calls.map((call) => ({ type: "toolCall", id: call.id, name: call.name, arguments: call.args ?? {} })),
    timestamp: id,
  };
}

function toolResultMessage(id: string, toolCallId: string, text: string, toolName = "bash") {
  return {
    role: "toolResult",
    toolCallId,
    toolName,
    content: [{ type: "text", text }],
    timestamp: id,
  };
}

function entry(id: string, message: any, messages = [message]) {
  return { sourceEntry: { id, type: "message", message }, messages };
}

function protect(turns = MASKING.keepRecentTurns) {
  return Array.from({ length: turns }, (_unused, index) =>
    entry(`recent-${index}`, { role: "assistant", content: [{ type: "text", text: "ok" }], timestamp: `recent-${index}` }));
}

test("below the effective trigger plans no drafts", () => {
  const entries = [
    entry("e1", assistantMessage("a1", [{ id: "c1", name: "bash", args: { command: "ls" } }])),
    entry("e2", toolResultMessage("r1", "c1", "x".repeat(8_000))),
    ...protect(),
  ];
  const plan = planMasking({ contextEntries: entries, currentTokens: 999, contextWindow: 0, masking: MASKING, estimate });
  assert.equal(plan.drafts.length, 0);
  assert.equal(plan.masked.length, 0);
});

test("above the trigger masks oldest first until context fits under target", () => {
  const entries = [
    entry("e1", assistantMessage("a1", [
      { id: "c1", name: "bash", args: { command: "one" } },
      { id: "c2", name: "bash", args: { command: "two" } },
      { id: "c3", name: "bash", args: { command: "three" } },
    ])),
    entry("e2", toolResultMessage("r1", "c1", "a".repeat(4_000))),
    entry("e3", toolResultMessage("r2", "c2", "b".repeat(4_000))),
    entry("e4", toolResultMessage("r3", "c3", "c".repeat(4_000))),
    entry("e5", toolResultMessage("r4", "c4", "recent result".repeat(200))),
    ...protect(),
  ];
  // needed = 2_000 - 400 = 1_600; each result clears ~980 tokens.
  const plan = planMasking({ contextEntries: entries, currentTokens: 2_000, contextWindow: 0, masking: MASKING, estimate });
  assert.deepEqual(plan.drafts.map((draft) => draft.targetId), ["e2", "e3"]);
});

test("results inside the recent keepRecentTurns window are never masked", () => {
  const entries = [
    entry("e1", toolResultMessage("r1", "c1", "a".repeat(4_000))),
    entry("recent-a", assistantMessage("ra", [{ id: "rc", name: "bash", args: {} }])),
    entry("recent-r", toolResultMessage("rr", "rc", "b".repeat(4_000))),
    entry("recent-b", assistantMessage("rb", [])),
  ];
  const plan = planMasking({ contextEntries: entries, currentTokens: 3_500, contextWindow: 0, masking: MASKING, estimate });
  assert.deepEqual(plan.drafts.map((draft) => draft.targetId), ["e1"]);
});

test("excluded tools, SKILL.md reads, already-masked and small results are skipped", () => {
  const entries = [
    entry("e1", assistantMessage("a1", [
      { id: "c1", name: "ask_user_question", args: { question: "x" } },
      { id: "c2", name: "read", args: { path: "/repo/skills/foo/SKILL.md" } },
      { id: "c3", name: "bash", args: { command: "old" } },
      { id: "c4", name: "bash", args: { command: "tiny" } },
    ])),
    entry("e2", toolResultMessage("r1", "c1", "a".repeat(4_000), "ask_user_question")),
    entry("e3", toolResultMessage("r2", "c2", "b".repeat(4_000), "read")),
    entry("e4", toolResultMessage("r3", "c3", `${MASK_PREFIX}bash result cleared to save context (~900 tokens).`), ),
    entry("e5", toolResultMessage("r4", "c4", "tiny")),
    entry("e6", toolResultMessage("r5", "c5", "d".repeat(4_000))),
    ...protect(),
  ];
  const masking = { ...MASKING, minClearTokens: 100, minEntryTokens: 100 };
  const plan = planMasking({ contextEntries: entries, currentTokens: 3_500, contextWindow: 0, masking, estimate });
  assert.deepEqual(plan.drafts.map((draft) => draft.targetId), ["e6"]);
  assert.equal(plan.masked[0]!.toolCallId, "c5");
});

test("when clearable volume stays under minClear nothing is masked", () => {
  const entries = [
    entry("e1", toolResultMessage("r1", "c1", "a".repeat(600))),
    ...protect(),
  ];
  const plan = planMasking({ contextEntries: entries, currentTokens: 3_500, contextWindow: 0, masking: MASKING, estimate });
  // ~130 tokens clearable, below minClearTokens (200) → not worth invalidating cache.
  assert.equal(plan.drafts.length, 0);
});

test("placeholder is deterministic, names the tool and call id, and truncates args at 300 chars", () => {
  const longArgs = { command: `run ${"v".repeat(600)}` };
  const entries = [
    entry("e1", assistantMessage("a1", [{ id: "c1", name: "bash", args: longArgs }])),
    entry("e2", toolResultMessage("r1", "c1", "a".repeat(4_000))),
    ...protect(),
  ];
  const masking = { ...MASKING, minClearTokens: 1 };
  const first = planMasking({ contextEntries: entries, currentTokens: 3_500, contextWindow: 0, masking, estimate });
  const second = planMasking({ contextEntries: entries, currentTokens: 3_500, contextWindow: 0, masking, estimate });
  assert.equal(first.drafts.length, 1);
  const placeholder = first.masked[0]!.placeholder;
  assert.equal(placeholder, second.masked[0]!.placeholder);
  assert.ok(placeholder.startsWith(MASK_PREFIX));
  assert.match(placeholder, /bash result cleared/);
  assert.match(placeholder, /"toolCallIds":\["c1"\]/);
  const argsText = placeholder.match(/Call args: (.*)\. Re-run/)?.[1] ?? "";
  assert.ok(argsText.length <= 301, `args excerpt ${argsText.length} exceeded 300+ellipsis`);
  assert.ok(placeholder.includes("…"));
});

test("a masked result carrying details.overflow.id gets an overflow_query hint in the placeholder", () => {
  const withOverflow = {
    ...toolResultMessage("r1", "c1", "a".repeat(4_000)),
    details: { overflow: { id: "0123456789ab", totalLines: 5, totalBytes: 4000, truncated: true } },
  };
  const withoutOverflow = toolResultMessage("r2", "c2", "b".repeat(4_000));
  const entries = [
    entry("e1", assistantMessage("a1", [{ id: "c1", name: "bash" }, { id: "c2", name: "bash" }])),
    entry("e2", withOverflow),
    entry("e3", withoutOverflow),
    ...protect(),
  ];
  const masking = { ...MASKING, minClearTokens: 1 };
  const plan = planMasking({ contextEntries: entries, currentTokens: 3_500, contextWindow: 0, masking, estimate });
  assert.equal(plan.masked.length, 2);
  const first = plan.masked.find((m) => m.toolCallId === "c1")!.placeholder;
  const second = plan.masked.find((m) => m.toolCallId === "c2")!.placeholder;
  assert.ok(first.includes(` Full original stored: overflow_query {"id":"0123456789ab"}.`));
  assert.ok(!second.includes("overflow_query"));
});

test("the small-window ratio caps both trigger and target", () => {
  const masking: MaskingConfig = {
    ...MASKING,
    triggerTokens: 200_000,
    triggerRatio: 0.6,
    targetTokens: 100_000,
    targetRatio: 0.35,
    minClearTokens: 50_000,
  };
  assert.equal(effectiveThreshold(masking.triggerTokens, masking.triggerRatio, 272_000), 163_200);
  assert.equal(effectiveThreshold(masking.targetTokens, masking.targetRatio, 272_000), 95_200);
  const entries = [
    entry("e1", toolResultMessage("r1", "c1", "a".repeat(300_000))),
    ...protect(),
  ];
  const below = planMasking({ contextEntries: entries, currentTokens: 163_199, contextWindow: 272_000, masking, estimate });
  assert.equal(below.drafts.length, 0);
  const above = planMasking({ contextEntries: entries, currentTokens: 163_200, contextWindow: 272_000, masking, estimate });
  assert.equal(above.drafts.length, 1);
});

test("ledger rebuild counts only masked targets still inside the active context", () => {
  const placeholder = (id: string) =>
    `${MASK_PREFIX}bash result cleared to save context (~100 tokens). Call args: {}. Re-run the call if still needed, or recover the original with context_tree_query {"toolCallIds":["${id}"]}.`;
  const branch = [
    { type: "message", id: "u1", parentId: null, timestamp: "t", message: { role: "user", content: "start", timestamp: 1 } },
    { type: "message", id: "r-old", parentId: "u1", timestamp: "t", message: toolResultMessage("m-old", "c-old", "o".repeat(2_000)) },
    { type: "context_edit", id: "e-old", parentId: "r-old", timestamp: "t", targetId: "r-old", replacement: { content: placeholder("c-old") } },
    { type: "message", id: "m-keep", parentId: "e-old", timestamp: "t", message: { role: "user", content: "kept question", timestamp: 2 } },
    { type: "compaction", id: "c1", parentId: "m-keep", timestamp: "t", summary: "summary", firstKeptEntryId: "m-keep", tokensBefore: 9_999 },
    { type: "message", id: "a-keep", parentId: "c1", timestamp: "t", message: assistantMessage("a-keep", []) },
    { type: "message", id: "r-new", parentId: "a-keep", timestamp: "t", message: toolResultMessage("m-new", "c-new", "n".repeat(2_000)) },
    { type: "context_edit", id: "e-new", parentId: "r-new", timestamp: "t", targetId: "r-new", replacement: { content: placeholder("c-new") } },
  ];
  const ledger = rebuildMaskedLedger(branch, "e-new", estimate);
  assert.equal(ledger.size, 1);
  assert.ok(ledger.has("c-new"));
  assert.ok(!ledger.has("c-old"));
});

test("user and assistant messages are never masked", () => {
  const entries = [
    entry("e1", { role: "user", content: "u".repeat(8_000), timestamp: "u1" }),
    entry("e2", { role: "assistant", content: [{ type: "text", text: "a".repeat(8_000) }], timestamp: "a1" }),
    ...protect(),
  ];
  const plan = planMasking({ contextEntries: entries, currentTokens: 9_000, contextWindow: 0, masking: { ...MASKING, keepRecentTurns: 1, minClearTokens: 1 }, estimate });
  assert.equal(plan.drafts.length, 0);
});