@narumitw/pi-codex-compact 0.47.0

package/src/codex-compact.ts

@@ -0,0 +1,283 @@
+import type { AgentMessage } from "@earendil-works/pi-agent-core";
+import type { Api, Context, Model, Tool } from "@earendil-works/pi-ai";
+import { hasApi } from "@earendil-works/pi-ai";
+import {
+	buildContextEntries,
+	buildSessionContext,
+	convertToLlm,
+	type ExtensionAPI,
+	type ExtensionContext,
+	type SessionBeforeCompactEvent,
+	sessionEntryToContextMessages,
+} from "@earendil-works/pi-coding-agent";
+import {
+	buildReplacementHistory,
+	type CodexCheckpointDetails,
+	checkpointMarker,
+	createCheckpointDetails,
+	fallbackSummary,
+	latestCheckpoint,
+	projectCheckpointContext,
+} from "./checkpoint.js";
+import { hasCheckpointMarker, rewriteCheckpointMarker } from "./protocol.js";
+import { requestRemoteCompaction } from "./remote.js";
+import {
+	type CodexCompactSettings,
+	type CodexCompactSettingsRuntime,
+	type CodexCompactSettingsState,
+	createCodexCompactSettingsRuntime,
+} from "./settings.js";
+import { showCodexCompactMenu } from "./settings-menu.js";
+
+const STATUS_KEY = "codex-compact";
+const EXPERIMENTAL_WARNING =
+	"Experimental: Codex Remote Compaction V2 uses an opaque, provider-specific checkpoint. Sessions require this extension and openai-codex for full replay.";
+
+function isSupportedModel(model: Model<Api> | undefined): model is Model<"openai-codex-responses"> {
+	return model?.provider === "openai-codex" && hasApi(model, "openai-codex-responses");
+}
+
+function activeCheckpoint(ctx: ExtensionContext) {
+	return latestCheckpoint(ctx.sessionManager.getBranch());
+}
+
+function isCheckpointCompatible(
+	details: CodexCheckpointDetails,
+	model: Model<Api> | undefined,
+): model is Model<"openai-codex-responses"> {
+	return isSupportedModel(model) && model.id === details.modelId;
+}
+
+function keptMessages(event: SessionBeforeCompactEvent): AgentMessage[] {
+	const leafId = event.branchEntries.at(-1)?.id ?? null;
+	const contextEntries = buildContextEntries(event.branchEntries, leafId);
+	const keptIndex = contextEntries.findIndex(
+		(entry) => entry.id === event.preparation.firstKeptEntryId,
+	);
+	if (keptIndex < 0) {
+		throw new Error("Pi compaction cut point is not present in the active context");
+	}
+	return contextEntries.slice(keptIndex).flatMap(sessionEntryToContextMessages);
+}
+
+function activeTools(pi: ExtensionAPI): Tool[] {
+	const enabled = new Set(pi.getActiveTools());
+	return pi
+		.getAllTools()
+		.filter((tool) => enabled.has(tool.name))
+		.map((tool) => ({
+			name: tool.name,
+			description: tool.description,
+			parameters: tool.parameters,
+		}));
+}
+
+function projectedCurrentMessages(
+	event: SessionBeforeCompactEvent,
+	model: Model<"openai-codex-responses">,
+): { messages: AgentMessage[]; prior?: CodexCheckpointDetails } {
+	const leafId = event.branchEntries.at(-1)?.id ?? null;
+	const session = buildSessionContext(event.branchEntries, leafId);
+	const prior = latestCheckpoint(event.branchEntries)?.details;
+	if (!prior) return { messages: session.messages };
+	if (prior.modelId !== model.id) {
+		throw new Error("The active opaque checkpoint belongs to a different Codex model");
+	}
+	const projected = projectCheckpointContext(session.messages, prior);
+	if (!projected) {
+		throw new Error("The previous opaque checkpoint could not be projected safely");
+	}
+	return { messages: projected, prior };
+}
+
+function notifyFailure(
+	ctx: ExtensionContext,
+	error: unknown,
+	settings: CodexCompactSettings,
+): void {
+	if (!ctx.hasUI || !settings.notifyOnFallback) return;
+	const message = error instanceof Error ? error.message : String(error);
+	ctx.ui.notify(`Codex remote compaction failed; using Pi compaction. ${message}`, "warning");
+}
+
+function sessionStillOwned(ctx: ExtensionContext, sessionId: string, signal: AbortSignal): boolean {
+	return !signal.aborted && ctx.sessionManager.getSessionId() === sessionId;
+}
+
+async function compactRemotely(
+	pi: ExtensionAPI,
+	event: SessionBeforeCompactEvent,
+	ctx: ExtensionContext,
+	settings: CodexCompactSettings,
+	fetch?: typeof globalThis.fetch,
+) {
+	const model = ctx.model;
+	if (!settings.enabled || !isSupportedModel(model)) return undefined;
+	const sessionId = ctx.sessionManager.getSessionId();
+	ctx.ui.setStatus(STATUS_KEY, "Codex remote compaction…");
+	try {
+		const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
+		if (!sessionStillOwned(ctx, sessionId, event.signal)) return { cancel: true };
+		if (!auth.ok || !auth.apiKey) {
+			throw new Error(auth.ok ? "OpenAI Codex OAuth token is unavailable" : auth.error);
+		}
+		const provider = ctx.modelRegistry.getProvider(model.provider);
+		if (!provider) throw new Error("OpenAI Codex provider is unavailable");
+		const current = projectedCurrentMessages(event, model);
+		const context: Context = {
+			systemPrompt: ctx.getSystemPrompt(),
+			messages: convertToLlm(current.messages),
+			tools: activeTools(pi),
+		};
+		const response = await requestRemoteCompaction({
+			provider,
+			model,
+			context,
+			apiKey: auth.apiKey,
+			headers: auth.headers,
+			env: auth.env,
+			signal: event.signal,
+			priorCheckpoint: current.prior
+				? {
+						marker: checkpointMarker(current.prior.checkpointId),
+						replacementHistory: current.prior.replacementHistory,
+					}
+				: undefined,
+			requestTimeoutMs: settings.requestTimeoutMs,
+			maxRetries: settings.maxRetries,
+			fetch,
+		});
+		if (!sessionStillOwned(ctx, sessionId, event.signal)) return { cancel: true };
+		const replacementHistory = buildReplacementHistory(response.promptInput, response.item, {
+			tokenBudget: settings.replacementTokenBudget,
+		});
+		const details = createCheckpointDetails({
+			modelId: model.id,
+			replacementHistory,
+			keptMessages: keptMessages(event),
+		});
+		return {
+			compaction: {
+				summary: fallbackSummary(details.checkpointId),
+				firstKeptEntryId: event.preparation.firstKeptEntryId,
+				tokensBefore: event.preparation.tokensBefore,
+				usage: response.usage,
+				details,
+			},
+		};
+	} catch (error) {
+		if (event.signal.aborted || ctx.sessionManager.getSessionId() !== sessionId) {
+			return { cancel: true };
+		}
+		notifyFailure(ctx, error, settings);
+		return undefined;
+	} finally {
+		if (ctx.sessionManager.getSessionId() === sessionId) ctx.ui.setStatus(STATUS_KEY, undefined);
+	}
+}
+
+export function createCodexCompactExtension(
+	options: { fetch?: typeof globalThis.fetch; settingsRuntime?: CodexCompactSettingsRuntime } = {},
+): (pi: ExtensionAPI) => void {
+	return (pi) => {
+		const providerWarnings = new Set<string>();
+		const settingsRuntime = options.settingsRuntime ?? createCodexCompactSettingsRuntime();
+		let sessionController = new AbortController();
+		let generation = 0;
+
+		pi.registerCommand("codex-compact", {
+			description: "Compact now or configure experimental Codex Remote Compaction V2",
+			handler: async (_args, ctx) => {
+				const ownerGeneration = generation;
+				await showCodexCompactMenu(settingsRuntime, ctx, {
+					signal: sessionController.signal,
+					isCurrent: () => ownerGeneration === generation && !sessionController.signal.aborted,
+				});
+			},
+		});
+
+		pi.on("session_start", async (_event, ctx) => {
+			sessionController.abort();
+			sessionController = new AbortController();
+			generation += 1;
+			const ownerGeneration = generation;
+			const sessionId = ctx.sessionManager.getSessionId();
+			providerWarnings.clear();
+			let state: Readonly<CodexCompactSettingsState>;
+			try {
+				state = await settingsRuntime.reload(sessionController.signal);
+			} catch (error) {
+				if (sessionController.signal.aborted || ownerGeneration !== generation) return;
+				if (ctx.hasUI) {
+					ctx.ui.notify(
+						`Could not load pi-codex-compact.json; using defaults. ${error instanceof Error ? error.message : String(error)}`,
+						"warning",
+					);
+				}
+				return;
+			}
+			if (
+				sessionController.signal.aborted ||
+				ownerGeneration !== generation ||
+				ctx.sessionManager.getSessionId() !== sessionId
+			) {
+				return;
+			}
+			if (ctx.hasUI) {
+				ctx.ui.notify(EXPERIMENTAL_WARNING, "warning");
+				if (state.kind === "invalid") {
+					ctx.ui.notify(
+						`Invalid pi-codex-compact.json; using defaults without overwriting it. ${state.issue}`,
+						"warning",
+					);
+				}
+			}
+		});
+
+		pi.on("session_before_compact", (event, ctx) =>
+			compactRemotely(pi, event, ctx, settingsRuntime.get().settings, options.fetch),
+		);
+
+		pi.on("context", (event, ctx) => {
+			if (!settingsRuntime.get().settings.enabled) return undefined;
+			const checkpoint = activeCheckpoint(ctx);
+			if (!checkpoint || !isCheckpointCompatible(checkpoint.details, ctx.model)) return undefined;
+			const messages = projectCheckpointContext(event.messages, checkpoint.details);
+			return messages ? { messages } : undefined;
+		});
+
+		pi.on("before_provider_request", (event, ctx) => {
+			if (!settingsRuntime.get().settings.enabled) return undefined;
+			const checkpoint = activeCheckpoint(ctx);
+			if (!checkpoint || !isCheckpointCompatible(checkpoint.details, ctx.model)) return undefined;
+			const marker = checkpointMarker(checkpoint.details.checkpointId);
+			if (!hasCheckpointMarker(event.payload, marker)) return undefined;
+			return rewriteCheckpointMarker(event.payload, marker, checkpoint.details.replacementHistory);
+		});
+
+		pi.on("model_select", (event, ctx) => {
+			if (!settingsRuntime.get().settings.enabled) return;
+			const checkpoint = activeCheckpoint(ctx);
+			if (!checkpoint || isCheckpointCompatible(checkpoint.details, event.model)) return;
+			const key = `${ctx.sessionManager.getSessionId()}:${event.model.provider}:${event.model.id}`;
+			if (providerWarnings.has(key)) return;
+			providerWarnings.add(key);
+			if (ctx.hasUI) {
+				ctx.ui.notify(
+					"The active Codex checkpoint cannot replay on this model; Pi will expose only its fallback marker and retained recent messages.",
+					"warning",
+				);
+			}
+		});
+
+		pi.on("session_shutdown", async (_event, ctx) => {
+			generation += 1;
+			sessionController.abort();
+			providerWarnings.clear();
+			ctx.ui.setStatus(STATUS_KEY, undefined);
+			await settingsRuntime.flush();
+		});
+	};
+}
+
+export default createCodexCompactExtension();

