my-pi/pi-compaction/test/extension-runtime.test.ts

import test from "node:test";
import assert from "node:assert/strict";
import { performance } from "node:perf_hooks";
import piCompaction from "../index.ts";

function harness() {
  const handlers = new Map<string, Function[]>();
  const tools = new Map<string, any>();
  const commands = new Map<string, any>();
  const entries: Array<{ customType: string; data: unknown }> = [];
  const toolDefinitions = [{ name: "bash", description: "shell", parameters: { type: "object" } }];
  const pi: any = {
    on(name: string, handler: Function) {
      const list = handlers.get(name) ?? [];
      list.push(handler);
      handlers.set(name, list);
    },
    appendEntry(customType: string, data: unknown) { entries.push({ customType, data }); },
    registerTool(tool: any) { tools.set(tool.name, tool); },
    registerCommand(name: string, command: any) { commands.set(name, command); },
    getActiveTools() { return ["bash"]; },
    getAllTools() { return toolDefinitions; },
  };
  return { pi, handlers, tools, commands, entries };
}

function context(branch: any[] = []) {
  return {
    sessionManager: {
      getSessionId: () => "session-test",
      getBranch: () => branch,
    },
    model: { provider: "test", api: "test", id: "model", contextWindow: 100_000 },
    modelRegistry: {},
    getSystemPrompt: () => "system",
    getContextUsage: () => ({ tokens: 1_000, contextWindow: 100_000, percent: 1 }),
    hasUI: false,
    abort() {},
  } as any;
}

test("one runtime projects tool evidence and recovers it through context_tree_query", async () => {
  const h = harness();
  piCompaction(h.pi);
  const ctx = context();
  await h.handlers.get("session_start")![0]!({}, ctx);

  const assistant = {
    role: "assistant",
    content: [{ type: "toolCall", id: "call-1", name: "bash", arguments: { command: "npm test" } }],
    timestamp: 10,
  };
  const result = {
    role: "toolResult",
    toolCallId: "call-1",
    toolName: "bash",
    content: [{ type: "text", text: "failure evidence" }],
    isError: true,
    timestamp: 11,
  };
  await h.handlers.get("turn_end")![0]!({ message: assistant, toolResults: [result] }, ctx);
  const projected = await h.handlers.get("context")![0]!({ messages: [assistant, result] }, ctx);

  assert.match(projected.messages[1].content[0].text, /pi-compaction tool=bash/);
  const cacheStart = performance.now();
  for (let iteration = 0; iteration < 1_000; iteration++) {
    await h.handlers.get("context")![0]!({ messages: [assistant, result] }, ctx);
  }
  const cacheAverageMs = (performance.now() - cacheStart) / 1_000;
  assert.ok(cacheAverageMs < 2, `full context handler cache-hit average ${cacheAverageMs.toFixed(3)}ms exceeded 2ms`);
  assert.equal(h.entries.length, 1);
  assert.equal(h.entries[0]!.customType, "pi-compaction");

  const query = h.tools.get("context_tree_query");
  const recovered = await query.execute("q", { toolCallIds: ["call-1"] }, new AbortController().signal, undefined, ctx);
  assert.match(recovered.content[0].text, /failure evidence/);

  await h.handlers.get("session_shutdown")![0]!({}, ctx);
});

test("model switching after a native checkpoint projects the portable checkpoint", async () => {
  const h = harness();
  piCompaction(h.pi);
  const details = {
    kind: "pi-compaction",
    version: 1,
    checkpoint: "native",
    modelKey: "openai-codex/openai-codex-responses/gpt-old",
    replacementHistory: [{ type: "compaction", encrypted_content: "opaque" }],
    portableSummary: "decisions and files from the old model",
    sourceFingerprint: "old",
  };
  const branch = [
    { type: "compaction", id: "c1", parentId: null, timestamp: new Date().toISOString(), summary: "marker", firstKeptEntryId: "m0", tokensBefore: 1_000, details },
    { type: "message", id: "m1", parentId: "c1", timestamp: new Date().toISOString(), message: { role: "user", content: "continue", timestamp: 20 } },
  ] as any[];
  const ctx = context(branch);
  await h.handlers.get("session_start")![0]!({}, ctx);
  const projected = await h.handlers.get("context")![0]!({ messages: [] }, ctx);
  assert.match(projected.messages[0].content[0].text, /decisions and files/);
  assert.equal(projected.messages[1].content, "continue");
  await h.handlers.get("session_shutdown")![0]!({}, ctx);
});