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

import test from "node:test";
import assert from "node:assert/strict";
import piCompaction from "../index.ts";
import { getCompactionService } from "../src/service.ts";
import { MASK_PREFIX } from "../src/masking.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[] = [], tokens = 1_000) {
  const compactCalls: any[] = [];
  const sessionId = `session-${Math.random().toString(36).slice(2)}`;
  const ctx = {
    sessionManager: {
      getSessionId: () => sessionId,
      getBranch: () => branch,
      getLeafId: () => "leaf-1",
    },
    model: { provider: "test", api: "test", id: "model", contextWindow: 100_000 },
    modelRegistry: {},
    getSystemPrompt: () => "system",
    getContextUsage: () => ({ tokens, contextWindow: 100_000, percent: 1 }),
    compact(options: unknown) { compactCalls.push(options); },
    isIdle: () => true,
    hasUI: false,
    abort() {},
  } as any;
  return { ctx, compactCalls };
}

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

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,
  };
}

// Default config: masking trigger = min(200K, 60% of window); at the harness
// 100K window that is 60K tokens and the target is 35K, so 60K used needs
// ~25K tokens of maskable results outside the 10-turn protected window.
function pressurizedContextEntries() {
  const entries = [
    projectedEntry("a1", assistantMessage("a1", [{ id: "c1", name: "bash", args: { command: "npm test" } }])),
    projectedEntry("r1", toolResultMessage("r1", "c1", `failure evidence\n${"x".repeat(120_000)}\nstack tail`)),
  ];
  for (let index = 0; index < 10; index++) {
    entries.push(projectedEntry(`keep-${index}`, assistantMessage(`keep-${index}`, [])));
  }
  return entries;
}