package/src/index.ts

@@ -0,0 +1 @@
+export { default } from "./codex-compact.js";

package/src/protocol.ts

@@ -0,0 +1,224 @@
+export const MAX_SSE_BYTES = 8 * 1024 * 1024;
+export const MAX_COMPACTION_ITEM_BYTES = 2 * 1024 * 1024;
+
+export type JsonObject = Record<string, unknown>;
+
+export class CodexCompactionProtocolError extends Error {
+	constructor(message: string) {
+		super(message);
+		this.name = "CodexCompactionProtocolError";
+	}
+}
+
+function isObject(value: unknown): value is JsonObject {
+	return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+function byteLength(value: unknown): number {
+	return Buffer.byteLength(JSON.stringify(value), "utf8");
+}
+
+export function isCompactionItem(value: unknown): value is JsonObject {
+	return (
+		isObject(value) &&
+		value.type === "compaction" &&
+		typeof value.encrypted_content === "string" &&
+		value.encrypted_content.length > 0
+	);
+}
+
+export function validateCompactionItem(
+	value: unknown,
+	maxBytes = MAX_COMPACTION_ITEM_BYTES,
+): JsonObject {
+	if (!isCompactionItem(value)) {
+		throw new CodexCompactionProtocolError(
+			"Remote response did not contain a valid compaction item",
+		);
+	}
+	if (byteLength(value) > maxBytes) {
+		throw new CodexCompactionProtocolError("Remote compaction item exceeded the size limit");
+	}
+	return structuredClone(value);
+}
+
+export interface CollectedCompaction {
+	item: JsonObject;
+	completedResponse?: JsonObject;
+}
+
+function compactionItemsFromEvent(event: JsonObject): unknown[] {
+	const items: unknown[] = [];
+	if (event.type === "response.output_item.done" && isObject(event.item)) {
+		items.push(event.item);
+	}
+	if (event.type === "response.completed" && isObject(event.response)) {
+		const output = event.response.output;
+		if (Array.isArray(output)) items.push(...output);
+	}
+	return items.filter((item) => isObject(item) && item.type === "compaction");
+}
+
+export async function collectCompactionSse(
+	stream: ReadableStream<Uint8Array>,
+	options: {
+		signal?: AbortSignal;
+		maxBytes?: number;
+		maxItemBytes?: number;
+	} = {},
+): Promise<CollectedCompaction> {
+	const maxBytes = options.maxBytes ?? MAX_SSE_BYTES;
+	const reader = stream.getReader();
+	const onAbort = () => {
+		void reader.cancel(new DOMException("Compaction aborted", "AbortError")).catch(() => undefined);
+	};
+	options.signal?.addEventListener("abort", onAbort, { once: true });
+	const decoder = new TextDecoder();
+	let bytes = 0;
+	let pending = "";
+	let dataLines: string[] = [];
+	let completedResponse: JsonObject | undefined;
+	const items = new Map<string, JsonObject>();
+
+	const checkAbort = () => {
+		if (options.signal?.aborted) throw new DOMException("Compaction aborted", "AbortError");
+	};
+	const dispatch = () => {
+		if (dataLines.length === 0) return;
+		const data = dataLines.join("\n");
+		dataLines = [];
+		if (data === "[DONE]") return;
+		let parsed: unknown;
+		try {
+			parsed = JSON.parse(data);
+		} catch {
+			throw new CodexCompactionProtocolError("Remote compaction returned malformed SSE JSON");
+		}
+		if (!isObject(parsed)) return;
+		if (parsed.type === "response.completed") {
+			completedResponse = isObject(parsed.response) ? parsed.response : {};
+		}
+		for (const candidate of compactionItemsFromEvent(parsed)) {
+			const item = validateCompactionItem(candidate, options.maxItemBytes);
+			items.set(JSON.stringify(item), item);
+		}
+	};
+	const processLine = (line: string) => {
+		if (line === "") {
+			dispatch();
+			return;
+		}
+		if (line.startsWith(":")) return;
+		if (line === "data") dataLines.push("");
+		else if (line.startsWith("data:")) dataLines.push(line.slice(5).replace(/^ /, ""));
+	};
+
+	try {
+		while (true) {
+			checkAbort();
+			const { done, value } = await reader.read();
+			if (done) break;
+			bytes += value.byteLength;
+			if (bytes > maxBytes) {
+				throw new CodexCompactionProtocolError("Remote compaction stream exceeded the size limit");
+			}
+			pending += decoder.decode(value, { stream: true });
+			let newline = pending.indexOf("\n");
+			while (newline !== -1) {
+				const rawLine = pending.slice(0, newline);
+				pending = pending.slice(newline + 1);
+				processLine(rawLine.endsWith("\r") ? rawLine.slice(0, -1) : rawLine);
+				newline = pending.indexOf("\n");
+			}
+		}
+		pending += decoder.decode();
+		if (pending.length > 0) processLine(pending.endsWith("\r") ? pending.slice(0, -1) : pending);
+		dispatch();
+		checkAbort();
+	} catch (error) {
+		await reader.cancel(error).catch(() => undefined);
+		throw error;
+	} finally {
+		options.signal?.removeEventListener("abort", onAbort);
+		reader.releaseLock();
+	}
+
+	if (!completedResponse) {
+		throw new CodexCompactionProtocolError(
+			"Remote compaction stream ended without response.completed",
+		);
+	}
+	if (items.size !== 1) {
+		throw new CodexCompactionProtocolError(
+			`Remote compaction returned ${items.size} distinct compaction items; expected exactly one`,
+		);
+	}
+	return { item: [...items.values()][0], completedResponse };
+}
+
+function markerTextFromItem(item: unknown): string | undefined {
+	if (!isObject(item) || item.role !== "user" || !Array.isArray(item.content)) return undefined;
+	if (item.content.length !== 1) return undefined;
+	const content = item.content[0];
+	if (!isObject(content) || content.type !== "input_text" || typeof content.text !== "string") {
+		return undefined;
+	}
+	return content.text;
+}
+
+export function rewriteCheckpointMarker(
+	payload: unknown,
+	marker: string,
+	replacementHistory: readonly unknown[],
+): JsonObject {
+	if (!isObject(payload) || !Array.isArray(payload.input)) {
+		throw new CodexCompactionProtocolError("OpenAI Codex payload is missing an input array");
+	}
+	const matches = payload.input
+		.map((item, index) => (markerTextFromItem(item) === marker ? index : -1))
+		.filter((index) => index >= 0);
+	if (matches.length !== 1) {
+		throw new CodexCompactionProtocolError(
+			`Provider payload contained ${matches.length} checkpoint markers; expected exactly one`,
+		);
+	}
+	const index = matches[0];
+	return {
+		...payload,
+		input: [
+			...payload.input.slice(0, index),
+			...structuredClone(replacementHistory),
+			...payload.input.slice(index + 1),
+		],
+	};
+}
+
+export function appendCompactionTrigger(payload: unknown): JsonObject {
+	if (!isObject(payload) || !Array.isArray(payload.input)) {
+		throw new CodexCompactionProtocolError("OpenAI Codex payload is missing an input array");
+	}
+	if (payload.input.some((item) => isObject(item) && item.type === "compaction_trigger")) {
+		throw new CodexCompactionProtocolError(
+			"Provider payload already contains a compaction trigger",
+		);
+	}
+	return { ...payload, input: [...payload.input, { type: "compaction_trigger" }] };
+}
+
+export function prepareRemoteCompactionPayload(
+	payload: unknown,
+	checkpoint?: { marker: string; replacementHistory: readonly unknown[] },
+): JsonObject {
+	const expanded = checkpoint
+		? rewriteCheckpointMarker(payload, checkpoint.marker, checkpoint.replacementHistory)
+		: payload;
+	return appendCompactionTrigger(expanded);
+}
+
+export function hasCheckpointMarker(payload: unknown, marker: string): boolean {
+	return (
+		isObject(payload) &&
+		Array.isArray(payload.input) &&
+		payload.input.some((item) => markerTextFromItem(item) === marker)
+	);
+}

package/src/remote.ts

@@ -0,0 +1,117 @@
+import type { Context, Model, Provider, Usage } from "@earendil-works/pi-ai";
+import {
+	CodexCompactionProtocolError,
+	type CollectedCompaction,
+	collectCompactionSse,
+	type JsonObject,
+	prepareRemoteCompactionPayload,
+} from "./protocol.js";
+
+interface PriorCheckpointPayload {
+	marker: string;
+	replacementHistory: readonly unknown[];
+}
+
+export interface RemoteCompactionRequest {
+	provider: Provider;
+	model: Model<"openai-codex-responses">;
+	context: Context;
+	apiKey?: string;
+	headers?: Record<string, string>;
+	env?: Record<string, string>;
+	signal: AbortSignal;
+	priorCheckpoint?: PriorCheckpointPayload;
+	requestTimeoutMs?: number;
+	maxRetries?: number;
+	fetch?: typeof globalThis.fetch;
+}
+
+export interface RemoteCompactionResponse {
+	item: JsonObject;
+	promptInput: JsonObject[];
+	usage: Usage;
+}
+
+const EMPTY_USAGE: Usage = {
+	input: 0,
+	output: 0,
+	cacheRead: 0,
+	cacheWrite: 0,
+	totalTokens: 0,
+	cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+};
+
+function isObject(value: unknown): value is JsonObject {
+	return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+function abortError(): DOMException {
+	return new DOMException("Compaction aborted", "AbortError");
+}
+
+export async function requestRemoteCompaction(
+	request: RemoteCompactionRequest,
+): Promise<RemoteCompactionResponse> {
+	if (request.signal.aborted) throw abortError();
+	let sentInput: JsonObject[] | undefined;
+	const inspections: Promise<
+		{ ok: true; value: CollectedCompaction } | { ok: false; error: unknown }
+	>[] = [];
+	const baseFetch = request.fetch ?? globalThis.fetch;
+	const inspectedFetch: typeof globalThis.fetch = async (input, init) => {
+		const response = await baseFetch(input, init);
+		if (!response.ok || !response.body) return response;
+		const [providerBody, inspectionBody] = response.body.tee();
+		const inspection = collectCompactionSse(inspectionBody, { signal: request.signal }).then(
+			(value) => ({ ok: true as const, value }),
+			(error: unknown) => ({ ok: false as const, error }),
+		);
+		inspections.push(inspection);
+		return new Response(providerBody, {
+			status: response.status,
+			statusText: response.statusText,
+			headers: response.headers,
+		});
+	};
+
+	const stream = request.provider.stream(request.model, request.context, {
+		apiKey: request.apiKey,
+		headers: request.headers,
+		env: request.env,
+		signal: request.signal,
+		transport: "sse",
+		cacheRetention: "none",
+		timeoutMs: request.requestTimeoutMs ?? 5 * 60 * 1000,
+		maxRetries: request.maxRetries ?? 2,
+		fetch: inspectedFetch,
+		onPayload: (payload) => {
+			const prepared = prepareRemoteCompactionPayload(payload, request.priorCheckpoint);
+			if (!Array.isArray(prepared.input) || !prepared.input.every(isObject)) {
+				throw new CodexCompactionProtocolError(
+					"Prepared compaction payload has invalid input items",
+				);
+			}
+			sentInput = structuredClone(prepared.input.slice(0, -1)) as JsonObject[];
+			return prepared;
+		},
+	});
+
+	let usage = EMPTY_USAGE;
+	for await (const event of stream) {
+		if (request.signal.aborted) throw abortError();
+		if (event.type === "error") {
+			throw new Error(event.error.errorMessage ?? "OpenAI Codex compaction request failed");
+		}
+		if (event.type === "done") usage = event.message.usage;
+	}
+	if (request.signal.aborted) throw abortError();
+	if (!sentInput)
+		throw new CodexCompactionProtocolError("Provider did not expose a request payload");
+	if (inspections.length === 0) {
+		throw new CodexCompactionProtocolError("Provider response did not expose an SSE body");
+	}
+	const inspection = await inspections.at(-1);
+	if (request.signal.aborted) throw abortError();
+	if (!inspection?.ok) throw inspection?.error ?? new Error("Remote compaction inspection failed");
+	return { item: inspection.value.item, promptInput: sentInput, usage };
+}