pi-openai-toolkit 0.11.1 → 0.13.0

package/src/extension-runtime.ts

@@ -1,21 +1,31 @@
 import type {
 	BeforeProviderRequestEvent,
+	ContextEvent,
 	CompactionResult,
 	ExtensionAPI,
 	ExtensionContext,
 	SessionBeforeCompactEvent,
 } from "@earendil-works/pi-coding-agent";
 import { loadToolkitConfig } from "./config";
+import {
+	codexContextProviderHeaders,
+	resolveCodexContextProvider,
+	isCodexGatewayModel,
+	isNativeCodexModel,
+} from "./context-management/codex-provider";
+import { routeContextNamespaceToolMessage } from "./context-management/namespace-tools";
+import { CodexContextWindowManager } from "./context-management/window-manager";
+import { registerContextManagementTools } from "./context-management/tools";
 import { writeDebugArtifact, writeReplayFailureArtifact } from "./debug";
 import { resolveLatestNativeCompactionEntry } from "./details-store";
 import { runNativeFallbackCompaction } from "./native-fallback";
 import {
 	rewriteResponsesPayloadWithNativeReplay,
+	removeNativeCompactionRetainedMessages,
 	serializeLiveTailToResponsesInput,
 } from "./payload-rewrite";
 import { getCompactionRequestExtras, rememberRequestContext } from "./request-context-cache";
 import { executeRemoteV2Compaction } from "./remote-v2-client";
-import { registerInlineCompactionRuntime } from "./inline-compaction";
 import {
 	resolveNativeCompactionEnvironment,
 	resolveRemoteCompactionExecution,
@@ -33,6 +43,19 @@ import {
 	type NativeCompactionRequestMeta,
 } from "./types";
 
+type CompactionDependencies = {
+	loadConfig: typeof loadToolkitConfig;
+	remoteCompact: typeof executeRemoteV2Compaction;
+	nativeFallback: typeof runNativeFallbackCompaction;
+	contextWindows: CodexContextWindowManager;
+};
+
+type RemoteContextActive = (
+	ctx: ExtensionContext,
+	config: CompactionConfig,
+	model?: ExtensionContext["model"],
+) => Promise<boolean>;
+
 type ResponsesCompactOutcome =
 	| { outcome: "success"; compaction: CompactionResult<NativeCompactionDetails> }
 	| { outcome: "aborted" }
@@ -80,6 +103,29 @@ function notifyWarning(ctx: ExtensionContext, message: string): void {
 	}
 }
 
+async function isRemoteContextActive(
+	ctx: ExtensionContext,
+	config: CompactionConfig,
+	model: ExtensionContext["model"] = ctx.model,
+): Promise<boolean> {
+	if (!config.enabled || config.contextManagement !== "remote") return false;
+	const resolution = await resolveCodexContextProvider(ctx, model, config.codexGatewayModels);
+	return resolution.ok;
+}
+
+function isCodexContextModel(
+	model: ExtensionContext["model"] | undefined,
+	config: CompactionConfig,
+): boolean {
+	return isNativeCodexModel(model) || isCodexGatewayModel(model, config.codexGatewayModels);
+}
+
+function notifyRemoteContextFailure(ctx: ExtensionContext, reason: string): void {
+	if (ctx.hasUI) {
+		ctx.ui.notify(`${COMPACTION_EXTENSION_ID}: Remote Context management inactive (${reason})`, "warning");
+	}
+}
+
 function cloneOpaqueWindow(window: readonly unknown[]): unknown[] {
 	return window.map((item) => structuredClone(item));
 }
@@ -98,6 +144,7 @@ async function runResponsesNativeCompact(
 	ctx: ExtensionContext,
 	config: CompactionConfig,
 	execution: RemoteCompactionExecution,
+	remoteCompact: typeof executeRemoteV2Compaction,
 ): Promise<ResponsesCompactOutcome> {
 	const { consumer, compactor } = execution;
 	const instructions = buildCompactionInstructions(ctx.getSystemPrompt(), event.customInstructions);
@@ -179,12 +226,12 @@ async function runResponsesNativeCompact(
 		model: consumer.model,
 		baseUrl: consumer.baseUrl,
 		sessionId: getSessionId(ctx),
-	});
+	}, compactor.currentModel);
 	if (extras) {
 		request = { ...request, ...extras };
 	}
 
