context-compactor

// Checks the rules that would ruin the plugin silently rather than loudly:
//  1. a cut never orphans a tool result from the request that produced it;
//  2. the verbatim tail is bounded in tokens, so one dump cannot overflow the window;
//  3. the rebuilt history always contains a user turn, which chat templates require;
//  4. chunk boundaries stay stable as the conversation grows — what the cache rests on.
// Stubs stand in for ChatMessage where only roles, tool requests and sizes matter; real Chat
// objects are used where the SDK's own validation is part of the test. Run: node test-split.js
const { Chat } = require("@lmstudio/sdk");
const {
  buildChat,
  countLeadingSystemMessages,
  findSplitIndex,
  hashOf,
  isSafeCut,
} = require("./dist/compaction");

let failures = 0;
function check(name, actual, expected) {
  const ok = actual === expected;
  if (!ok) failures++;
  console.log(`${ok ? "PASS" : "FAIL"}  ${name}` + (ok ? "" : `\n      got ${actual}, expected ${expected}`));
}

// "assistant*" marks an assistant message awaiting tool results. `size` is the token weight the
// fake counter reports, so a test can model a tool result that dwarfs a whole conversation.
const messages = (roles, sizes = {}) =>
  roles.map((role, i) => ({
    getRole: () => (role === "assistant*" ? "assistant" : role),
    getToolCallRequests: () => (role === "assistant*" ? [{ name: "some_tool" }] : []),
    getText: () => `${role} message ${i}`,
    getToolCallResults: () => [],
    size: sizes[i] ?? 10,
  }));

const counter = async m => m.size;
const splitOf = (msgs, systemCount, budget) => findSplitIndex(msgs, systemCount, budget, counter);

