feat: unify context compaction

pi-compaction/index.ts

@@ -0,0 +1,590 @@
+import { createHash, randomUUID } from "node:crypto";
+import { Type } from "typebox";
+import {
+  buildSessionContext,
+  convertToLlm,
+  estimateTokens,
+  generateSummaryWithUsage,
+  sessionEntryToContextMessages,
+  truncateHead,
+  DEFAULT_MAX_BYTES,
+  DEFAULT_MAX_LINES,
+  type ExtensionAPI,
+  type ExtensionContext,
+  type SessionEntry,
+} from "@earendil-works/pi-coding-agent";
+import { DEFAULT_CONFIG, loadConfig } from "./src/config.js";
+import {
+  buildCodexHeaders,
+  buildCompactionRequestBody,
+  buildReplacementHistory,
+  buildToolPayload,
+  callRemoteCompaction,
+  findNativeCheckpoint,
+  isJsonObject,
+  isOpenAICodexModel,
+  mergeFeatureHeader,
+  messagesToResponseItems,
+  modelKey,
+  NATIVE_COMPACTION_KIND,
+  NATIVE_COMPACTION_VERSION,
+  resolveCodexResponsesUrl,
+  stripInputFromPayload,
+  type JsonObject,
+  type NativeCompactionDetails,
+  type ResponseItem,
+} from "./src/native-compaction.js";
+import { SegmentIndex } from "./src/segment-index.js";
+import { SemanticWorker } from "./src/semantic-worker.js";
+import { shouldCompact } from "./src/token-accounting.js";
+import {
+  CompactionServiceRegistry,
+  publishCompactionService,
+  unpublishCompactionService,
+} from "./src/service.js";
+import {
+  COMPACTION_CUSTOM_TYPE,
+  COMPACTION_VERSION,
+  QUERY_TOOL_NAME,
+  type PiCompactionConfig,
+  type ProjectionSnapshot,
+  type ToolHistoryRecord,
+} from "./src/types.js";
+
+interface SessionRuntime {
+  index: SegmentIndex;
+  config: PiCompactionConfig;
+  codexCache?: { fingerprint: string; modelKey: string; toolsHash: string; items: ResponseItem[] };
+}
+
+interface CachedPayload {
+  modelKey: string;
+  payload: JsonObject;
+}
+
+function errorMessage(error: unknown): string {
+  return error instanceof Error ? error.message : String(error);
+}
+
+function hash(value: string): string {
+  return createHash("sha256").update(value).digest("hex").slice(0, 20);
+}
+
+interface EnvironmentValue {
+  hash: string;
+  tokens: number;
+  contextWindow: number;
+}
+
+interface EnvironmentCacheEntry extends EnvironmentValue {
+  prompt: string;
+  modelIdentity: string;
+  activeNames: string[];
+  tools: unknown[];
+}
+
+function environment(
+  pi: ExtensionAPI,
+  ctx: ExtensionContext,
+  cache: Map<string, EnvironmentCacheEntry>,
+): EnvironmentValue {
+  const sessionId = ctx.sessionManager.getSessionId();
+  const prompt = ctx.getSystemPrompt();
+  const activeNames = pi.getActiveTools();
+  const active = new Set(activeNames);
+  const tools = pi.getAllTools().filter((tool) => active.has(tool.name));
+  const contextWindow = ctx.model?.contextWindow ?? 0;
+  const modelIdentity = ctx.model ? modelKey(ctx.model) : "none";
+  const previous = cache.get(sessionId);
+  if (
+    previous
+    && previous.prompt === prompt
+    && previous.contextWindow === contextWindow
+    && previous.modelIdentity === modelIdentity
+    && previous.activeNames.length === activeNames.length
+    && previous.activeNames.every((name, index) => name === activeNames[index])
+    && previous.tools.length === tools.length
+    && previous.tools.every((tool, index) => tool === tools[index])
+  ) return previous;
+
+  const serializedTools = JSON.stringify(tools.map((tool) => ({
+    name: tool.name,
+    description: tool.description,
+    parameters: tool.parameters,
+  })));
+  const text = `${modelIdentity}\n${prompt}\n${serializedTools}`;
+  const value: EnvironmentCacheEntry = {
+    prompt,
+    modelIdentity,
+    activeNames: [...activeNames],
+    tools,
+    hash: hash(text),
+    tokens: Math.ceil(text.length / 4),
+    contextWindow,
+  };
+  cache.set(sessionId, value);
+  return value;
+}
+
+function portableSummary(messages: readonly any[], maxChars = 24_000): string {
+  const sections: string[] = ["Portable context checkpoint generated by pi-compaction."];
+  for (const message of messages) {
+    if (message?.role === "toolResult") {
+      const text = Array.isArray(message.content)
+        ? message.content.flatMap((part: any) => part?.type === "text" ? [part.text] : []).join("\n")
+        : "";
+      if (text) sections.push(`[Tool ${message.toolName ?? "unknown"}] ${text}`);
+      continue;
+    }
+    if (message?.role === "user") {
+      const text = typeof message.content === "string"
+        ? message.content
+        : Array.isArray(message.content)
+          ? message.content.flatMap((part: any) => part?.type === "text" ? [part.text] : []).join("\n")
+          : "";
+      if (text) sections.push(`[User] ${text}`);
+      continue;
+    }
+    if (message?.role === "assistant" && Array.isArray(message.content)) {
+      const text = message.content.flatMap((part: any) => part?.type === "text" ? [part.text] : []).join("\n");
+      if (text) sections.push(`[Assistant] ${text}`);
+    }
+  }
+  const combined = sections.join("\n\n");
+  if (combined.length <= maxChars) return combined;
+  const head = Math.floor(maxChars * 0.35);
+  const tail = maxChars - head - 32;
+  return `${combined.slice(0, head)}\n\n… checkpoint truncated …\n\n${combined.slice(-tail)}`;
+}
+
+function asPortableMessage(summary: string): any {
+  return {
+    role: "user",
+    content: [{ type: "text", text: summary }],
+    timestamp: Date.now(),
+  };
+}
+
+function setFeatureHeader(headers: Record<string, string | null>): void {
+  const existing = Object.entries(headers).find(([name]) => name.toLowerCase() === "x-codex-beta-features");
+  if (existing) headers[existing[0]] = mergeFeatureHeader(existing[1]);
+  else headers["x-codex-beta-features"] = mergeFeatureHeader(undefined);
+}
+
+function cleanHeaders(headers: Record<string, string | null> | undefined): Record<string, string> | undefined {
+  if (!headers) return undefined;
+  return Object.fromEntries(Object.entries(headers).filter((entry): entry is [string, string] => typeof entry[1] === "string"));
+}
+
+function activeToolsHash(pi: ExtensionAPI): string {
+  return hash(pi.getActiveTools().join("\0"));
+}
+
+export default function piCompactionExtension(pi: ExtensionAPI): void {
+  const runtimes = new Map<string, SessionRuntime>();
+  const payloads = new Map<string, CachedPayload>();
+  const environments = new Map<string, EnvironmentCacheEntry>();
+  const worker = new SemanticWorker();
+  const service = new CompactionServiceRegistry();
+  publishCompactionService(service);
+
+  const runtimeFor = (ctx: ExtensionContext): SessionRuntime => {
+    const sessionId = ctx.sessionManager.getSessionId();
+    let runtime = runtimes.get(sessionId);
+    if (!runtime) {
+      const index = new SegmentIndex(sessionId);
+      index.hydrate(ctx.sessionManager.getBranch() as SessionEntry[], DEFAULT_CONFIG);
+      runtime = { index, config: structuredClone(DEFAULT_CONFIG) };
+      runtimes.set(sessionId, runtime);
+      service.setSession(sessionId, (ids) => index.query(ids));
+    }
+    return runtime;
+  };
+
+  const persist = (record: unknown): void => {
+    pi.appendEntry(COMPACTION_CUSTOM_TYPE, record);
+  };
+
+  const drainSemantic = (ctx: ExtensionContext, runtime: SessionRuntime): void => {
+    const sessionId = ctx.sessionManager.getSessionId();
+    for (const result of worker.drain(sessionId)) {
+      const record = runtime.index.applySemantic(result.sourceKey, result.sourceHash, result.summary);
+      if (record) persist(record);
+    }
+    runtime.index.setSemanticJobs(worker.count(sessionId));
+  };
+
+  const projectMessages = (
+    ctx: ExtensionContext,
+    runtime: SessionRuntime,
+    messages: readonly any[],
+  ): ProjectionSnapshot => {
+    const env = environment(pi, ctx, environments);
+    runtime.index.setSemanticJobs(worker.count(ctx.sessionManager.getSessionId()));
+    const snapshot = runtime.index.project(messages, {
+      environmentHash: env.hash,
+      environmentTokens: env.tokens,
+      contextWindow: env.contextWindow,
+    });
+    service.setSnapshot(snapshot);
+    return snapshot;
+  };
+
+  const branchMessages = (
+    ctx: ExtensionContext,
+    runtime: SessionRuntime,
+    branch: SessionEntry[],
+    excludeLastAssistantError = false,
+  ): { snapshot: ProjectionSnapshot; checkpoint: ReturnType<typeof findNativeCheckpoint> } => {
+    const checkpoint = findNativeCheckpoint(branch);
+    let messages: any[];
+    if (checkpoint.status === "valid") {
+      messages = branch
+        .slice(checkpoint.checkpoint.entryIndex + 1)
+        .flatMap((entry) => sessionEntryToContextMessages(entry));
+      if (checkpoint.checkpoint.details.modelKey !== (ctx.model ? modelKey(ctx.model) : "")) {
+        messages = [asPortableMessage(checkpoint.checkpoint.details.portableSummary), ...messages];
+      }
+    } else {
+      messages = buildSessionContext(branch).messages;
+    }
+    if (excludeLastAssistantError) {
+      const index = messages.findLastIndex((message) => message.role === "assistant");
+      if (index >= 0 && messages[index]?.stopReason === "error") messages = messages.toSpliced(index, 1);
+    }
+    return { snapshot: projectMessages(ctx, runtime, messages), checkpoint };
+  };
+
+  const codexInput = (
+    piModel: any,
+    runtime: SessionRuntime,
+    snapshot: ProjectionSnapshot,
+    checkpoint?: NativeCompactionDetails,
+  ): ResponseItem[] => {
+    const toolsHash = activeToolsHash(pi);
+    const key = modelKey(piModel);
+    const cached = runtime.codexCache;
+    if (cached && cached.fingerprint === snapshot.fingerprint && cached.modelKey === key && cached.toolsHash === toolsHash) {
+      return cached.items;
+    }
+    const tail = messagesToResponseItems(piModel, convertToLlm(snapshot.messages as any), pi.getAllTools());
+    const items = checkpoint && checkpoint.modelKey === key
+      ? [...structuredClone(checkpoint.replacementHistory), ...tail]
+      : tail;
+    runtime.codexCache = { fingerprint: snapshot.fingerprint, modelKey: key, toolsHash, items };
+    snapshot.codexItems = items;
+    return items;
+  };
+
+  pi.on("session_start", async (_event, ctx) => {
+    publishCompactionService(service);
+    const sessionId = ctx.sessionManager.getSessionId();
+    const config = await loadConfig();
+    const index = new SegmentIndex(sessionId);
+    index.hydrate(ctx.sessionManager.getBranch() as SessionEntry[], config);
+    runtimes.set(sessionId, { index, config });
+    service.setSession(sessionId, (ids) => index.query(ids));
+  });
+
+  pi.on("session_shutdown", (_event, ctx) => {
+    const sessionId = ctx.sessionManager.getSessionId();
+    worker.cancelSession(sessionId);
+    runtimes.delete(sessionId);
+    payloads.delete(sessionId);
+    environments.delete(sessionId);
+    service.deleteSession(sessionId);
+    if (runtimes.size === 0) unpublishCompactionService(service);
+  });
+
+  pi.on("session_tree", (_event, ctx) => {
+    const runtime = runtimeFor(ctx);
+    runtime.index.hydrate(ctx.sessionManager.getBranch() as SessionEntry[], runtime.config);
+    runtime.codexCache = undefined;
+    service.clearSnapshot(ctx.sessionManager.getSessionId());
+  });
+
+  pi.on("session_compact", (_event, ctx) => {
+    const runtime = runtimeFor(ctx);
+    runtime.index.hydrate(ctx.sessionManager.getBranch() as SessionEntry[], runtime.config);
+    runtime.codexCache = undefined;
+    service.clearSnapshot(ctx.sessionManager.getSessionId());
+    payloads.delete(ctx.sessionManager.getSessionId());
+  });
+
+  pi.on("model_select", (_event, ctx) => {
+    const sessionId = ctx.sessionManager.getSessionId();
+    payloads.delete(sessionId);
+    environments.delete(sessionId);
+    const runtime = runtimes.get(sessionId);
+    if (runtime) runtime.codexCache = undefined;
+  });
+
+  pi.on("turn_end", (event, ctx) => {
+    const runtime = runtimeFor(ctx);
+    drainSemantic(ctx, runtime);
+    const previous = runtime.index.getSnapshot();
+    const captured = runtime.index.captureTurn(event.message, event.toolResults, runtime.config);
+    for (const record of captured) persist(runtime.index.toPersisted(record));
+    const pressure = previous && previous.contextWindow > 0
+      ? previous.upperBound / previous.contextWindow >= runtime.config.softRatio
+      : false;
+    const semanticCandidates = pressure
+      ? runtime.index.allRecords()
+      : captured.filter((record) => record.rawText.length >= runtime.config.semantic.maxInputChars);
+    for (const record of semanticCandidates) {
+      worker.enqueue(ctx.sessionManager.getSessionId(), record, ctx, runtime.config);
+    }
+    runtime.index.setSemanticJobs(worker.count(ctx.sessionManager.getSessionId()));
+
+    const usage = ctx.getContextUsage();
+    if (usage?.tokens && previous) {
+      const env = environment(pi, ctx, environments);
+      runtime.index.recordUsage(
+        usage.tokens,
+        previous.messageEstimate + estimateTokens(event.message as any),
+        env.hash,
+      );
+    }
+  });
+
+  pi.on("context", (event, ctx) => {
+    const runtime = runtimeFor(ctx);
+    const branch = ctx.sessionManager.getBranch() as SessionEntry[];
+    const checkpoint = findNativeCheckpoint(branch);
+    if (checkpoint.status === "invalid") {
+      ctx.abort();
+      if (ctx.hasUI) ctx.ui.notify("pi-compaction checkpoint is malformed; request blocked.", "error");
+      return { messages: [] };
+    }
+    if (checkpoint.status === "valid") {
+      const projected = branchMessages(ctx, runtime, branch);
+      return { messages: projected.snapshot.messages as any };
+    }
+    return { messages: projectMessages(ctx, runtime, event.messages).messages as any };
+  });
+
+  pi.on("before_provider_headers", (event, ctx) => {
+    if (isOpenAICodexModel(ctx.model)) setFeatureHeader(event.headers);
+  });
+
+  pi.on("before_provider_request", (event, ctx) => {
+    const model = ctx.model;
+    if (!model || !isOpenAICodexModel(model) || !isJsonObject(event.payload)) return undefined;
+    const runtime = runtimeFor(ctx);
+    const sessionId = ctx.sessionManager.getSessionId();
+    payloads.set(sessionId, { modelKey: modelKey(model), payload: stripInputFromPayload(event.payload) });
+    const checkpoint = findNativeCheckpoint(ctx.sessionManager.getBranch() as SessionEntry[]);
+    if (checkpoint.status !== "valid" || checkpoint.checkpoint.details.modelKey !== modelKey(model)) return undefined;
+    const snapshot = runtime.index.getSnapshot();
+    if (!snapshot) return undefined;
+    const payload: JsonObject = {
+      ...event.payload,
+      input: codexInput(model, runtime, snapshot, checkpoint.checkpoint.details),
+    };
+    delete payload.messages;
+    delete payload.previous_response_id;
+    return payload;
+  });
+
+  pi.on("session_before_compact", async (event, ctx) => {
+    const runtime = runtimeFor(ctx);
+    if (!runtime.config.enabled) return undefined;
+    drainSemantic(ctx, runtime);
+    const { snapshot, checkpoint } = branchMessages(
+      ctx,
+      runtime,
+      event.branchEntries as SessionEntry[],
+      event.reason === "overflow" && event.willRetry,
+    );
+
+    if (!shouldCompact(event.reason, snapshot, runtime.config)) return { cancel: true };
+
+    const model = ctx.model;
+    if (!model) return undefined;
+    if (!runtime.config.native.enabled || !isOpenAICodexModel(model)) {
+      try {
+        const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
+        if (!auth.ok) return undefined;
+        const result = await generateSummaryWithUsage(
+          snapshot.messages as any,
+          model,
+          16_384,
+          auth.apiKey,
+          cleanHeaders(auth.headers),
+          event.signal,
+          event.customInstructions,
+          undefined,
+          runtime.config.semantic.thinking,
+        );
+        return {
+          compaction: {
+            summary: result.text,
+            firstKeptEntryId: event.preparation.firstKeptEntryId,
+            tokensBefore: snapshot.effectiveTokens,
+            usage: result.usage,
+            details: {
+              kind: COMPACTION_CUSTOM_TYPE,
+              version: COMPACTION_VERSION,
+              checkpoint: "portable",
+              sourceFingerprint: snapshot.fingerprint,
+            },
+          },
+        };
+      } catch (error) {
+        if (!event.signal.aborted) console.error("pi-compaction: portable checkpoint failed:", error);
+        return undefined;
+      }
+    }
+
+    try {
+      const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
+      if (!auth.ok || !auth.apiKey) throw new Error(auth.ok ? "Codex authentication unavailable" : auth.error);
+      const currentCheckpoint = checkpoint.status === "valid" ? checkpoint.checkpoint.details : undefined;
+      const input = codexInput(model, runtime, snapshot, currentCheckpoint);
+      const sessionId = ctx.sessionManager.getSessionId();
+      const cached = payloads.get(sessionId);
+      const body = buildCompactionRequestBody({
+        basePayload: cached?.modelKey === modelKey(model) ? cached.payload : undefined,
+        model,
+        input,
+        instructions: ctx.getSystemPrompt(),
+        tools: buildToolPayload(pi.getAllTools(), pi.getActiveTools()),
+        sessionId,
+      });
+      if (ctx.hasUI) ctx.ui.notify("pi-compaction: creating Codex native checkpoint…", "info");
+      const remote = await callRemoteCompaction({
+        url: resolveCodexResponsesUrl(model.baseUrl),
+        headers: buildCodexHeaders({ apiKey: auth.apiKey, headers: cleanHeaders(auth.headers), sessionId }),
+        body,
+        signal: event.signal,
+      });
+      const details: NativeCompactionDetails = {
+        kind: NATIVE_COMPACTION_KIND,
+        version: NATIVE_COMPACTION_VERSION,
+        checkpoint: "native",
+        modelKey: modelKey(model),
+        replacementHistory: buildReplacementHistory(
+          input,
+          remote.compactionItem,
+          Math.min(64_000, Math.max(4_096, Math.floor(snapshot.contextWindow * Math.max(0.05, runtime.config.targetRatio - 0.3)))),
+        ),
+        portableSummary: portableSummary(snapshot.messages),
+        sourceFingerprint: snapshot.fingerprint,
+      };
+      return {
+        compaction: {
+          summary: `Pi native checkpoint ${randomUUID()}`,
+          firstKeptEntryId: event.preparation.firstKeptEntryId,
+          tokensBefore: snapshot.effectiveTokens,
+          usage: remote.usage,
+          details,
+        },
+      };
+    } catch (error) {
+      if (!event.signal.aborted && ctx.hasUI) {
+        ctx.ui.notify(`pi-compaction native checkpoint failed; falling back to Pi: ${errorMessage(error)}`, "warning");
+      }
+      return undefined;
+    }
+  });
+
+  pi.on("session_before_tree", async (event, ctx) => {
+    const model = ctx.model;
+    if (!event.preparation.userWantsSummary || !model) return undefined;
+    const runtime = runtimeFor(ctx);
+    if (!runtime.config.enabled) return undefined;
+    const messages = event.preparation.entriesToSummarize.flatMap((entry) => sessionEntryToContextMessages(entry));
+    const snapshot = projectMessages(ctx, runtime, messages);
+    try {
+      const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
+      if (!auth.ok) return undefined;
+      const result = await generateSummaryWithUsage(
+        snapshot.messages as any,
+        model,
+        16_384,
+        auth.apiKey,
+        cleanHeaders(auth.headers),
+        event.signal,
+        event.preparation.customInstructions,
+        undefined,
+        runtime.config.semantic.thinking,
+        undefined,
+        undefined,
+        undefined,
+        undefined,
+      );
+      return {
+        summary: {
+          summary: result.text,
+          usage: result.usage,
+          details: {
+            kind: COMPACTION_CUSTOM_TYPE,
+            version: COMPACTION_VERSION,
+            checkpoint: "portable",
+            sourceFingerprint: snapshot.fingerprint,
+          },
+        },
+      };
+    } catch (error) {
+      if (!event.signal.aborted) console.error("pi-compaction: branch summary failed:", error);
+      return undefined;
+    }
+  });
+
+  pi.registerTool({
+    name: QUERY_TOOL_NAME,
+    label: "Query Compacted Tool History",
+    description: "Recover original tool outputs by source key or tool call id after pi-compaction replaced them with bounded summaries.",
+    promptSnippet: "Recover original compacted tool evidence",
+    promptGuidelines: [
+      "Use context_tree_query only when a pi-compaction summary lacks evidence needed for the current task.",
+    ],
+    parameters: Type.Object({
+      toolCallIds: Type.Array(Type.String(), { minItems: 1, maxItems: 20 }),
+    }),
+    async execute(_id, params, _signal, _update, ctx) {
+      const records = service.query(ctx.sessionManager.getSessionId(), params.toolCallIds);
+      const blocks = records.map((record: ToolHistoryRecord) => {
+        const perResultBytes = Math.max(1_024, Math.floor(DEFAULT_MAX_BYTES / Math.max(1, records.length)));
+        const args = truncateHead(JSON.stringify(record.args), { maxLines: 20, maxBytes: 2_048 }).content;
+        const truncated = truncateHead(record.rawText, {
+          maxLines: Math.max(20, Math.floor(DEFAULT_MAX_LINES / Math.max(1, records.length))),
+          maxBytes: perResultBytes,
+        });
+        return [
+          `## ${record.sourceKey}`,
+          `Tool: ${record.toolName}`,
+          `Args: ${args}`,
+          `Status: ${record.isError ? "ERROR" : "OK"}`,
+          "",
+          truncated.content,
+          ...(truncated.truncated ? [`[truncated: ${truncated.outputLines}/${truncated.totalLines} lines]`] : []),
+        ].join("\n");
+      });
+      const missing = params.toolCallIds.filter((id: string) => !records.some((record) => record.sourceKey === id || record.toolCallId === id));
+      for (const id of missing) blocks.push(`## ${id}\n(not found)`);
+      return {
+        content: [{ type: "text", text: blocks.join("\n\n---\n\n") }],
+        details: { sourceKeys: records.map((record) => record.sourceKey), missing },
+      };
+    },
+  });
+
+  pi.registerCommand("compaction-status", {
+    description: "Show effective/raw context usage and pi-compaction state",
+    handler: async (_args, ctx) => {
+      const metrics = service.getMetrics(ctx.sessionManager.getSessionId());
+      const text = metrics
+        ? [
+            `Effective: ${metrics.effectiveTokens}/${metrics.contextWindow} (${metrics.percent.toFixed(1)}%, ${metrics.confidence})`,
+            `Raw: ${metrics.rawTokens}/${metrics.contextWindow}`,
+            `Reclaimed: ${metrics.reclaimedTokens}`,
+            `Segments: ${metrics.segmentCount}; semantic jobs: ${metrics.semanticJobs}; generation: ${metrics.generation}`,
+          ].join("\n")
+        : "pi-compaction has no projection snapshot yet.";
+      if (ctx.hasUI) ctx.ui.notify(text, "info");
+      else process.stderr.write(`${text}\n`);
+    },
+  });
+
+}