Fix Codex image generation output handling

README.md

@@ -14,25 +14,28 @@ It provides:
 
 `pi-codex-image` dynamically routes tools based on the currently selected model:
 
-- `image_generation` is active only when the current provider is `openai-codex` and the selected model advertises image input support.
+- `image_generation` is active when the current provider is `openai-codex`.
 - `view_image` is active for any selected model that advertises image input support.
 - Switching models triggers the router again, adding these tools when supported and removing them when unsupported.
 - Existing non-image active tools are preserved while image tools are added or removed.
 
-## Native image generation
+## Image generation
 
-The `image_generation` tool mirrors `pi-codex-conversion`'s native-tool approach:
+The `image_generation` tool is exposed as a normal Pi function tool. It accepts a prompt, calls the Codex Responses image generation endpoint, extracts the returned `image_generation_call.result`, and writes the PNG locally.
 
-1. The agent sees a function-style tool named `image_generation`.
-2. Before the provider request is sent, the extension rewrites that function tool into the OpenAI Codex Responses native tool:
+Generated images are saved under:
 
-```json
-{ "type": "image_generation", "output_format": "png" }
+```text
+.pi/openai-codex-images/
 ```
 
-3. The local function body is intentionally not used; if it executes locally, it throws an explanatory error.
+The newest image is also mirrored to:
+
+```text
+.pi/openai-codex-images/latest.png
+```
 
-Generated image handling is therefore delegated to the active OpenAI Codex Responses provider, matching the referenced adapter's routing semantics.
+This local tool implementation avoids depending on Pi core support for native `image_generation_call` stream items.
 
 ## View images
 

extensions/codex-image.ts

@@ -1,49 +1,37 @@
-import { stat } from "node:fs/promises";
-import { isAbsolute, resolve } from "node:path";
+import { mkdir, readFile, stat, writeFile } from "node:fs/promises";
+import { homedir } from "node:os";
+import { isAbsolute, join, resolve } from "node:path";
 import {
 	createReadTool,
 	type AgentToolResult,
 	type ExtensionAPI,
 	type ExtensionContext,
 	type ToolDefinition,
 } from "@earendil-works/pi-coding-agent";
-import { Type, type TSchema } from "typebox";
+import { Type, type Static, type TSchema } from "typebox";
 
 const STATUS_KEY = "codex-image";
 const STATUS_TEXT = "Codex image";
 const IMAGE_GENERATION_TOOL_NAME = "image_generation";
 const VIEW_IMAGE_TOOL_NAME = "view_image";
 const CODEX_IMAGE_TOOL_NAMES = [IMAGE_GENERATION_TOOL_NAME, VIEW_IMAGE_TOOL_NAME];
-const IMAGE_GENERATION_UNSUPPORTED_MESSAGE = "image_generation is only available with image-capable openai-codex models";
-const IMAGE_GENERATION_LOCAL_EXECUTION_MESSAGE = "image_generation is a native openai-codex provider tool and should not execute locally";
+const IMAGE_GENERATION_UNSUPPORTED_MESSAGE = "image_generation is only available with openai-codex models";
 const VIEW_IMAGE_UNSUPPORTED_MESSAGE = "view_image is not allowed because the current model does not support image inputs";
 const DETAIL_DESCRIPTION = "Use `original` to preserve the file's original resolution; omit for default resized behavior.";
+const DEFAULT_CODEX_BASE_URL = "https://chatgpt.com/backend-api/codex/responses";
+const JWT_CLAIM_PATH = "https://api.openai.com/auth";
 
-const IMAGE_GENERATION_PARAMETERS = Type.Unsafe<Record<string, never>>({
-	type: "object",
-	additionalProperties: false,
+const IMAGE_GENERATION_PARAMETERS = Type.Object({
+	prompt: Type.String({ description: "Detailed image prompt describing what to generate." }),
 });
 
+type ImageGenerationParams = Static<typeof IMAGE_GENERATION_PARAMETERS>;
+
 interface ExtensionState {
 	enabled: boolean;
 	previousToolNames?: string[];
 }
 
-interface FunctionToolPayload {
-	type?: unknown;
-	name?: unknown;
-}
-
-interface ResponsesPayload {
-	tools?: unknown[];
-	[key: string]: unknown;
-}
-
-interface ResponsesImageGenerationTool {
-	type: "image_generation";
-	output_format: "png";
-}
-
 interface ViewImageParams {
 	path: string;
 	detail?: string;
@@ -58,6 +46,11 @@ interface ViewImageReaders {
 	original: ViewImageReader;
 }
 
+interface CodexAuth {
+	access?: string;
+	accountId?: string;
+}
+
 type ViewImageParameters = ReturnType<typeof createViewImageParameters>;
 
 function supportsImageInputs(model: ExtensionContext["model"]): boolean {
@@ -73,7 +66,7 @@ function isCodexImageContext(ctx: ExtensionContext): boolean {
 }
 
 function supportsNativeImageGeneration(model: ExtensionContext["model"]): boolean {
-	return isOpenAICodexModel(model) && supportsImageInputs(model);
+	return isOpenAICodexModel(model);
 }
 
 function supportsOriginalImageDetail(model: ExtensionContext["model"]): boolean {
@@ -83,39 +76,146 @@ function supportsOriginalImageDetail(model: ExtensionContext["model"]): boolean
 	return supportsImageInputs(model) && (provider.includes("codex") || api.includes("codex") || id.includes("codex"));
 }
 
-function isImageGenerationFunctionTool(tool: unknown): tool is FunctionToolPayload {
-	return !!tool && typeof tool === "object" && (tool as FunctionToolPayload).type === "function" && (tool as FunctionToolPayload).name === IMAGE_GENERATION_TOOL_NAME;
+function parseJwtAccountId(token: string): string | undefined {
+	try {
+		const [, payload] = token.split(".");
+		if (!payload) return undefined;
+		const decoded = JSON.parse(Buffer.from(payload, "base64url").toString("utf8"));
+		return decoded?.[JWT_CLAIM_PATH]?.chatgpt_account_id;
+	} catch {
+		return undefined;
+	}
 }
 
-function rewriteNativeImageGenerationTool(payload: unknown, model: ExtensionContext["model"]): unknown {
-	if (!supportsNativeImageGeneration(model) || !payload || typeof payload !== "object") return payload;
-	const tools = (payload as ResponsesPayload).tools;
-	if (!Array.isArray(tools)) return payload;
+async function readCodexAuth(): Promise<{ token: string; accountId: string }> {
+	const authPath = join(homedir(), ".pi", "agent", "auth.json");
+	const raw = await readFile(authPath, "utf8");
+	const auth = JSON.parse(raw)?.["openai-codex"] as CodexAuth | undefined;
+	const token = auth?.access;
+	if (!token) throw new Error(`Missing openai-codex OAuth access token in ${authPath}`);
+	const accountId = auth.accountId || parseJwtAccountId(token);
+	if (!accountId) throw new Error("Unable to determine ChatGPT account id for openai-codex image generation");
+	return { token, accountId };
+}
+
+function parseSseEvents(text: string): unknown[] {
+	const events: unknown[] = [];
+	let dataLines: string[] = [];
+	const flush = () => {
+		if (dataLines.length === 0) return;
+		const data = dataLines.join("\n");
+		dataLines = [];
+		if (data === "[DONE]") return;
+		try {
+			events.push(JSON.parse(data));
+		} catch {
+			// Ignore non-JSON keepalive/debug chunks.
+		}
+	};
+	for (const line of text.split(/\r?\n/)) {
+		if (line === "") {
+			flush();
+			continue;
+		}
+		if (line.startsWith("data: ")) dataLines.push(line.slice(6));
+	}
+	flush();
+	return events;
+}
 
-	let rewritten = false;
-	const nextTools = tools.map((tool) => {
-		if (!isImageGenerationFunctionTool(tool)) return tool;
-		rewritten = true;
-		const nativeTool: ResponsesImageGenerationTool = { type: "image_generation", output_format: "png" };
-		return nativeTool;
+function findGeneratedImage(events: unknown[]): string | undefined {
+	for (const event of events) {
+		if (!event || typeof event !== "object") continue;
+		const item = (event as { item?: unknown }).item;
+		if (!item || typeof item !== "object") continue;
+		const typedItem = item as { type?: unknown; result?: unknown };
+		if (typedItem.type === "image_generation_call" && typeof typedItem.result === "string" && typedItem.result.length > 0) {
+			return typedItem.result;
+		}
+	}
+	return undefined;
+}
+
+async function generateCodexImage(prompt: string, modelId: string, signal?: AbortSignal): Promise<string> {
+	const { token, accountId } = await readCodexAuth();
+	const body = {
+		model: modelId,
+		store: false,
+		stream: true,
+		instructions: "Generate the requested image. Return no extra text.",
+		input: [{ role: "user", content: [{ type: "input_text", text: prompt }] }],
+		tools: [{ type: "image_generation", output_format: "png" }],
+		tool_choice: "auto",
+	};
+
+	const response = await fetch(DEFAULT_CODEX_BASE_URL, {
+		method: "POST",
+		headers: {
+			Authorization: `Bearer ${token}`,
+			"chatgpt-account-id": accountId,
+			originator: "pi",
+			"OpenAI-Beta": "responses=experimental",
+			accept: "text/event-stream",
+			"content-type": "application/json",
+			"User-Agent": "pi-codex-image",
+		},
+		body: JSON.stringify(body),
+		signal,
 	});
 
-	return rewritten ? { ...(payload as ResponsesPayload), tools: nextTools } : payload;
+	const text = await response.text();
+	if (!response.ok) {
+		throw new Error(`Codex image generation failed (${response.status} ${response.statusText}): ${text.slice(0, 1000)}`);
+	}
+
+	const imageBase64 = findGeneratedImage(parseSseEvents(text));
+	if (!imageBase64) {
+		throw new Error("Codex image generation completed without an image_generation_call result");
+	}
+	return imageBase64;
+}
+
+async function saveGeneratedImage(cwd: string, imageBase64: string): Promise<{ path: string; latestPath: string }> {
+	const outputDir = resolve(cwd, ".pi", "openai-codex-images");
+	await mkdir(outputDir, { recursive: true });
+	const timestamp = new Date().toISOString().replace(/[:.]/g, "-");
+	const imagePath = join(outputDir, `${timestamp}.png`);
+	const latestPath = join(outputDir, "latest.png");
+	const data = Buffer.from(imageBase64, "base64");
+	await writeFile(imagePath, data);
+	await writeFile(latestPath, data);
+	return { path: imagePath, latestPath };
+}
+
+function prepareImageGenerationArguments(args: unknown): ImageGenerationParams {
+	if (args && typeof args === "object") {
+		const record = args as Record<string, unknown>;
+		const prompt = record.prompt ?? record.description ?? record.query;
+		if (typeof prompt === "string" && prompt.trim().length > 0) return { prompt: prompt.trim() };
+	}
+	return { prompt: String(args ?? "").trim() };
 }
 
 function createImageGenerationTool(): ToolDefinition<typeof IMAGE_GENERATION_PARAMETERS> {
 	const description =
-		"Generate an image. Outputs are saved under `.pi/openai-codex-images/` and mirrored to `.pi/openai-codex-images/latest.png`.";
+		"Generate a PNG image from a prompt. Outputs are saved under `.pi/openai-codex-images/` and mirrored to `.pi/openai-codex-images/latest.png`.";
 	return {
 		name: IMAGE_GENERATION_TOOL_NAME,
 		label: IMAGE_GENERATION_TOOL_NAME,
 		description,
-		promptSnippet: description,
+		promptSnippet: `${description} When the user asks to generate an image, call this tool with a concise but complete prompt.`,
 		parameters: IMAGE_GENERATION_PARAMETERS,
-		prepareArguments: () => ({}),
-		async execute(_toolCallId, _params, _signal, _onUpdate, ctx) {
+		prepareArguments: prepareImageGenerationArguments,
+		async execute(_toolCallId, params, signal, _onUpdate, ctx) {
 			if (!supportsNativeImageGeneration(ctx.model)) throw new Error(IMAGE_GENERATION_UNSUPPORTED_MESSAGE);
-			throw new Error(IMAGE_GENERATION_LOCAL_EXECUTION_MESSAGE);
+			const prompt = params.prompt.trim();
+			if (!prompt) throw new Error("image_generation requires a non-empty prompt");
+			const imageBase64 = await generateCodexImage(prompt, ctx.model.id, signal);
+			const saved = await saveGeneratedImage(ctx.cwd, imageBase64);
+			return {
+				content: [{ type: "text", text: `Generated image saved to ${saved.path}\nLatest mirror: ${saved.latestPath}` }],
+				details: saved,
+			};
 		},
 	};
 }
@@ -274,8 +374,4 @@ export default function codexImage(pi: ExtensionAPI) {
 	pi.on("model_select", async (_event, ctx) => {
 		syncCodexImageTools(pi, ctx, state);
 	});
-
-	pi.on("before_provider_request", async (event, ctx) => {
-		return rewriteNativeImageGenerationTool(event.payload, ctx.model);
-	});
 }

skills/image-generation/SKILL.md

@@ -1,6 +1,6 @@
 ---
 name: image-generation
-description: Generate images with the Codex native image_generation tool and inspect local images with view_image when the current model supports those capabilities. Use when the user asks for AI-created bitmap visuals, photos, illustrations, mockups, variants, reference-image inspection, or local image viewing.
+description: Generate images with the Codex image_generation tool and inspect local images with view_image when the current model supports those capabilities. Use when the user asks for AI-created bitmap visuals, photos, illustrations, mockups, variants, reference-image inspection, or local image viewing.
 ---
 
 # Image Generation
@@ -11,7 +11,7 @@ Use this skill when a user wants a generated bitmap image rather than a code-nat
 
 1. Use `image_generation` for agent-driven raster image generation when it is available.
 2. Use `view_image` to inspect local image files before or after generation when the user references an existing image or asks you to verify the output.
-3. `image_generation` is available only on image-capable `openai-codex` models; if it is missing, tell the user to switch to an image-capable OpenAI Codex model.
+3. `image_generation` is available on `openai-codex` models.
 4. `view_image` is available only when the current model supports image inputs.
 5. Treat visible text as a supported capability: when the user asks for a cover, poster, social graphic, title, label, or typography, include the exact requested wording in the prompt instead of removing it or adding `no text`.
 6. Ask only if a missing detail blocks success; otherwise make reasonable prompt-shaping choices.

skills/image-generation/references/configuration.md

@@ -6,21 +6,19 @@ The package provides the `image_generation` and `view_image` extension tools plu
 
 Tool availability is dynamic and follows the current selected Pi model.
 
-- `image_generation`: enabled only for `openai-codex` models that advertise image input support.
+- `image_generation`: enabled for `openai-codex` models.
 - `view_image`: enabled for models that advertise image input support.
 - `view_image.detail = "original"`: exposed only for image-capable Codex-family models.
 
 When the model changes, the extension re-runs routing and updates active tools. Non-image tools that were already active are preserved.
 
-## Native routing
+## Image generation routing
 
-`image_generation` is exposed to the agent as a normal function tool, then rewritten in `before_provider_request` to the native OpenAI Codex Responses tool:
+`image_generation` is exposed to the agent as a normal function tool with a `prompt` parameter. When executed, it calls the Codex Responses image generation endpoint, extracts the returned `image_generation_call.result`, and writes the PNG locally under `.pi/openai-codex-images/`.
 
-```json
-{ "type": "image_generation", "output_format": "png" }
-```
+The latest generated image is also mirrored to `.pi/openai-codex-images/latest.png`.
 
-The local function should not execute. If it does, the current provider/model is unsupported or the request was not rewritten.
+This avoids depending on Pi core provider parsing for native `image_generation_call` response stream items.
 
 ## Install