async function main() {
  //                0         1       2            3       4             5       6            7       8             9       10
  const toolTail = messages([
    "system", "user", "assistant", "user", "assistant*", "tool", "assistant", "user", "assistant*", "tool", "assistant",
  ]);

  console.log("-- cut safety --");
  check("system messages are counted", countLeadingSystemMessages(toolTail), 1);
  check("refuses to open the window on a tool result", isSafeCut(toolTail, 5), false);
  check("allows a cut that keeps a tool pair together", isSafeCut(toolTail, 4), true);
  check("accepts a completed turn", isSafeCut(toolTail, 3), true);
  check("the chosen cut is always safe", isSafeCut(toolTail, await splitOf(toolTail, 1, 1000)), true);
  check("refuses to cut when nothing can be compacted", await splitOf(messages(["system", "user", "assistant"]), 1, 0), -1);
  check("prefers a user turn when the conversation offers one", toolTail[await splitOf(toolTail, 1, 1000)].getRole(), "user");

  const manySystem = messages(["system", "system", "user", "assistant", "user", "assistant"]);
  const preambleSplit = await splitOf(manySystem, 2, 1000);
  check("never cuts into the system preamble", preambleSplit > 2 || preambleSplit === -1, true);

  // An agentic run — one question, then a long tool chain — fills the context fastest, so it is the
  // one that most needs compacting. Its only user turn sits at the top where cutting compacts
  // nothing, so the cut falls mid-chain and buildChat supplies the turn the template needs.
  const agentic = messages(["system", "user", "assistant*", "tool", "assistant*", "tool", "assistant*", "tool", "assistant"]);
  const agenticSplit = await splitOf(agentic, 1, 40);
  check("compacts a single-turn tool chain rather than let it overflow", agenticSplit > 1, true);
  check("and cuts it at a safe point", isSafeCut(agentic, agenticSplit), true);

  console.log("\n-- the tail is bounded in tokens, not messages --");

  // The failure that shipped: a tail of eight "recent messages" holding a disassembly dump. Counted
  // as messages it looked small; measured in tokens it was larger than the whole window, and the
  // compaction produced a prompt that still overflowed.
  const dump = messages(
    ["system", "user", "assistant", "user", "assistant", "tool", "assistant", "user", "assistant"],
    { 5: 60000 },
  );
  const dumpSplit = await splitOf(dump, 1, 6000);
  const tailTokens = dump.slice(dumpSplit).reduce((n, m) => n + m.size, 0);
  check("cuts past a huge tool result instead of preserving it", dumpSplit > 5, true);
  check("and the tail it keeps fits the budget", tailTokens <= 6000, true);
  check("refuses when not even the last message fits", await splitOf(dump, 1, 1), -1);

  // The earliest fitting cut wins: that keeps the most verbatim context.
  const plainRun = messages(["system", "user", "assistant", "user", "assistant", "user", "assistant"]);
  check("keeps as much verbatim as the budget allows", await splitOf(plainRun, 1, 1000), 3);

  console.log("\n-- the rebuilt history always satisfies the template --");

  // The rule Qwen's template enforces with raise_exception('No user query found in messages.').
  const hasUserTurn = chat => chat.getMessagesArray().some(m => m.getRole() === "user");
  for (const [name, roles, split] of [
    ["agentic tail, no user turn left", ["system", "user", "assistant", "assistant", "assistant", "assistant"], 3],
    ["tail that already has a user turn", ["system", "user", "assistant", "user", "assistant"], 3],
    ["tail of a single assistant turn", ["system", "user", "assistant", "assistant"], 3],
  ]) {
    const history = Chat.from(roles.map((r, i) => ({ role: r, content: `${r} ${i}` })));
    const rebuilt = buildChat(history, history.getMessagesArray(), 1, {
      splitIndex: split,
      summary: "state",
      memoryNote: "note",
    });
    check(`a user turn survives compaction (${name})`, hasUserTurn(rebuilt), true);
  }

  const plain = Chat.from([{ role: "system", content: "s" }, { role: "user", content: "u" }]);
  check(
    "returns the history unchanged when there is nothing to inject",
    buildChat(plain, plain.getMessagesArray(), 1, {}) === plain,
    true,
  );

  console.log("\n-- tool results are capped before they reach the model --");

  const { capResult, capTools } = require("./dist/toolCap");

  // 6 MB of binary from filesystem's read_file is what killed a real conversation: a quarter of a
  // million tokens arriving mid-reply, long after compaction has had its say.
  const huge = "x".repeat(6_000_000);
  const capped = capResult(huge, 4000);
  check("truncates a result larger than the whole context", capped.length < 20000, true);
  check("and says so, so the model does not simply retry", capped.includes("Truncated by context-compactor"), true);
  check("leaves a small result untouched", capResult("small", 4000), "small");
  check("passes through null", capResult(null, 4000), null);
  check("leaves a small object untouched", capResult({ a: 1 }, 4000).a, 1);

  // The shape belongs to MCP, not to us: a tool result is validated against
  // { content: [{ type: "text", text }] } on the way back. Returning a truncated JSON rendering
  // instead of the object gets the call rejected — the tool having worked perfectly.
  const mcp = capResult({ content: [{ type: "text", text: huge }], isError: false }, 4000);
  check("keeps the MCP result shape", Array.isArray(mcp.content) && mcp.content[0].type === "text", true);
  check("keeps sibling fields", mcp.isError, false);
  check("and the text field stays a string", typeof mcp.content[0].text, "string");
  check("but a truncated one", mcp.content[0].text.length < 20000, true);
  check("preserves non-string leaves", capResult({ n: 42, b: true, z: null }, 4000).n, 42);

  // Many medium strings must be bounded together, not each on its own.
  const many = capResult({ content: Array.from({ length: 20 }, () => ({ type: "text", text: "z".repeat(5000) })) }, 1000);
  const totalChars = many.content.reduce((n, p) => n + p.text.length, 0);
  check("shares one budget across the whole structure", totalChars < 4000 + 20 * 250, true);

  // The wrapper must be transparent: same name, same schema, only the result bounded.
  const mkTools = (impl, opts) => capTools([{ name: "read_file", description: "d", implementation: impl }], opts);
  const [wrapped] = mkTools(async () => huge, { maxPerResult: 4000, budget: 100000 });
  check("the wrapped tool keeps its name", wrapped.name, "read_file");
  check("and its result comes back bounded", (await wrapped.implementation({})).length < 20000, true);

  // The failure the per-result cap could not catch: fifteen obedient 4k results in one turn.
  const [rationed] = mkTools(async () => "y".repeat(40000), { maxPerResult: 4000, budget: 10000 });
  let total = 0;
  let starved = 0;
  for (let i = 0; i < 15; i++) {
    const r = await rationed.implementation({});
    if (r.includes("no context left")) starved++;
    total += Math.ceil(r.length / 4);
  }
  check("stops spending once the turn's budget is gone", starved > 0, true);
  check("and the whole turn stays inside the budget", total <= 10000 + 15 * 100, true);

  // Starved calls must still run the tool and still return its shape: a write must write.
  let ran = false;
  const [broke] = mkTools(
    async () => {
      ran = true;
      return { content: [{ type: "text", text: huge }] };
    },
    { maxPerResult: 4000, budget: 0 },
  );
  const starvedResult = await broke.implementation({});
  check("runs the tool even with no budget left", ran, true);
  check("and still returns a valid MCP shape", starvedResult.content[0].type, "text");
  check("carrying only the notice", starvedResult.content[0].text.includes("no context left"), true);

  console.log("\n-- chunk boundary stability (what the cache rests on) --");

  // Mirrors chunkByBudget's greedy walk with a deterministic cost, to assert the property the cache
  // depends on: growing the history must not disturb the chunks already sealed.
  function chunkGreedy(items, budget, cost) {
    const chunks = [];
    let current = [];
    let total = 0;
    for (const item of items) {
      const c = cost(item);
      if (current.length > 0 && total + c > budget) {
        chunks.push(current);
        current = [];
        total = 0;
      }
      current.push(item);
      total += c;
    }
    if (current.length > 0) chunks.push(current);
    return chunks;
  }

  const cost = m => m.length;
  const early = ["aaaa", "bbbb", "cc", "dddddd", "ee", "ffff"];
  const later = [...early, "gggg", "hh", "iiii", "jj"];
  const chunksEarly = chunkGreedy(early, 10, cost);
  const chunksLater = chunkGreedy(later, 10, cost);
  const sealedEarly = chunksEarly.slice(0, -1).map(c => hashOf(c.join("|")));
  const sealedLater = chunksLater.slice(0, sealedEarly.length).map(c => hashOf(c.join("|")));

  check("sealed chunks keep identical hashes as the conversation grows", JSON.stringify(sealedEarly), JSON.stringify(sealedLater));
  check("growing the history adds chunks rather than reshuffling them", chunksLater.length >= chunksEarly.length, true);
  check("hashing is content-addressed, not order-of-call dependent", hashOf("abc"), hashOf("abc"));
  check("different content yields a different key", hashOf("abc") === hashOf("abd"), false);

  console.log(failures === 0 ? "\nAll checks passed." : `\n${failures} check(s) failed.`);
  process.exit(failures === 0 ? 0 : 1);
}

main();