test("turn_end masks old tool results under pressure and preserves prior drafts", async () => {
  const h = harness();
  piCompaction(h.pi);
  const { ctx } = context([], 60_000);
  await h.handlers.get("session_start")![0]!({}, ctx);

  const prior = { type: "custom", customType: "other-extension", data: {} };
  const result = await h.handlers.get("turn_end")![0]!({
    entries: [prior],
    message: assistantMessage("a1", []),
    toolResults: [],
    context: { contextEntries: pressurizedContextEntries(), contextMessages: [] },
  }, ctx);

  assert.ok(result?.entries);
  assert.strictEqual(result.entries[0], prior);
  const edits = result.entries.slice(1);
  assert.equal(edits.length, 1);
  assert.equal(edits[0]!.type, "context_edit");
  assert.equal(edits[0]!.targetId, "r1");
  assert.ok(edits[0]!.replacement.content.startsWith(MASK_PREFIX));
  // Masking is append-only context_edit; nothing is written as custom entries.
  assert.equal(h.entries.length, 0);

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

test("turn_end below the masking trigger returns undefined", async () => {
  const h = harness();
  piCompaction(h.pi);
  const { ctx } = context([], 40_000);
  await h.handlers.get("session_start")![0]!({}, ctx);

  const result = await h.handlers.get("turn_end")![0]!({
    entries: [],
    message: assistantMessage("a1", []),
    toolResults: [],
    context: { contextEntries: pressurizedContextEntries(), contextMessages: [] },
  }, ctx);
  assert.equal(result, undefined);
  await h.handlers.get("session_shutdown")![0]!({}, ctx);
});

test("context hook returns messages untouched on the non-checkpoint path", async () => {
  const h = harness();
  piCompaction(h.pi);
  const { ctx } = context();
  await h.handlers.get("session_start")![0]!({}, ctx);

  const messages = [assistantMessage("a1", []), toolResultMessage("r1", "c1", "visible\nfile-a.ts")];
  const projected = await h.handlers.get("context")![0]!({ messages }, ctx);
  assert.equal(projected, undefined);
  await h.handlers.get("session_shutdown")![0]!({}, ctx);
});

test("agent_settled compacts once above the summary threshold and never repeats on the same leaf", async () => {
  const h = harness();
  piCompaction(h.pi);
  const { ctx, compactCalls } = context([], 95_000);
  await h.handlers.get("session_start")![0]!({}, ctx);

  h.handlers.get("agent_settled")![0]!({ type: "agent_settled", aborted: false }, ctx);
  assert.equal(compactCalls.length, 1);
  assert.equal(typeof compactCalls[0].onError, "function");

  h.handlers.get("agent_settled")![0]!({ type: "agent_settled", aborted: false }, ctx);
  assert.equal(compactCalls.length, 1);

  ctx.sessionManager.getLeafId = () => "leaf-2";
  h.handlers.get("agent_settled")![0]!({ type: "agent_settled", aborted: false }, ctx);
  assert.equal(compactCalls.length, 2);
  await h.handlers.get("session_shutdown")![0]!({}, ctx);
});

test("agent_settled stays idle below the summary threshold", async () => {
  const h = harness();
  piCompaction(h.pi);
  const { ctx, compactCalls } = context([], 50_000);
  await h.handlers.get("session_start")![0]!({}, ctx);

  h.handlers.get("agent_settled")![0]!({ type: "agent_settled", aborted: false }, ctx);
  assert.equal(compactCalls.length, 0);
  await h.handlers.get("session_shutdown")![0]!({}, ctx);
});

test("context_tree_query recovers the original tool output from the branch", async () => {
  const h = harness();
  piCompaction(h.pi);
  const branch = [
    { type: "message", id: "a1", parentId: null, timestamp: "t", message: assistantMessage("a1", [{ id: "c1", name: "bash", args: { command: "npm test" } }]) },
    { type: "message", id: "r1", parentId: "a1", timestamp: "t", message: toolResultMessage("r1", "c1", "failure evidence\nstack tail") },
  ];
  const { ctx } = context(branch);
  await h.handlers.get("session_start")![0]!({}, ctx);

  const query = h.tools.get("context_tree_query")!;
  const recovered = await query.execute("q", { toolCallIds: ["c1"] }, new AbortController().signal, undefined, ctx);
  assert.match(recovered.content[0].text, /failure evidence/);
  assert.match(recovered.content[0].text, /stack tail/);
  assert.deepEqual(recovered.details.missing, []);

  const missing = await query.execute("q", { toolCallIds: ["nope"] }, new AbortController().signal, undefined, ctx);
  assert.match(missing.content[0].text, /not found/);
  await h.handlers.get("session_shutdown")![0]!({}, ctx);
});

test("service query returns placeholders only for masked tool call ids", async () => {
  const h = harness();
  piCompaction(h.pi);
  const { ctx } = context([], 60_000);
  await h.handlers.get("session_start")![0]!({}, ctx);
  const sessionId = ctx.sessionManager.getSessionId();

  await h.handlers.get("turn_end")![0]!({
    entries: [],
    message: assistantMessage("a1", []),
    toolResults: [],
    context: { contextEntries: pressurizedContextEntries(), contextMessages: [] },
  }, ctx);

  const records = getCompactionService()!.query(sessionId, ["c1", "never-masked"]);
  assert.equal(records.length, 1);
  assert.equal(records[0]!.toolCallId, "c1");
  assert.ok(records[0]!.deterministicSummary.startsWith(MASK_PREFIX));
  assert.equal(typeof records[0]!.resultTimestamp, "number");
  await h.handlers.get("session_shutdown")![0]!({}, ctx);
});

test("hydrating a branch with this extension's context_edit rebuilds masking metrics", async () => {
  const h = harness();
  piCompaction(h.pi);
  const placeholder = `${MASK_PREFIX}bash result cleared to save context (~1_000 tokens). Call args: {}. Re-run the call if still needed, or recover the original with context_tree_query {"toolCallIds":["c1"]}.`;
  const branch = [
    { type: "message", id: "a1", parentId: null, timestamp: "t", message: assistantMessage("a1", [{ id: "c1", name: "bash", args: { command: "npm test" } }]) },
    { type: "message", id: "r1", parentId: "a1", timestamp: "t", message: toolResultMessage("r1", "c1", "x".repeat(8_000)) },
    { type: "context_edit", id: "e1", parentId: "r1", timestamp: "t", targetId: "r1", replacement: { content: placeholder } },
  ];
  const { ctx } = context(branch);
  await h.handlers.get("session_start")![0]!({}, ctx);
  const sessionId = ctx.sessionManager.getSessionId();

  await h.handlers.get("context")![0]!({ messages: [toolResultMessage("r1", "c1", "x".repeat(8_000))] }, ctx);
  const metrics = getCompactionService()!.getMetrics(sessionId);
  assert.equal(metrics!.maskedResults, 1);
  assert.ok(metrics!.reclaimedTokens > 0);
  assert.ok(metrics!.rawTokens > metrics!.effectiveTokens);
  await h.handlers.get("session_shutdown")![0]!({}, ctx);
});

test("session_before_compact no longer cancels threshold compactions", async () => {
  const h = harness();
  piCompaction(h.pi);
  const { ctx } = context();
  await h.handlers.get("session_start")![0]!({}, ctx);
  ctx.modelRegistry.getApiKeyAndHeaders = async () => ({ ok: false, error: "no credentials" });

  const result = await h.handlers.get("session_before_compact")![0]!({
    reason: "threshold", branchEntries: [], preparation: { firstKeptEntryId: "u1" },
    signal: new AbortController().signal,
  }, ctx);
  assert.equal(result, undefined);
  await h.handlers.get("session_shutdown")![0]!({}, ctx);
});

test("native checkpoint projection applies context edits made after the checkpoint", async () => {
  const h = harness();
  piCompaction(h.pi);
  const details = {
    kind: "pi-compaction",
    version: 1,
    checkpoint: "native",
    modelKey: "test:test:model",
    replacementHistory: [{ type: "compaction", encrypted_content: "opaque" }],
    portableSummary: "checkpoint summary",
    sourceFingerprint: "fp",
  };
  const branch = [
    { type: "compaction", id: "c1", parentId: null, timestamp: "t", summary: "marker", firstKeptEntryId: "r1", tokensBefore: 1_000, details },
    { type: "message", id: "r1", parentId: "c1", timestamp: "t", message: toolResultMessage("r1", "c1", "ORIGINAL-CONTENT-HERE") },
    { type: "context_edit", id: "e1", parentId: "r1", timestamp: "t", targetId: "r1", replacement: { content: `${MASK_PREFIX}bash result cleared to save context (~5 tokens). Call args: {}.` } },
    { type: "message", id: "u1", parentId: "e1", timestamp: "t", message: { role: "user", content: "next question", timestamp: 30 } },
  ] as any[];
  const { ctx } = context(branch);
  await h.handlers.get("session_start")![0]!({}, ctx);

  const projected = await h.handlers.get("context")![0]!({ messages: [] }, ctx);
  const serialized = JSON.stringify(projected.messages);
  assert.doesNotMatch(serialized, /ORIGINAL-CONTENT-HERE/);
  assert.match(serialized, /result cleared to save context/);
  assert.match(serialized, /next question/);
  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);
});

for (const provider of ["sub-api", "openai-codex"]) {
  test(`${provider} native compaction uses the selected model and restores its checkpoint`, async (t) => {
    const h = harness();
    piCompaction(h.pi);
    const branch: any[] = [{
      type: "message", id: "u1", parentId: null,
      message: { role: "user", content: "Remember the test decision.", timestamp: 1 },
    }];
    const { ctx } = context(branch);
    ctx.model = {
      provider, api: provider === "sub-api" ? "openai-responses" : "openai-codex-responses",
      id: "selected-model", input: ["text"], contextWindow: 100_000,
      baseUrl: provider === "sub-api" ? "https://gateway.example/v1" : "https://codex.example/codex",
    };
    const jwt = Buffer.from(JSON.stringify({
      "https://api.openai.com/auth": { chatgpt_account_id: "test-account" },
    })).toString("base64url");
    const apiKey = provider === "sub-api" ? "ordinary-test-key" : `header.${jwt}.signature`;
    ctx.modelRegistry.getApiKeyAndHeaders = async (model: unknown) => {
      assert.equal(model, ctx.model);
      return { ok: true, apiKey, headers: { "x-test-route": "preserved" } };
    };
    const requests: Array<{ url: string; headers: Headers; body: any }> = [];
    t.mock.method(globalThis, "fetch", async (url: string, init: RequestInit) => {
      requests.push({ url: String(url), headers: new Headers(init.headers), body: JSON.parse(String(init.body)) });
      const events = [
        { type: "response.output_item.done", item: { type: "compaction", encrypted_content: "opaque-checkpoint" } },
        { type: "response.completed", response: { usage: { input_tokens: 12, output_tokens: 3 } } },
      ];
      return new Response(events.map((event) => `data: ${JSON.stringify(event)}\n\n`).join(""), {
        headers: { "content-type": "text/event-stream" },
      });
    });
    try {
      await h.handlers.get("context")![0]!({ messages: branch.map((entry) => entry.message) }, ctx);
      const headers: Record<string, string> = { "x-codex-beta-features": "existing-feature" };
      await h.handlers.get("before_provider_headers")![0]!({ headers }, ctx);
      assert.match(headers["x-codex-beta-features"]!, /existing-feature/);
      assert.notEqual(headers["x-codex-beta-features"], "existing-feature");
      await h.handlers.get("before_provider_request")![0]!({ payload: {
        model: "stale-cached-model", input: [{ role: "user", content: "STALE-RAW-HISTORY" }],
        store: true, stream: false,
      } }, ctx);
      const result = await h.handlers.get("session_before_compact")![0]!({
        reason: "manual", branchEntries: branch, signal: new AbortController().signal,
        preparation: { firstKeptEntryId: "u1" },
      }, ctx);
      assert.ok(result?.compaction, "native compaction should return a checkpoint");
      assert.equal(requests.length, 1);
      const request = requests[0]!;
      assert.equal(request.url, provider === "sub-api"
        ? "https://gateway.example/v1/responses" : "https://codex.example/codex/responses");
      assert.equal(request.headers.get("authorization"), `Bearer ${apiKey}`);
      assert.equal(request.headers.get("x-test-route"), "preserved");
      assert.equal(request.headers.get("chatgpt-account-id"), provider === "sub-api" ? null : "test-account");
      assert.equal(request.body.model, "selected-model");
      assert.equal(request.body.stream, true);
      assert.equal(request.body.store, false);
      assert.deepEqual(request.body.input.at(-1), { type: "compaction_trigger" });
      assert.match(JSON.stringify(request.body.input), /Remember the test decision/);
      assert.doesNotMatch(JSON.stringify(request.body.input), /STALE-RAW-HISTORY/);
      const details = result.compaction.details;
      assert.equal(details.checkpoint, "native");
      assert.equal(details.modelKey, `${provider}:${ctx.model.api}:selected-model`);
      assert.ok(details.replacementHistory.some((item: any) => item.encrypted_content === "opaque-checkpoint"));
      assert.ok(details.portableSummary.length > 0);
      branch.push(
        { type: "compaction", id: "c1", parentId: "u1", ...result.compaction },
        { type: "message", id: "u2", parentId: "c1", message: { role: "user", content: "Continue after compaction.", timestamp: 2 } },
      );
      const projected = await h.handlers.get("context")![0]!({ messages: [] }, ctx);
      assert.doesNotMatch(JSON.stringify(projected.messages), /opaque-checkpoint/);
      const payload = await h.handlers.get("before_provider_request")![0]!({
        payload: { model: ctx.model.id, input: [], messages: [], previous_response_id: "stale-response" },
      }, ctx);
      assert.ok(payload.input.some((item: any) => item.encrypted_content === "opaque-checkpoint"));
      assert.match(JSON.stringify(payload.input), /Continue after compaction/);
      assert.equal(payload.previous_response_id, undefined);
      assert.equal(payload.messages, undefined);
      ctx.model = { ...ctx.model, id: "another-model" };
      const switched = await h.handlers.get("context")![0]!({ messages: [] }, ctx);
      assert.match(JSON.stringify(switched.messages), /Remember the test decision/);
      assert.equal(await h.handlers.get("before_provider_request")![0]!({ payload: { input: [] } }, ctx), undefined);
    } finally {
      await h.handlers.get("session_shutdown")![0]!({}, ctx);
    }
  });
}

test("SubAPI native failures and cancellation return control to Pi without a checkpoint", async (t) => {
  const h = harness();
  piCompaction(h.pi);
  const branch: any[] = [{
    type: "message", id: "u1", parentId: null,
    message: { role: "user", content: "Test only.", timestamp: 1 },
  }];
  const { ctx } = context(branch);
  ctx.model = { provider: "sub-api", api: "openai-responses", id: "test-model", input: ["text"], baseUrl: "https://gateway.example/v1" };
  ctx.modelRegistry.getApiKeyAndHeaders = async () => ({ ok: true, apiKey: "test-key" });
  const responses = [
    new Response('{"error":{"message":"unsupported"}}', { status: 400 }),
    new Response('data: {"type":"response.completed","response":{}}\n\n', { headers: { "content-type": "text/event-stream" } }),
  ];
  let calls = 0;
  t.mock.method(globalThis, "fetch", async (_url: string, init: RequestInit) => {
    init.signal?.throwIfAborted();
    calls++;
    return responses.shift()!;
  });
  try {
    for (let attempt = 0; attempt < 2; attempt++) {
      assert.equal(await h.handlers.get("session_before_compact")![0]!({
        reason: "manual", branchEntries: branch, preparation: { firstKeptEntryId: "u1" },
        signal: new AbortController().signal,
      }, ctx), undefined);
    }
    const controller = new AbortController();
    controller.abort();
    assert.equal(await h.handlers.get("session_before_compact")![0]!({
      reason: "manual", branchEntries: branch, preparation: { firstKeptEntryId: "u1" }, signal: controller.signal,
    }, ctx), undefined);
    assert.equal(calls, 2);
    assert.equal(branch.some((entry) => entry.type === "compaction"), false);
  } finally {
    await h.handlers.get("session_shutdown")![0]!({}, ctx);
  }
});

test("other providers and SubAPI Chat Completions keep their provider requests unchanged", async () => {
  const h = harness();
  piCompaction(h.pi);
  const { ctx } = context();
  try {
    for (const model of [
      { provider: "other", api: "openai-responses", id: "test" },
      { provider: "sub-api", api: "openai-completions", id: "test" },
    ]) {
      ctx.model = model;
      const headers: Record<string, string> = {};
      await h.handlers.get("before_provider_headers")![0]!({ headers }, ctx);
      assert.equal(headers["x-codex-beta-features"], undefined);
      assert.equal(await h.handlers.get("before_provider_request")![0]!({ payload: { input: [] } }, ctx), undefined);
    }
  } finally {
    await h.handlers.get("session_shutdown")![0]!({}, ctx);
  }
});

test("native runtime delegates authentication to each provider without blocking SubAPI header-only auth", async (t) => {
  const cases: Array<{ name: string; provider?: string; auth: any; succeeds: boolean; authorization?: string; gatewayAuthorization?: string }> = [
    { name: "explicit header", auth: { ok: true, apiKey: "dummy-key", headers: { Authorization: "Bearer explicit" } }, succeeds: true, authorization: "Bearer explicit" },
    { name: "header only", auth: { ok: true, headers: { authorization: "Bearer header-only" } }, succeeds: true, authorization: "Bearer header-only" },
    { name: "proxy header only", auth: { ok: true, headers: { "CF-AIG-Authorization": "Bearer proxy" } }, succeeds: true, gatewayAuthorization: "Bearer proxy" },
    { name: "blank header", auth: { ok: true, headers: { Authorization: " " } }, succeeds: false },
    { name: "missing credentials", auth: { ok: true }, succeeds: false },
    { name: "registry failure", auth: { ok: false, error: "credentials unavailable" }, succeeds: false },
    { name: "Codex header only", provider: "openai-codex", auth: { ok: true, headers: { Authorization: "Bearer header-only" } }, succeeds: false },
    { name: "Codex invalid JWT", provider: "openai-codex", auth: { ok: true, apiKey: "invalid", headers: { Authorization: "Bearer header-only" } }, succeeds: false },
  ];
  for (const scenario of cases) {
    await t.test(scenario.name, async (t) => {
      const h = harness();
      piCompaction(h.pi);
      const branch: any[] = [{
        type: "message", id: "u1", parentId: null,
        message: { role: "user", content: "Synthetic authentication test.", timestamp: 1 },
      }];
      const { ctx } = context(branch);
      const provider = scenario.provider ?? "sub-api";
      ctx.model = {
        provider, api: provider === "sub-api" ? "openai-responses" : "openai-codex-responses",
        id: "test-model", input: ["text"], baseUrl: "https://gateway.example/v1",
      };
      ctx.modelRegistry.getApiKeyAndHeaders = async () => scenario.auth;
      let calls = 0;
      t.mock.method(globalThis, "fetch", async (_url: string, init: RequestInit) => {
        calls++;
        const headers = new Headers(init.headers);
        assert.equal(headers.get("authorization"), scenario.authorization ?? null);
        assert.equal(headers.get("cf-aig-authorization"), scenario.gatewayAuthorization ?? null);
        const events = [
          { type: "response.output_item.done", item: { type: "compaction", encrypted_content: "test-checkpoint" } },
          { type: "response.completed", response: {} },
        ];
        return new Response(events.map((event) => `data: ${JSON.stringify(event)}\n\n`).join(""), {
          headers: { "content-type": "text/event-stream" },
        });
      });
      try {
        const result = await h.handlers.get("session_before_compact")![0]!({
          reason: "manual", branchEntries: branch, preparation: { firstKeptEntryId: "u1" },
          signal: new AbortController().signal,
        }, ctx);
        assert.equal(Boolean(result?.compaction), scenario.succeeds);
        assert.equal(calls, scenario.succeeds ? 1 : 0);
      } finally {
        await h.handlers.get("session_shutdown")![0]!({}, ctx);
      }
    });
  }
});