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 { applyContextEdits, planMasking, rebuildMaskedLedger, type MaskedLedgerEntry, } from "./src/masking.js"; import { buildNativeCompactionTransport, isNativeCompactionModel } from "./src/native-provider.js"; import { buildCompactionRequestBody, buildReplacementHistory, buildToolPayload, callRemoteCompaction, findNativeCheckpoint, isJsonObject, mergeFeatureHeader, messagesToResponseItems, modelKey, NATIVE_COMPACTION_KIND, NATIVE_COMPACTION_VERSION, stripInputFromPayload, type JsonObject, type NativeCompactionDetails, type ResponseItem, } from "./src/native-compaction.js"; import { contentText, SegmentIndex } from "./src/segment-index.js"; import { CompactionServiceRegistry, publishCompactionService, unpublishCompactionService, } from "./src/service.js"; import { effectiveThreshold } from "./src/token-accounting.js"; import { COMPACTION_CUSTOM_TYPE, COMPACTION_VERSION, QUERY_TOOL_NAME, type MaskedResultRecord, type PiCompactionConfig, type ProjectionSnapshot, } from "./src/types.js"; interface SessionRuntime { index: SegmentIndex; config: PiCompactionConfig; codexCache?: { fingerprint: string; modelKey: string; toolsHash: string; items: ResponseItem[] }; masked: Map; lastAutoCompactLeafId: string | null; } interface CachedPayload { modelKey: string; payload: JsonObject; } function errorMessage(error: unknown): string { return error instanceof Error ? error.message : String(error); } interface EnvironmentValue { hash: string; tokens: number; contextWindow: number; } interface EnvironmentCacheEntry extends EnvironmentValue { prompt: string; modelIdentity: string; activeNames: string[]; tools: unknown[]; } function hash(value: string): string { return createHash("sha256").update(value).digest("hex").slice(0, 20); } function environment( pi: ExtensionAPI, ctx: ExtensionContext, cache: Map, ): 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): 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 | undefined): Record | 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(); const payloads = new Map(); const environments = new Map(); const service = new CompactionServiceRegistry(); publishCompactionService(service); const maskedRecords = (sessionId: string, ids: string[]): MaskedResultRecord[] => { const masked = runtimes.get(sessionId)?.masked; if (!masked) return []; return ids.flatMap((id) => { const entry = masked.get(id); return entry ? [{ toolCallId: id, resultTimestamp: entry.resultTimestamp, deterministicSummary: entry.placeholder, }] : []; }); }; const runtimeFor = (ctx: ExtensionContext): SessionRuntime => { const sessionId = ctx.sessionManager.getSessionId(); let runtime = runtimes.get(sessionId); if (!runtime) { runtime = { index: new SegmentIndex(sessionId), config: structuredClone(DEFAULT_CONFIG), masked: rebuildLedger(ctx), lastAutoCompactLeafId: null, }; runtimes.set(sessionId, runtime); service.setSession(sessionId, (ids) => maskedRecords(sessionId, ids)); } return runtime; }; const ledgerStats = (runtime: SessionRuntime): { maskedResults: number; reclaimedTokens: number } => { let reclaimedTokens = 0; for (const entry of runtime.masked.values()) { reclaimedTokens += Math.max(0, entry.originalTokens - entry.placeholderTokens); } return { maskedResults: runtime.masked.size, reclaimedTokens }; }; const projectMessages = ( ctx: ExtensionContext, runtime: SessionRuntime, messages: readonly any[], ): ProjectionSnapshot => { const env = environment(pi, ctx, environments); const stats = ledgerStats(runtime); const snapshot = runtime.index.project(messages, { environmentHash: env.hash, environmentTokens: env.tokens, contextWindow: env.contextWindow, maskedResults: stats.maskedResults, reclaimedTokens: stats.reclaimedTokens, }); service.setSnapshot(snapshot); return snapshot; }; const rebuildLedger = (ctx: ExtensionContext): Map => rebuildMaskedLedger( ctx.sessionManager.getBranch(), ctx.sessionManager.getLeafId(), (message) => estimateTokens(message as any), ); const hydrateRuntime = (ctx: ExtensionContext, runtime: SessionRuntime): void => { runtime.index.hydrate(); runtime.masked = rebuildLedger(ctx); runtime.lastAutoCompactLeafId = null; runtime.codexCache = undefined; }; const branchMessages = ( ctx: ExtensionContext, runtime: SessionRuntime, branch: SessionEntry[], excludeLastAssistantError = false, ): { snapshot: ProjectionSnapshot; checkpoint: ReturnType } => { const checkpoint = findNativeCheckpoint(branch); let messages: any[]; if (checkpoint.status === "valid") { // sessionEntryToContextMessages sees only raw entries; applyContextEdits // additionally replays the latest context_edit per target (omissions and // placeholder replacements included). messages = applyContextEdits(branch.slice(checkpoint.checkpoint.entryIndex + 1)); 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 runtime = runtimes.get(sessionId) ?? { index: new SegmentIndex(sessionId), config, masked: new Map(), lastAutoCompactLeafId: null, }; runtime.config = config; runtime.index.hydrate(); runtime.masked = rebuildLedger(ctx); runtime.lastAutoCompactLeafId = null; runtimes.set(sessionId, runtime); service.setSession(sessionId, (ids) => maskedRecords(sessionId, ids)); }); pi.on("session_shutdown", (_event, ctx) => { const sessionId = ctx.sessionManager.getSessionId(); 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); hydrateRuntime(ctx, runtime); service.clearSnapshot(ctx.sessionManager.getSessionId()); }); pi.on("session_compact", (_event, ctx) => { const runtime = runtimeFor(ctx); hydrateRuntime(ctx, runtime); 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); if (!runtime.config.enabled) return undefined; const previous = runtime.index.getSnapshot(); 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, ); } const currentTokens = usage?.tokens ?? (event.context?.contextMessages ?? []).reduce((sum, message) => sum + estimateTokens(message as any), 0) + environment(pi, ctx, environments).tokens; const plan = planMasking({ contextEntries: event.context?.contextEntries ?? [], currentTokens, contextWindow: ctx.model?.contextWindow ?? 0, masking: runtime.config.masking, estimate: (message) => estimateTokens(message as any), }); if (!plan.drafts.length) return undefined; for (const result of plan.masked) { runtime.masked.set(result.toolCallId, { placeholder: result.placeholder, originalTokens: result.originalTokens, placeholderTokens: result.placeholderTokens, resultTimestamp: result.resultTimestamp, }); } return { entries: [...(event.entries ?? []), ...plan.drafts] }; }); pi.on("agent_settled", (_event, ctx) => { const runtime = runtimeFor(ctx); if (!runtime.config.enabled || ctx.isIdle?.() === false) return; const snapshot = runtime.index.getSnapshot(); const tokens = ctx.getContextUsage()?.tokens ?? snapshot?.effectiveTokens; if (tokens == null) return; const contextWindow = ctx.model?.contextWindow ?? snapshot?.contextWindow ?? 0; const threshold = effectiveThreshold( runtime.config.summary.triggerTokens, runtime.config.summary.triggerRatio, contextWindow, ); if (tokens < threshold) return; const leafId = ctx.sessionManager.getLeafId(); if (leafId !== null && runtime.lastAutoCompactLeafId === leafId) return; runtime.lastAutoCompactLeafId = leafId; ctx.compact({ onError: (error) => { if (ctx.hasUI) ctx.ui.notify(`pi-compaction auto summary failed: ${errorMessage(error)}`, "warning"); }, }); }); 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 }; } // Non-checkpoint path only refreshes the shared snapshot; the message list // is returned to Pi unchanged — masking is append-only via context_edit. projectMessages(ctx, runtime, event.messages); return undefined; }); pi.on("before_provider_headers", (event, ctx) => { if (isNativeCompactionModel(ctx.model)) setFeatureHeader(event.headers); }); pi.on("before_provider_request", (event, ctx) => { const model = ctx.model; if (!model || !isNativeCompactionModel(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; const { snapshot, checkpoint } = branchMessages( ctx, runtime, event.branchEntries as SessionEntry[], event.reason === "overflow" && event.willRetry, ); const model = ctx.model; if (!model) return undefined; if (!runtime.config.native.enabled || !isNativeCompactionModel(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.summary.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) throw new Error(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 ${model.provider} native checkpoint...`, "info"); const remote = await callRemoteCompaction({ ...buildNativeCompactionTransport({ model, 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 * 0.25))), ), 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.summary.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 the original output of a tool result that pi-compaction cleared from context, by tool call id.", promptSnippet: "Recover original tool output cleared from context", promptGuidelines: [ "Use context_tree_query when you see a \"[pi-compaction] Earlier\" placeholder and still need the original tool output.", ], parameters: Type.Object({ toolCallIds: Type.Array(Type.String(), { minItems: 1, maxItems: 20 }), }), async execute(_id, params, _signal, _update, ctx) { const branch = ctx.sessionManager.getBranch() as SessionEntry[]; const calls = new Map(); const results = new Map(); for (const entry of branch) { const message = (entry as any)?.message; if ((entry as any)?.type !== "message" || !message) continue; if (message.role === "assistant" && Array.isArray(message.content)) { for (const block of message.content) { if (block?.type === "toolCall" && typeof block.id === "string") { calls.set(block.id, { toolName: String(block.name ?? "unknown"), args: block.arguments ?? block.input ?? block.args ?? {}, }); } } } else if (message.role === "toolResult" && typeof message.toolCallId === "string") { results.set(message.toolCallId, message); } } const found: string[] = []; const blocks: string[] = []; for (const id of params.toolCallIds) { const result = results.get(id); if (!result) continue; found.push(id); const call = calls.get(id) ?? { toolName: String(result.toolName ?? "unknown"), args: {} }; const perResultBytes = Math.max(1_024, Math.floor(DEFAULT_MAX_BYTES / Math.max(1, params.toolCallIds.length))); const args = truncateHead(JSON.stringify(call.args ?? {}), { maxLines: 20, maxBytes: 2_048 }).content; const rawText = contentText(result.content); const truncated = truncateHead(rawText, { maxLines: Math.max(20, Math.floor(DEFAULT_MAX_LINES / Math.max(1, params.toolCallIds.length))), maxBytes: perResultBytes, }); blocks.push([ `## ${id}`, `Tool: ${call.toolName}`, `Args: ${args}`, `Status: ${result.isError ? "ERROR" : "OK"}`, "", truncated.content, ...(truncated.truncated ? [`[truncated: ${truncated.outputLines}/${truncated.totalLines} lines]`] : []), ].join("\n")); } const missing = params.toolCallIds.filter((id: string) => !results.has(id)); for (const id of missing) blocks.push(`## ${id}\n(not found)`); return { content: [{ type: "text", text: blocks.join("\n\n---\n\n") }], details: { toolCallIds: found, 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}; masked tool results: ${metrics.maskedResults}; 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`); }, }); }