spb/khaelor Public

1/**2 * KHAELOR3 * File: tests/context/compaction.test.ts4 * Description: Compaction tests — prune selection, pairing-safe cuts, prune-then-summarize pipeline ordering.5 *6 * Author: Simon-Pierre Boucher7 * Contact: contact@spboucher.ai8 */910import { describe, expect, it } from "vitest";11import {12  DEFAULT_COMPACTION_OPTIONS,13  PRUNE_PLACEHOLDER,14  collectRunningProcesses,15  compressEvents,16  parseCheckpoint,17  selectCompactionCut,18  selectPruneCandidates,19} from "../../src/context/index.js";20import type { CompactionOptions } from "../../src/context/index.js";21import { isPairingSafeCut } from "../../src/session/index.js";22import { FAKE_CHECKPOINT_TEXT, FakeModelClient, StreamBuilder, sampleSession } from "./fixtures.js";2324const opts = (overrides: Partial<CompactionOptions>): CompactionOptions => ({25  ...DEFAULT_COMPACTION_OPTIONS,26  ...overrides,27});2829describe("selectPruneCandidates", () => {30  it("prunes old tool results while protecting the newest N tokens of tool output", () => {31    const big = "x".repeat(2_000); // ~500 estimated tokens each32    const session = sampleSession({ bigToolOutput: big });33    // Make toolu_2's result big too, then add a third big recent result.34    const b = new StreamBuilder();35    for (const event of session.events) {36      if (event.type === "tool.completed" && event.payload.toolUseId === "toolu_2") {37        b.add({ type: "tool.completed", payload: { ...event.payload, modelText: big } });38      } else {39        b.add({ type: event.type, payload: event.payload } as never);40      }41    }42    b.add({43      type: "tool.requested",44      payload: { requestId: "req_4", blockIndex: 0, toolUseId: "toolu_3", toolName: "read", input: {} },45    });46    b.add({47      type: "tool.completed",48      payload: { toolUseId: "toolu_3", modelText: big, durationMs: 1, ui: { kind: "read", summary: "r" } },49    });5051    const decision = selectPruneCandidates(b.events, opts({ protectRecentToolTokens: 600 }));52    expect(decision.toolUseIds).toEqual(["toolu_1", "toolu_2"]);53    expect(decision.placeholder).toBe(PRUNE_PLACEHOLDER);54    expect(decision.tokensReclaimedEstimate).toBeGreaterThan(900);55  });5657  it("skips results already pruned and tiny results not worth a placeholder", () => {58    const big = "x".repeat(2_000);59    const session = sampleSession({ bigToolOutput: big });60    session.add({61      type: "context.pruned",62      payload: { toolUseIds: ["toolu_1"], placeholder: PRUNE_PLACEHOLDER, tokensReclaimedEstimate: 470 },63    });64    // toolu_2's result is tiny; toolu_1 already pruned → nothing left.65    const decision = selectPruneCandidates(session.events, opts({ protectRecentToolTokens: 0 }));66    expect(decision.toolUseIds).toEqual([]);67    expect(decision.tokensReclaimedEstimate).toBe(0);68  });69});7071describe("selectCompactionCut", () => {72  it("protects the first user message (head) and the latest user message (tail)", () => {73    const session = sampleSession();74    const cut = selectCompactionCut(session.events, opts({ protectTailTokens: 1 }));75    expect(cut).not.toBeNull();76    expect(cut as object).toEqual({ fromSeq: 4, toSeq: 14 });77    expect(isPairingSafeCut(session.events, cut as { fromSeq: number; toSeq: number })).toBe(true);78  });7980  it("never orphans a tool_use/tool_result pair — toSeq walks down to a safe boundary", () => {81    const b = sampleSession();82    // A tool round-trip whose result lands inside the protected tail.83    b.add({84      type: "tool.requested",85      payload: { requestId: "req_5", blockIndex: 0, toolUseId: "toolu_9", toolName: "bash", input: { cmd: "npm test" } },86    }); // seq 1787    b.add({88      type: "tool.completed",89      payload: {90        toolUseId: "toolu_9",91        modelText: "148 passed — full suite green, nothing else to report here",92        durationMs: 5,93        ui: { kind: "exec", summary: "npm test" },94      },95    }); // seq 1896    b.add({ type: "user.message-created", payload: { text: "Great, now document it", mentions: [] } }); // seq 199798    // Tail protection reaches back into seq 18 (the tool_result) but not seq 17:99    // a naive toSeq of 17 would orphan toolu_9, so the cut must retreat to 16.100    const cut = selectCompactionCut(b.events, opts({ protectTailTokens: 8 }));101    expect(cut).not.toBeNull();102    expect((cut as { toSeq: number }).toSeq).toBeLessThan(17);103    expect(isPairingSafeCut(b.events, cut as { fromSeq: number; toSeq: number })).toBe(true);104  });105106  it("returns null when there is no compactable middle", () => {107    const b = new StreamBuilder();108    b.add({ type: "user.message-created", payload: { text: "hello", mentions: [] } });109    expect(selectCompactionCut(b.events, opts({}))).toBeNull();110    expect(selectCompactionCut([], opts({}))).toBeNull();111  });112});113114describe("compressEvents pipeline", () => {115  it("prunes first: the summarizer sees recorded placeholders, never the reclaimed output", async () => {116    const big = "GIANT_TOOL_OUTPUT ".repeat(500);117    const session = sampleSession({ bigToolOutput: big });118    // Simulate the kernel: record the prune decision as a durable event first.119    const decision = selectPruneCandidates(session.events, opts({ protectRecentToolTokens: 10 }));120    expect(decision.toolUseIds).toContain("toolu_1");121    session.add({ type: "context.pruned", payload: decision });122123    const cut = selectCompactionCut(session.events, opts({ protectTailTokens: 1 }));124    expect(cut).not.toBeNull();125126    const client = new FakeModelClient(FAKE_CHECKPOINT_TEXT);127    const result = await compressEvents(128      { modelClient: client, model: "aux-model", maxOutputTokens: 2_000 },129      {130        events: session.events,131        cut: cut as { fromSeq: number; toSeq: number },132        trigger: "proactive-token-budget",133        tokensBefore: 150_000,134      },135    );136137    // Ordering: exactly one summarization call, made AFTER pruning was applied.138    expect(client.requests).toHaveLength(1);139    const sent = JSON.stringify(client.requests[0]);140    expect(sent).toContain(PRUNE_PLACEHOLDER);141    expect(sent).not.toContain("GIANT_TOOL_OUTPUT");142143    expect(result.summaryModel).toBe("aux-model");144    expect(result.trigger).toBe("proactive-token-budget");145    expect(result.tokensBefore).toBe(150_000);146    const parsed = parseCheckpoint(result.checkpointYaml);147    expect(parsed.ok).toBe(true);148    if (parsed.ok) expect(parsed.value.objective).toBe("Fix the bug in the session store");149  });150151  it("uses a small output budget and passes the AbortSignal through (cancellable)", async () => {152    const session = sampleSession();153    const cut = selectCompactionCut(session.events, opts({ protectTailTokens: 1 }));154    const client = new FakeModelClient(FAKE_CHECKPOINT_TEXT);155    const options = opts({ protectTailTokens: 1, summaryMaxOutputTokens: 777 });156    const controller = new AbortController();157    await compressEvents(158      { modelClient: client, model: "aux-model", maxOutputTokens: options.summaryMaxOutputTokens },159      {160        events: session.events,161        cut: cut as { fromSeq: number; toSeq: number },162        trigger: "user-command",163        tokensBefore: 0,164        signal: controller.signal,165      },166      options,167    );168    expect(client.requests[0]?.maxOutputTokens).toBe(777);169    expect(client.signals[0]).toBe(controller.signal);170171    const aborted = new AbortController();172    aborted.abort();173    await expect(174      compressEvents(175        { modelClient: new FakeModelClient(FAKE_CHECKPOINT_TEXT), model: "aux-model", maxOutputTokens: 100 },176        {177          events: session.events,178          cut: cut as { fromSeq: number; toSeq: number },179          trigger: "user-command",180          tokensBefore: 0,181          signal: aborted.signal,182        },183        options,184      ),185    ).rejects.toThrow();186  });187188  it("falls back to current_state when the model reply is unstructured, and derives running_processes from events", async () => {189    const session = sampleSession();190    session.add({191      type: "process.started",192      payload: { processId: "proc_9", pid: 1, command: "npm run dev", cwd: "/p", toolUseId: "t" },193    });194    expect(collectRunningProcesses(session.events)).toEqual([195      { id: "proc_9", command: "npm run dev", status: "running" },196    ]);197198    const cut = selectCompactionCut(session.events, opts({ protectTailTokens: 1 }));199    const client = new FakeModelClient("Sorry, here is a prose summary instead.");200    const result = await compressEvents(201      { modelClient: client, model: "aux-model", maxOutputTokens: 100 },202      {203        events: session.events,204        cut: cut as { fromSeq: number; toSeq: number },205        trigger: "reactive-overflow",206        tokensBefore: 10,207      },208    );209    const parsed = parseCheckpoint(result.checkpointYaml);210    expect(parsed.ok).toBe(true);211    if (parsed.ok) {212      expect(parsed.value.currentState).toBe("Sorry, here is a prose summary instead.");213      expect(parsed.value.runningProcesses).toEqual([214        { id: "proc_9", command: "npm run dev", status: "running" },215      ]);216    }217  });218219  it("refuses a pairing-unsafe cut", async () => {220    const session = sampleSession();221    await expect(222      compressEvents(223        { modelClient: new FakeModelClient(FAKE_CHECKPOINT_TEXT), model: "aux-model", maxOutputTokens: 100 },224        { events: session.events, cut: { fromSeq: 4, toSeq: 5 }, trigger: "user-command", tokensBefore: 0 },225      ),226    ).rejects.toThrow(/pairing/);227  });228});229