diff --git a/README.md b/README.md index 38d5e81..14c6090 100644 --- a/README.md +++ b/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 diff --git a/extensions/codex-image.ts b/extensions/codex-image.ts index 4def765..1c0fa18 100644 --- a/extensions/codex-image.ts +++ b/extensions/codex-image.ts @@ -1,5 +1,6 @@ -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, @@ -7,43 +8,30 @@ import { 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>({ - type: "object", - additionalProperties: false, +const IMAGE_GENERATION_PARAMETERS = Type.Object({ + prompt: Type.String({ description: "Detailed image prompt describing what to generate." }), }); +type ImageGenerationParams = Static; + 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; 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 { + 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; + 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 { 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); - }); } diff --git a/skills/image-generation/SKILL.md b/skills/image-generation/SKILL.md index 1542f74..f52eff8 100644 --- a/skills/image-generation/SKILL.md +++ b/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. diff --git a/skills/image-generation/references/configuration.md b/skills/image-generation/references/configuration.md index ae69b5b..af00ce2 100644 --- a/skills/image-generation/references/configuration.md +++ b/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