pi-openai-toolkit 0.14.5 → 0.14.7

package/src/extension-runtime.ts

@@ -18,6 +18,7 @@ import { loadHistoryNotesThreadHint } from "./context-management/history-notes";
 import { CodexContextWindowManager } from "./context-management/window-manager";
 import { registerContextManagementTools } from "./context-management/tools";
 import { writeDebugArtifact, writeReplayFailureArtifact } from "./debug";
+import { CODEX_GATEWAY_FORWARD_HEADERS } from "./responses-headers";
 import { resolveLatestNativeCompactionEntry } from "./details-store";
 import { runNativeFallbackCompaction } from "./native-fallback";
 import {
@@ -398,6 +399,7 @@ async function handleSessionBeforeCompact(
 		{
 			enabled: config.enabled,
 			responsesApis: config.responsesApis,
+			codexGatewayModels: config.contextManagement === "remote" ? config.gatewayContextModels : [],
 		},
 		config.remoteCompactModel,
 	);
@@ -536,9 +538,15 @@ async function handleContext(
 		}
 
 		// 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.
+		// older compaction-replay path. Gateway traffic must also fail closed when
+		// its transport capability is unavailable; otherwise Pi could send an
+		// ordinary unscoped request and lose CPA OAuth/session affinity.
 		if (isCodexContextModel(ctx.model, config)) {
+			if (isCodexGatewayModel(ctx.model, config.gatewayContextModels)) {
+				notifyRemoteContextFailure(ctx, "codex-context-unavailable");
+				ctx.abort();
+				return undefined;
+			}
 			const visibleMessages = contextWindows.project(event.messages, "off");
 			return visibleMessages.length === event.messages.length && visibleMessages.every((message, index) => message === event.messages[index])
 				? undefined
@@ -554,6 +562,7 @@ async function handleContext(
 	const resolution = await resolveNativeCompactionEnvironment(ctx, {
 		enabled: config.enabled,
 		responsesApis: config.responsesApis,
+		codexGatewayModels: config.contextManagement === "remote" ? config.gatewayContextModels : [],
 	});
 	if (!resolution.ok) return undefined;
 	const branchEntries = ctx.sessionManager.getBranch();
@@ -601,8 +610,14 @@ async function handleBeforeProviderRequest(
 		}
 	}
 	if (isCodexContextModel(ctx.model, config)) {
-		// Keep Codex Remote mutually exclusive with the legacy replay
-		// pipeline when authentication or tool ownership is unavailable.
+		// Keep Codex Remote mutually exclusive with the legacy replay pipeline
+		// when authentication or tool ownership is unavailable. Gateway traffic
+		// cannot continue as an ordinary Responses request because that would
+		// silently lose the required CPA session/account selection.
+		if (isCodexGatewayModel(ctx.model, config.gatewayContextModels)) {
+			notifyRemoteContextFailure(ctx, "codex-context-unavailable");
+			ctx.abort();
+		}
 		return undefined;
 	}
 
@@ -611,6 +626,7 @@ async function handleBeforeProviderRequest(
 		{
 			enabled: config.enabled,
 			responsesApis: config.responsesApis,
+			codexGatewayModels: config.contextManagement === "remote" ? config.gatewayContextModels : [],
 		},
 		event.payload,
 	);
@@ -877,7 +893,14 @@ export default function registerCompactionExtension(
 	pi.on("before_provider_headers", async (event, ctx) => {
 		const config = dependencies.loadConfig().config.compaction;
 		if (!isCodexContextModel(ctx.model, config)) return;
-		if (!(await remoteContextActive(ctx, config))) return;
+		const active = await remoteContextActive(ctx, config);
+		if (!active) {
+			if (isCodexGatewayModel(ctx.model, config.gatewayContextModels)) {
+				notifyRemoteContextFailure(ctx, "codex-context-unavailable");
+				ctx.abort();
+			}
+			return;
+		}
 		const provider = await resolveCodexContextProvider(ctx, ctx.model, config.gatewayContextModels);
 		if (provider.ok && provider.provider.kind === "codex-gateway") {
 			const sessionId = getSessionId(ctx);
@@ -885,6 +908,10 @@ export default function registerCompactionExtension(
 				sessionId,
 				clientRequestId: sessionId,
 			});
+			const allowedGatewayHeaders = new Set<string>(CODEX_GATEWAY_FORWARD_HEADERS);
+			for (const existing of Object.keys(event.headers)) {
+				if (!allowedGatewayHeaders.has(existing.toLowerCase())) delete event.headers[existing];
+			}
 			for (const name of [
 				"authorization",
 				"originator",
@@ -901,15 +928,13 @@ export default function registerCompactionExtension(
 				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) => {
+	pi.on("message_end", async (event, ctx) => {
 		const config = dependencies.loadConfig().config.compaction;
 		if (!isCodexContextModel(ctx.model, config) || !tools.isRegistered) return undefined;
+		if (!(await remoteContextActive(ctx, config))) return undefined;
 		const message = routeContextNamespaceToolMessage(event.message);
 		return message === event.message ? undefined : { message };
 	});

package/src/image-generation/artifacts.ts

@@ -1,237 +1,237 @@
-import { constants as fsConstants } from "node:fs";
-import {
-	link,
-	mkdir,
-	open,
-	realpath,
-	stat,
-	unlink,
-} from "node:fs/promises";
-import { randomUUID } from "node:crypto";
-import {
-	basename,
-	dirname,
-	extname,
-	isAbsolute,
-	join,
-	relative,
-	resolve,
-	sep,
-} from "node:path";
-import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
-import {
-	MAX_IMAGE_PATH_CHARS,
-	ImageGenerationError,
-	sanitizeImageDiagnostic,
-} from "./types";
-
-export type ExplicitOutputPlan = {
-	path: string;
-};
-
-function isPathInsideOrEqual(root: string, candidate: string): boolean {
-	const rel = relative(root, candidate);
-	return rel === "" || (!rel.startsWith(`..${sep}`) && rel !== ".." && !isAbsolute(rel));
-}
-
-function sanitizePathPart(value: string, fallback: string): string {
-	const sanitized = value
-		.trim()
-		.replace(/[^A-Za-z0-9._-]+/g, "-")
-		.replace(/^[.-]+|[.-]+$/g, "")
-		.slice(0, 120);
-	return sanitized || fallback;
-}
-
-async function realDirectory(path: string): Promise<string> {
-	return realpath(path).catch(() => resolve(path));
-}
-
-async function resolveWriteTarget(candidate: string): Promise<string> {
-	const suffix: string[] = [];
-	let current = resolve(candidate);
-	while (true) {
-		try {
-			const info = await stat(current);
-			if (suffix.length > 0 && !info.isDirectory()) {
-				throw new ImageGenerationError(
-					"output-path-invalid",
-					`Output parent is not a directory: ${current}`,
-				);
-			}
-			const existing = await realpath(current);
-			return resolve(existing, ...suffix);
-		} catch (error) {
-			if (error instanceof ImageGenerationError) throw error;
-			if (
-				error &&
-				typeof error === "object" &&
-				"code" in error &&
-				error.code !== "ENOENT"
-			) {
-				throw new ImageGenerationError(
-					"output-path-invalid",
-					`Unable to inspect output path: ${candidate}`,
-				);
-			}
-			const parent = dirname(current);
-			if (parent === current) {
-				throw new ImageGenerationError(
-					"output-path-invalid",
-					`Unable to resolve output path: ${candidate}`,
-				);
-			}
-			suffix.unshift(basename(current));
-			current = parent;
-		}
-	}
-}
-
-function throwIfAborted(signal: AbortSignal | undefined): void {
-	if (signal?.aborted) {
-		throw new ImageGenerationError("aborted", "Image generation was cancelled.");
-	}
-}
-
-export function getGeneratedImagesRoot(agentDir: string): string {
-	return join(resolve(agentDir), "generated-images");
-}
-
-export async function prepareExplicitOutputPath(args: {
-	rawPath?: string;
-	agentDir: string;
-	ctx: Pick<ExtensionContext, "cwd" | "hasUI" | "isProjectTrusted" | "ui">;
-	signal?: AbortSignal;
-}): Promise<ExplicitOutputPlan | undefined> {
-	if (args.rawPath === undefined) return undefined;
-	const trimmed = args.rawPath.trim();
-	if (!trimmed || trimmed.length > MAX_IMAGE_PATH_CHARS || extname(trimmed).toLowerCase() !== ".png") {
-		throw new ImageGenerationError(
-			"output-path-invalid",
-			"outputPath must be a non-empty .png file path.",
-		);
-	}
-
-	throwIfAborted(args.signal);
-	const lexical = isAbsolute(trimmed) ? resolve(trimmed) : resolve(args.ctx.cwd, trimmed);
-	const target = await resolveWriteTarget(lexical);
-	const targetInfo = await stat(target).catch(() => undefined);
-	if (targetInfo) {
-		throw new ImageGenerationError(
-			"output-path-invalid",
-			`Output path already exists and will not be overwritten: ${target}`,
-		);
-	}
-
-	const [agentRoot, projectRoot] = await Promise.all([
-		realDirectory(args.agentDir),
-		realDirectory(args.ctx.cwd),
-	]);
-	if (isPathInsideOrEqual(agentRoot, target)) return { path: target };
-
-	if (isPathInsideOrEqual(projectRoot, target)) {
-		if (!args.ctx.isProjectTrusted()) {
-			throw new ImageGenerationError(
-				"output-path-invalid",
-				"Writing generated images inside an untrusted project is not allowed.",
-			);
-		}
-		return { path: target };
-	}
-
-	if (!args.ctx.hasUI) {
-		throw new ImageGenerationError(
-			"output-path-approval-required",
-			"Writing outside the Pi agent root or trusted project requires interactive approval.",
-		);
-	}
-	const approved = await args.ctx.ui.confirm(
-		"Write generated image outside safe roots?",
-		[
-			"The generated PNG will be copied to this external path without overwriting an existing file:",
-			"",
-			target,
-			"",
-			"Continue?",
-		].join("\n"),
-		{ signal: args.signal },
-	);
-	throwIfAborted(args.signal);
-	if (!approved) {
-		throw new ImageGenerationError("output-path-declined", "External output path was declined.");
-	}
-	return { path: target };
-}
-
-async function publishBytesNoOverwrite(bytes: Uint8Array, target: string): Promise<void> {
-	const parent = dirname(target);
-	await mkdir(parent, { recursive: true });
-	const tempPath = join(parent, `.${basename(target)}.${randomUUID()}.tmp`);
-	let handle: Awaited<ReturnType<typeof open>> | undefined;
-	try {
-		handle = await open(tempPath, fsConstants.O_CREAT | fsConstants.O_EXCL | fsConstants.O_WRONLY, 0o600);
-		await handle.writeFile(bytes);
-		await handle.sync();
-		await handle.close();
-		handle = undefined;
-		await link(tempPath, target);
-	} finally {
-		await handle?.close().catch(() => undefined);
-		await unlink(tempPath).catch(() => undefined);
-	}
-}
-
-function isAlreadyExistsError(error: unknown): boolean {
-	return !!error && typeof error === "object" && "code" in error && error.code === "EEXIST";
-}
-
-export async function saveCanonicalImage(args: {
-	bytes: Uint8Array;
-	agentDir: string;
-	sessionId: string;
-	imageCallId: string;
-}): Promise<string> {
-	const root = getGeneratedImagesRoot(args.agentDir);
-	const sessionPart = sanitizePathPart(args.sessionId, "session");
-	const imagePart = sanitizePathPart(args.imageCallId, "image_generation");
-	const directory = join(root, sessionPart);
-
-	for (let index = 1; index <= 1000; index += 1) {
-		const suffix = index === 1 ? "" : `-${index}`;
-		const target = join(directory, `${imagePart}${suffix}.png`);
-		try {
-			await publishBytesNoOverwrite(args.bytes, target);
-			return target;
-		} catch (error) {
-			if (isAlreadyExistsError(error)) continue;
-			throw new ImageGenerationError(
-				"artifact-write-failed",
-				sanitizeImageDiagnostic(
-					error instanceof Error ? error.message : error,
-					"Failed to persist the generated image artifact.",
-				),
-			);
-		}
-	}
-	throw new ImageGenerationError(
-		"artifact-write-failed",
-		"Failed to reserve a unique generated image artifact path.",
-	);
-}
-
-export async function copyImageToExplicitPath(args: {
-	bytes: Uint8Array;
-	plan: ExplicitOutputPlan;
-}): Promise<void> {
-	const currentTarget = await resolveWriteTarget(args.plan.path);
-	if (currentTarget !== args.plan.path) {
-		throw new Error("Output path changed after approval.");
-	}
-	await publishBytesNoOverwrite(args.bytes, args.plan.path);
-}
-
-export const _artifactTest = {
-	isPathInsideOrEqual,
-	publishBytesNoOverwrite,
-	resolveWriteTarget,
-};
+import { constants as fsConstants } from "node:fs";
+import {
+	link,
+	mkdir,
+	open,
+	realpath,
+	stat,
+	unlink,
+} from "node:fs/promises";
+import { randomUUID } from "node:crypto";
+import {
+	basename,
+	dirname,
+	extname,
+	isAbsolute,
+	join,
+	relative,
+	resolve,
+	sep,
+} from "node:path";
+import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
+import {
+	MAX_IMAGE_PATH_CHARS,
+	ImageGenerationError,
+	sanitizeImageDiagnostic,
+} from "./types";
+
+export type ExplicitOutputPlan = {
+	path: string;
+};
+
+function isPathInsideOrEqual(root: string, candidate: string): boolean {
+	const rel = relative(root, candidate);
+	return rel === "" || (!rel.startsWith(`..${sep}`) && rel !== ".." && !isAbsolute(rel));
+}
+
+function sanitizePathPart(value: string, fallback: string): string {
+	const sanitized = value
+		.trim()
+		.replace(/[^A-Za-z0-9._-]+/g, "-")
+		.replace(/^[.-]+|[.-]+$/g, "")
+		.slice(0, 120);
+	return sanitized || fallback;
+}
+
+async function realDirectory(path: string): Promise<string> {
+	return realpath(path).catch(() => resolve(path));
+}
+
+async function resolveWriteTarget(candidate: string): Promise<string> {
+	const suffix: string[] = [];
+	let current = resolve(candidate);
+	while (true) {
+		try {
+			const info = await stat(current);
+			if (suffix.length > 0 && !info.isDirectory()) {
+				throw new ImageGenerationError(
+					"output-path-invalid",
+					`Output parent is not a directory: ${current}`,
+				);
+			}
+			const existing = await realpath(current);
+			return resolve(existing, ...suffix);
+		} catch (error) {
+			if (error instanceof ImageGenerationError) throw error;
+			if (
+				error &&
+				typeof error === "object" &&
+				"code" in error &&
+				error.code !== "ENOENT"
+			) {
+				throw new ImageGenerationError(
+					"output-path-invalid",
+					`Unable to inspect output path: ${candidate}`,
+				);
+			}
+			const parent = dirname(current);
+			if (parent === current) {
+				throw new ImageGenerationError(
+					"output-path-invalid",
+					`Unable to resolve output path: ${candidate}`,
+				);
+			}
+			suffix.unshift(basename(current));
+			current = parent;
+		}
+	}
+}
+
+function throwIfAborted(signal: AbortSignal | undefined): void {
+	if (signal?.aborted) {
+		throw new ImageGenerationError("aborted", "Image generation was cancelled.");
+	}
+}
+
+export function getGeneratedImagesRoot(agentDir: string): string {
+	return join(resolve(agentDir), "generated-images");
+}
+
+export async function prepareExplicitOutputPath(args: {
+	rawPath?: string;
+	agentDir: string;
+	ctx: Pick<ExtensionContext, "cwd" | "hasUI" | "isProjectTrusted" | "ui">;
+	signal?: AbortSignal;
+}): Promise<ExplicitOutputPlan | undefined> {
+	if (args.rawPath === undefined) return undefined;
+	const trimmed = args.rawPath.trim();
+	if (!trimmed || trimmed.length > MAX_IMAGE_PATH_CHARS || extname(trimmed).toLowerCase() !== ".png") {
+		throw new ImageGenerationError(
+			"output-path-invalid",
+			"outputPath must be a non-empty .png file path.",
+		);
+	}
+
+	throwIfAborted(args.signal);
+	const lexical = isAbsolute(trimmed) ? resolve(trimmed) : resolve(args.ctx.cwd, trimmed);
+	const target = await resolveWriteTarget(lexical);
+	const targetInfo = await stat(target).catch(() => undefined);
+	if (targetInfo) {
+		throw new ImageGenerationError(
+			"output-path-invalid",
+			`Output path already exists and will not be overwritten: ${target}`,
+		);
+	}
+
+	const [agentRoot, projectRoot] = await Promise.all([
+		realDirectory(args.agentDir),
+		realDirectory(args.ctx.cwd),
+	]);
+	if (isPathInsideOrEqual(agentRoot, target)) return { path: target };
+
+	if (isPathInsideOrEqual(projectRoot, target)) {
+		if (!args.ctx.isProjectTrusted()) {
+			throw new ImageGenerationError(
+				"output-path-invalid",
+				"Writing generated images inside an untrusted project is not allowed.",
+			);
+		}
+		return { path: target };
+	}
+
+	if (!args.ctx.hasUI) {
+		throw new ImageGenerationError(
+			"output-path-approval-required",
+			"Writing outside the Pi agent root or trusted project requires interactive approval.",
+		);
+	}
+	const approved = await args.ctx.ui.confirm(
+		"Write generated image outside safe roots?",
+		[
+			"The generated PNG will be copied to this external path without overwriting an existing file:",
+			"",
+			target,
+			"",
+			"Continue?",
+		].join("\n"),
+		{ signal: args.signal },
+	);
+	throwIfAborted(args.signal);
+	if (!approved) {
+		throw new ImageGenerationError("output-path-declined", "External output path was declined.");
+	}
+	return { path: target };
+}
+
+async function publishBytesNoOverwrite(bytes: Uint8Array, target: string): Promise<void> {
+	const parent = dirname(target);
+	await mkdir(parent, { recursive: true });
+	const tempPath = join(parent, `.${basename(target)}.${randomUUID()}.tmp`);
+	let handle: Awaited<ReturnType<typeof open>> | undefined;
+	try {
+		handle = await open(tempPath, fsConstants.O_CREAT | fsConstants.O_EXCL | fsConstants.O_WRONLY, 0o600);
+		await handle.writeFile(bytes);
+		await handle.sync();
+		await handle.close();
+		handle = undefined;
+		await link(tempPath, target);
+	} finally {
+		await handle?.close().catch(() => undefined);
+		await unlink(tempPath).catch(() => undefined);
+	}
+}
+
+function isAlreadyExistsError(error: unknown): boolean {
+	return !!error && typeof error === "object" && "code" in error && error.code === "EEXIST";
+}
+
+export async function saveCanonicalImage(args: {
+	bytes: Uint8Array;
+	agentDir: string;
+	sessionId: string;
+	imageCallId: string;
+}): Promise<string> {
+	const root = getGeneratedImagesRoot(args.agentDir);
+	const sessionPart = sanitizePathPart(args.sessionId, "session");
+	const imagePart = sanitizePathPart(args.imageCallId, "image_generation");
+	const directory = join(root, sessionPart);
+
+	for (let index = 1; index <= 1000; index += 1) {
+		const suffix = index === 1 ? "" : `-${index}`;
+		const target = join(directory, `${imagePart}${suffix}.png`);
+		try {
+			await publishBytesNoOverwrite(args.bytes, target);
+			return target;
+		} catch (error) {
+			if (isAlreadyExistsError(error)) continue;
+			throw new ImageGenerationError(
+				"artifact-write-failed",
+				sanitizeImageDiagnostic(
+					error instanceof Error ? error.message : error,
+					"Failed to persist the generated image artifact.",
+				),
+			);
+		}
+	}
+	throw new ImageGenerationError(
+		"artifact-write-failed",
+		"Failed to reserve a unique generated image artifact path.",
+	);
+}
+
+export async function copyImageToExplicitPath(args: {
+	bytes: Uint8Array;
+	plan: ExplicitOutputPlan;
+}): Promise<void> {
+	const currentTarget = await resolveWriteTarget(args.plan.path);
+	if (currentTarget !== args.plan.path) {
+		throw new Error("Output path changed after approval.");
+	}
+	await publishBytesNoOverwrite(args.bytes, args.plan.path);
+}
+
+export const _artifactTest = {
+	isPathInsideOrEqual,
+	publishBytesNoOverwrite,
+	resolveWriteTarget,
+};