-	const compactResult = await executeRemoteV2Compaction({
+	const compactResult = await remoteCompact({
 		runtime: compactor,
 		request,
 		signal: event.signal,
@@ -287,8 +334,13 @@ async function runResponsesNativeCompact(
 	return { outcome: "success", compaction };
 }
 
-async function handleSessionBeforeCompact(event: SessionBeforeCompactEvent, ctx: ExtensionContext) {
-	const { config: toolkitConfig } = loadToolkitConfig();
+async function handleSessionBeforeCompact(
+	event: SessionBeforeCompactEvent,
+	ctx: ExtensionContext,
+	dependencies: CompactionDependencies,
+	remoteContextActive: RemoteContextActive,
+) {
+	const { config: toolkitConfig } = dependencies.loadConfig();
 	const config = toolkitConfig.compaction;
 	if (!config.enabled) {
 		return undefined;
@@ -315,6 +367,28 @@ async function handleSessionBeforeCompact(event: SessionBeforeCompactEvent, ctx:
 		return { cancel: true };
 	}
 
+	// Remote Context management owns this eligible Codex session. It persists a
+	// no-summary boundary and deliberately never calls remote_compaction_v2.
+	if (config.contextManagement === "remote" && isCodexContextModel(ctx.model, config)) {
+		if (await remoteContextActive(ctx, config)) {
+			try {
+				dependencies.contextWindows.synchronize(ctx);
+				return dependencies.contextWindows.prepareCompaction(event);
+			} catch {
+				notifyRemoteContextFailure(ctx, "malformed-window-state");
+				return { cancel: true };
+			}
+		}
+		// A configured native Codex Remote session must never re-enter this
+		// extension's Remote V2 compaction chain, and Pi's native compaction is
+		// deliberately disabled for Remote-managed models: cancelling here keeps
+		// the boundary the only rollover mechanism. The inactive reason is
+		// surfaced so the user can fix activation instead of losing context to a
+		// silent native summary.
+		notifyRemoteContextFailure(ctx, "native-codex-context-unavailable");
+		return { cancel: true };
+	}
+
 	// Branch 1: Responses-family APIs use remote_compaction_v2 on the normal Responses stream.
 	let remoteAttempted = false;
 	const resolution = await resolveRemoteCompactionExecution(
@@ -327,7 +401,7 @@ async function handleSessionBeforeCompact(event: SessionBeforeCompactEvent, ctx:
 	);
 	if (resolution.ok) {
 		remoteAttempted = true;
-		const responsesOutcome = await runResponsesNativeCompact(event, ctx, config, resolution.execution);
+		const responsesOutcome = await runResponsesNativeCompact(event, ctx, config, resolution.execution, dependencies.remoteCompact);
 		if (responsesOutcome.outcome === "success") {
 			return { compaction: responsesOutcome.compaction };
 		}
@@ -362,7 +436,7 @@ async function handleSessionBeforeCompact(event: SessionBeforeCompactEvent, ctx:
 	// Branch 2: run pi's native compaction method. A failed remote request is compacted by the
 	// remote producer itself; a model that cannot use remote v2 at all uses nativeFallback.model.
 	const fallbackModelSpec = remoteAttempted ? config.remoteCompactModel : config.nativeFallback.model;
-	const fallback = await runNativeFallbackCompaction({ ctx, event, config, modelSpec: fallbackModelSpec });
+	const fallback = await dependencies.nativeFallback({ ctx, event, config, modelSpec: fallbackModelSpec });
 	if (fallback.ok) {
 		if (ctx.hasUI) {
 			ctx.ui.notify(
@@ -414,13 +488,122 @@ async function handleSessionBeforeCompact(event: SessionBeforeCompactEvent, ctx:
 	return undefined;
 }
 
-async function handleBeforeProviderRequest(event: BeforeProviderRequestEvent, ctx: ExtensionContext) {
-	const { config: toolkitConfig } = loadToolkitConfig();
+async function handleContext(
+	event: ContextEvent,
+	ctx: ExtensionContext,
+	pi: ExtensionAPI,
+	loadConfig: typeof loadToolkitConfig,
+	contextWindows: CodexContextWindowManager,
+	remoteContextActive: RemoteContextActive,
+) {
+	const { config: { compaction: config } } = loadConfig();
+	if (!config.enabled) {
+		const visibleMessages = contextWindows.project(event.messages, "off");
+		return visibleMessages.length === event.messages.length && visibleMessages.every((message, index) => message === event.messages[index])
+			? undefined
+			: { messages: visibleMessages };
+	}
+
+	if (config.contextManagement === "remote") {
+		try {
+			contextWindows.synchronize(ctx);
+		} catch {
+			notifyRemoteContextFailure(ctx, "malformed-window-state");
+			ctx.abort();
+			return undefined;
+		}
+
+		const remoteActive = await remoteContextActive(ctx, config);
+		if (remoteActive) {
+			try {
+				contextWindows.recordBudget(
+					pi,
+					ctx,
+					true,
+					config.contextReminderThresholdPercent,
+				);
+				const projected = contextWindows.project(event.messages, "remote");
+				return projected.length === event.messages.length && projected.every((message, index) => message === event.messages[index])
+					? undefined
+					: { messages: projected };
+			} catch (error) {
+				notifyRemoteContextFailure(ctx, "malformed-window-state");
+				ctx.abort();
+				return undefined;
+			}
+		}
+
+		// An eligible Codex Remote session is never allowed to fall through to the
+		// older compaction-replay path. Keep its internal markers hidden, but let
+		// Pi's normal request construction handle the inactive session.
+		if (isCodexContextModel(ctx.model, config)) {
+			const visibleMessages = contextWindows.project(event.messages, "off");
+			return visibleMessages.length === event.messages.length && visibleMessages.every((message, index) => message === event.messages[index])
+				? undefined
+				: { messages: visibleMessages };
+		}
+	}
+
+	// Inactive/ineligible Remote mode must not expose internal window markers to a
+	// gateway or another provider. The existing compaction replay path remains the
+	// owner for this safe fallback.
+	const visibleMessages = contextWindows.project(event.messages, "off");
+	const replayEvent = visibleMessages === event.messages ? event : { ...event, messages: visibleMessages };
+	const resolution = await resolveNativeCompactionEnvironment(ctx, {
+		enabled: config.enabled,
+		responsesApis: config.responsesApis,
+	});
+	if (!resolution.ok) return undefined;
+	const branchEntries = ctx.sessionManager.getBranch();
+	const latest = resolveLatestNativeCompactionEntry(branchEntries, { baseUrl: resolution.runtime.baseUrl });
+	if (!latest.ok) return undefined;
+	const result = removeNativeCompactionRetainedMessages({ messages: replayEvent.messages, branchEntries, compactionEntry: latest.entry });
+	if (!result.ok) {
+		writeReplayFailureArtifact({ reason: result.reason, compactionEntryId: latest.entry.id }, config, ctx);
+		if (ctx.hasUI) ctx.ui.notify(`${COMPACTION_EXTENSION_ID}: replay failed (${result.reason}); request aborted`, "error");
+		ctx.abort();
+		return undefined;
+	}
+	return result.messages === replayEvent.messages ? undefined : { messages: result.messages };
+}
+
+async function handleBeforeProviderRequest(
+	event: BeforeProviderRequestEvent,
+	ctx: ExtensionContext,
+	loadConfig: typeof loadToolkitConfig,
+	contextWindows: CodexContextWindowManager,
+	remoteContextActive: RemoteContextActive,
+) {
+	const { config: toolkitConfig } = loadConfig();
 	const config = toolkitConfig.compaction;
 	if (!config.enabled) {
 		return undefined;
 	}
 
+	if (config.contextManagement === "remote") {
+		try {
+			contextWindows.synchronize(ctx);
+		} catch {
+			notifyRemoteContextFailure(ctx, "malformed-window-state");
+			ctx.abort();
+			return undefined;
+		}
+	}
+	if (config.contextManagement === "remote" && await remoteContextActive(ctx, config)) {
+		try {
+			return contextWindows.rewritePayload(event.payload, ctx);
+		} catch {
+			notifyRemoteContextFailure(ctx, "malformed-request-state");
+			ctx.abort();
+			return undefined;
+		}
+	}
+	if (config.contextManagement === "remote" && isCodexContextModel(ctx.model, config)) {
+		// Keep Codex Remote mutually exclusive with the legacy replay
+		// pipeline when authentication or tool ownership is unavailable.
+		return undefined;
+	}
+
 	const resolution = await resolveNativeCompactionEnvironment(
 		ctx,
 		{
@@ -570,14 +753,63 @@ async function handleBeforeProviderRequest(event: BeforeProviderRequestEvent, ct
 	return rewrite.rewrittenPayload;
 }
 
-export default function (pi: ExtensionAPI) {
-	const inlineCompaction = registerInlineCompactionRuntime(pi, loadToolkitConfig);
-
-	pi.on("session_start", (_event, ctx) => {
-		const { config: toolkitConfig, source, warnings } = loadToolkitConfig();
+export default function registerCompactionExtension(
+	pi: ExtensionAPI,
+	overrides: Partial<CompactionDependencies> = {},
+) {
+	const loadConfig = overrides.loadConfig ?? loadToolkitConfig;
+	const contextWindows = overrides.contextWindows ?? new CodexContextWindowManager(
+		undefined,
+		() => loadConfig().config.compaction.codexGatewayModels,
+	);
+	const dependencies: CompactionDependencies = {
+		loadConfig,
+		remoteCompact: executeRemoteV2Compaction,
+		nativeFallback: runNativeFallbackCompaction,
+		contextWindows,
+		...overrides,
+	};
+	let tools!: ReturnType<typeof registerContextManagementTools>;
+	tools = registerContextManagementTools(
+		pi,
+		contextWindows,
+		async (ctx) => {
+			const config = dependencies.loadConfig().config.compaction;
+			return tools.isRegistered && await isRemoteContextActive(ctx, config);
+		},
+		() => dependencies.loadConfig().config.compaction.codexGatewayModels,
+	);
+	const remoteContextActive: RemoteContextActive = async (ctx, config, model = ctx.model) =>
+		tools.isRegistered && await isRemoteContextActive(ctx, config, model);
+	const syncTools = async (ctx: ExtensionContext, model = ctx.model): Promise<boolean> => {
+		const config = dependencies.loadConfig().config.compaction;
+		const active = await isRemoteContextActive(ctx, config, model);
+		return tools.sync(active);
+	};
+	pi.on("session_start", async (_event, ctx) => {
+		const active = await syncTools(ctx);
+		const { config: toolkitConfig, source, warnings } = dependencies.loadConfig();
 		const config = toolkitConfig.compaction;
 		if (!config.enabled) return;
 
+		let activationReason: string | undefined;
+		if (config.contextManagement === "remote") {
+			if (active) {
+				try {
+					contextWindows.ensureInitialized(pi, ctx, true);
+				} catch {
+					notifyRemoteContextFailure(ctx, "malformed-window-state");
+				}
+			} else if (!tools.isRegistered) {
+				activationReason = "tool-name-conflict";
+				notifyRemoteContextFailure(ctx, activationReason);
+			} else {
+				const remoteResolution = await resolveCodexContextProvider(ctx, ctx.model, config.codexGatewayModels);
+				activationReason = remoteResolution.ok ? "codex-context-unavailable" : remoteResolution.reason;
+				notifyRemoteContextFailure(ctx, activationReason);
+			}
+		}
+
 		if (warnings.length > 0 && ctx.hasUI && config.debug) {
 			ctx.ui.notify(`${COMPACTION_EXTENSION_ID}: ${warnings[0]}`, "warning");
 		}
@@ -589,6 +821,12 @@ export default function (pi: ExtensionAPI) {
 				config,
 				configSource: source,
 				warnings,
+				activation: {
+					active,
+					contextManagement: config.contextManagement,
+					model: ctx.model ? `${ctx.model.provider}/${ctx.model.id}` : undefined,
+					...(activationReason ? { reason: activationReason } : {}),
+				},
 			},
 			config,
 			ctx,
@@ -604,8 +842,57 @@ export default function (pi: ExtensionAPI) {
 		}
 	});
 
-	pi.on("session_before_compact", handleSessionBeforeCompact);
-	pi.on("before_provider_request", handleBeforeProviderRequest);
-	pi.on("turn_end", (event, ctx) => inlineCompaction.handlePublicTurnEnd(event, ctx));
-	pi.on("session_shutdown", (_event, ctx) => inlineCompaction.dispose(ctx));
+	pi.on("context", (event, ctx) => handleContext(event, ctx, pi, dependencies.loadConfig, contextWindows, remoteContextActive));
+	pi.on("session_before_compact", (event, ctx) => handleSessionBeforeCompact(event, ctx, dependencies, remoteContextActive));
+	pi.on("session_compact", (event, _ctx) => contextWindows.recordCompaction(event.compactionEntry.details));
+	pi.on("session_shutdown", () => {
+		contextWindows.reset();
+		tools.reset();
+	});
+	pi.on("model_select", async (event, ctx) => {
+		await syncTools(ctx, event.model);
+	});
+	pi.on("before_agent_start", async (_event, ctx) => {
+		await syncTools(ctx);
+	});
+	pi.on("before_provider_request", (event, ctx) => handleBeforeProviderRequest(event, ctx, dependencies.loadConfig, contextWindows, remoteContextActive));
+	pi.on("before_provider_headers", async (event, ctx) => {
+		const config = dependencies.loadConfig().config.compaction;
+		if (config.contextManagement !== "remote" || !isCodexContextModel(ctx.model, config)) return;
+		if (!(await remoteContextActive(ctx, config))) return;
+		const provider = await resolveCodexContextProvider(ctx, ctx.model, config.codexGatewayModels);
+		if (provider.ok && provider.provider.kind === "codex-gateway") {
+			const sessionId = getSessionId(ctx);
+			const gatewayHeaders = codexContextProviderHeaders(provider.provider, {
+				sessionId,
+				clientRequestId: sessionId,
+			});
+			for (const name of [
+				"authorization",
+				"originator",
+				"user-agent",
+				"version",
+				"session-id",
+				"x-client-request-id",
+				"x-codex-affinity-scope",
+				"x-codex-model",
+			]) {
+				for (const existing of Object.keys(event.headers)) {
+					if (existing.toLowerCase() === name) delete event.headers[existing];
+				}
+				const value = gatewayHeaders.get(name);
+				if (value) event.headers[name] = value;
+			}
+			for (const existing of Object.keys(event.headers)) {
+				if (["cookie", "chatgpt-account-id", "x-api-key"].includes(existing.toLowerCase())) delete event.headers[existing];
+			}
+		}
+		contextWindows.rewriteHeaders(event.headers, ctx);
+	});
+	pi.on("message_end", (event, ctx) => {
+		const config = dependencies.loadConfig().config.compaction;
+		if (config.contextManagement !== "remote" || !isCodexContextModel(ctx.model, config) || !tools.isRegistered) return undefined;
+		const message = routeContextNamespaceToolMessage(event.message);
+		return message === event.message ? undefined : { message };
+	});
 }

package/src/payload-rewrite.ts

@@ -1,3 +1,4 @@
+import { sessionEntryToContextMessages } from "@earendil-works/pi-coding-agent";
 import type { AgentMessage } from "@earendil-works/pi-agent-core";
 import type { Api, Model } from "@earendil-works/pi-ai";
 import type {
@@ -33,7 +34,8 @@ export type NativeReplayPayloadRewriteFailureReason =
 	| "first-kept-entry-not-found"
 	| "invalid-compacted-window"
 	| "unexpected-compaction-after-boundary"
-	| "compaction-summary-not-found";
+	| "compaction-summary-not-found"
+	| "retained-context-mismatch";
 
 export type NativeReplayPayloadRewriteFailure = {
 	ok: false;
@@ -311,6 +313,41 @@ export function collectLiveTailMessages(entries: readonly SessionEntry[]): Agent
 	return collectReplayMessages(entries);
 }
 
+/**
+ * Remote V2 covers the complete pre-compaction context. Pi 0.85.1 also retains
+ * recent messages, so remove those verified copies before provider conversion
+ * (where IDs/fields may change across models). Persisted session entries are
+ * untouched. Preserve transient additions and all post-compaction messages.
+ */
+export function removeNativeCompactionRetainedMessages(args: {
+	messages: AgentMessage[];
+	branchEntries: readonly SessionEntry[];
+	compactionEntry: NativeCompactionEntry;
+}): { ok: true; messages: AgentMessage[] } | NativeReplayPayloadRewriteFailure {
+	const boundary = findCompactionBoundaryIndex(args.branchEntries, args.compactionEntry.id);
+	if (boundary === undefined) return { ok: false, reason: "compaction-boundary-not-found" };
+	const firstKept = findEntryIndexByIdBeforeBoundary(args.branchEntries, args.compactionEntry.firstKeptEntryId, boundary);
+	if (firstKept === undefined) return { ok: false, reason: "first-kept-entry-not-found" };
+	const summary = sessionEntryToContextMessages(args.compactionEntry)[0];
+	const summaryIndex = args.messages.findIndex((message) => areEquivalentValues(message, summary));
+	if (summaryIndex < 0) return { ok: false, reason: "compaction-summary-not-found" };
+	const retained = args.branchEntries.slice(firstKept, boundary).flatMap(sessionEntryToContextMessages);
+	const removed = new Set<number>();
+	let cursor = summaryIndex + 1;
+	for (const expected of retained) {
+		const index = args.messages.findIndex((message, index) => index >= cursor && areEquivalentValues(message, expected));
+		if (index >= 0) {
+			removed.add(index);
+			cursor = index + 1;
+		}
+	}
+	// Already removed by another context handler: leave it alone. Partial
+	// matches are ambiguous; never remove half a call/result batch.
+	if (removed.size === 0) return { ok: true, messages: args.messages };
+	if (removed.size !== retained.length) return { ok: false, reason: "retained-context-mismatch" };
+	return { ok: true, messages: args.messages.filter((_, index) => !removed.has(index)) };
+}
+
 export function serializeLiveTailToResponsesInput<TApi extends Api>(args: {
 	model: Model<TApi>;
 	entries: readonly SessionEntry[];
@@ -370,15 +407,8 @@ function buildNativeReplaySegmentsInternal<TApi extends Api>(args: {
 		};
 	}
 
-	const preCompactionEntries = args.branchEntries.slice(firstKeptEntryIndex, boundaryIndex);
-	const postCompactionEntries = args.branchEntries.slice(boundaryIndex + 1);
-
-	// Pi-authored compaction summary item. Pi 0.84.3 renders the context after a
-	// compaction as [leading prompts][summary item][all post-compaction messages]:
-	// pre-compaction messages never appear in the payload (entry.retainedTail is
-	// never persisted by appendCompaction), so the summary item itself is the only
-	// anchor and the only thing that gets replaced. Everything else is preserved
-	// verbatim, so transient messages injected by other extensions stay intact.
+	// The context hook has removed Pi's verified retained copies. Replace only
+	// the summary anchor here; preserve provider fields and transient/post items.
 	const summaryIndex = findCompactionSummaryIndex(args.payload.input, args.compactionEntry.summary);
 	if (summaryIndex < 0) {
 		return {