Pi extension for OpenAI image generation

core.test.ts

import { afterEach, describe, expect, test } from "bun:test";
import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import {
	buildImageRequest,
	codexBaseUrl,
	IMAGE_DIR,
	imageEndpoint,
	parseImageResponse,
	saveImages,
} from "./core.ts";

const temporaryDirectories: string[] = [];
const directMutation = async <T>(_path: string, operation: () => Promise<T>) => operation();

const temporaryDirectory = async () => {
	const directory = await mkdtemp(join(tmpdir(), "pi-imagegen-test-"));
	temporaryDirectories.push(directory);
	return directory;
};

afterEach(async () => {
	await Promise.all(temporaryDirectories.splice(0).map((directory) => rm(directory, { recursive: true, force: true })));
});

describe("OpenAI image generation core", () => {
	test("normalizes Codex base URLs and action endpoints", () => {
		expect(codexBaseUrl("https://chatgpt.com/backend-api/"))
			.toBe("https://chatgpt.com/backend-api/codex");
		expect(codexBaseUrl("https://example.com"))
			.toBe("https://example.com/api/codex");
		expect(codexBaseUrl("https://example.com/api/codex/responses"))
			.toBe("https://example.com/api/codex");
		expect(imageEndpoint("edit", "https://example.com/api"))
			.toBe("https://example.com/api/codex/images/edits");
	});

	test("builds generation and local-file edit requests", async () => {
		const cwd = await temporaryDirectory();
		await writeFile(join(cwd, "source.png"), Buffer.from("source image"));

		expect(await buildImageRequest({
			prompt: "a mountain",
			action: "generate",
			images: [],
			cwd,
		})).toEqual({
			prompt: "a mountain",
			model: "gpt-image-2",
			background: "auto",
			quality: "auto",
			size: "auto",
		});

		const edit = await buildImageRequest({
			prompt: "add snow",
			action: "edit",
			images: ["@source.png", "https://example.com/reference.png"],
			cwd,
		});
		expect(edit.images).toEqual([
			{ image_url: `data:image/png;base64,${Buffer.from("source image").toString("base64")}` },
			{ image_url: "https://example.com/reference.png" },
		]);
	});

	test("rejects edits without input images", async () => {
		await expect(buildImageRequest({
			prompt: "change it",
			action: "edit",
			images: [],
			cwd: "/tmp",
		})).rejects.toThrow("at least one image");
	});

	test("parses metadata and rejects malformed responses", () => {
		expect(parseImageResponse(JSON.stringify({
			data: [{ b64_json: "aW1hZ2U=" }],
			background: "opaque",
			quality: "high",
			size: "1024x1024",
		}))).toEqual({
			data: [{ b64_json: "aW1hZ2U=" }],
			background: "opaque",
			quality: "high",
			size: "1024x1024",
		});
		expect(() => parseImageResponse("not json")).toThrow("invalid JSON");
		expect(() => parseImageResponse('{"data":[]}')).toThrow("returned no images");
	});

	test("saves unique output under the supplied home directory and refreshes latest.png", async () => {
		const home = await temporaryDirectory();
		const details = await saveImages(home, {
			data: [
				{ b64_json: Buffer.from("first").toString("base64") },
				{ b64_json: Buffer.from("second").toString("base64") },
			],
			quality: "high",
		}, directMutation);

		expect(details.images).toHaveLength(2);
		expect(details.path).toMatch(/^~\/\.pi\/openai-codex-images\/ig_[a-f0-9]+\.png$/);
		expect(details.latest_path).toBe(`~/${IMAGE_DIR}/latest.png`);
		expect(details.quality).toBe("high");
		expect(await readFile(details.images[0].absolute_path, "utf8")).toBe("first");
		expect(await readFile(details.images[1].absolute_path, "utf8")).toBe("second");
		expect(await readFile(details.images[0].latest_absolute_path, "utf8")).toBe("first");
	});
});

core.ts

import { randomUUID } from "node:crypto";
import { mkdir, readFile, writeFile } from "node:fs/promises";
import { extname, join, relative, resolve } from "node:path";

export const DEFAULT_CODEX_BASE_URL = "https://chatgpt.com/backend-api";
export const IMAGE_DIR = ".pi/openai-codex-images";
export const IMAGE_MODEL = "gpt-image-2";

export type ImagegenAction = "generate" | "edit";

export type ImageRequestParams = {
	prompt: string;
	action: ImagegenAction;
	images: string[];
	cwd: string;
};

export type ImageResponse = {
	data: Array<{ b64_json: string }>;
	background?: string;
	quality?: string;
	size?: string;
};

export type SavedImage = {
	path: string;
	absolute_path: string;
	latest_path: string;
	latest_absolute_path: string;
};

export type ImagegenDetails = {
	path: string;
	latest_path: string;
	images: SavedImage[];
	background?: string;
	quality?: string;
	size?: string;
};

export type FileMutationQueue = <T>(path: string, operation: () => Promise<T>) => Promise<T>;

const MIME_TYPES: Record<string, string> = {
	".gif": "image/gif",
	".jpeg": "image/jpeg",
	".jpg": "image/jpeg",
	".png": "image/png",
	".webp": "image/webp",
};

const asObject = (value: unknown): Record<string, unknown> | undefined =>
	value !== null && typeof value === "object" && !Array.isArray(value)
		? value as Record<string, unknown>
		: undefined;

const optionalString = (value: unknown) => typeof value === "string" ? value : undefined;

const normalizedInputPath = (value: string) => value.startsWith("@") ? value.slice(1) : value;

const imageUrl = async (value: string, cwd: string) => {
	if (value.startsWith("data:image/") || value.startsWith("http://") || value.startsWith("https://")) {
		return value;
	}

	const inputPath = normalizedInputPath(value);
	const absolutePath = resolve(cwd, inputPath);
	const bytes = await readFile(absolutePath);
	const mimeType = MIME_TYPES[extname(absolutePath).toLowerCase()] ?? "application/octet-stream";
	return `data:${mimeType};base64,${bytes.toString("base64")}`;
};

export const codexBaseUrl = (configured = process.env.PI_CODEX_BASE_URL ?? DEFAULT_CODEX_BASE_URL) => {
	const normalized = configured.trim().replace(/\/+$/, "");
	if (!normalized) throw new Error("PI_CODEX_BASE_URL cannot be empty");

	try {
		const url = new URL(normalized);
		if (url.pathname === "" || url.pathname === "/") return `${normalized}/api/codex`;
	} catch {
		// Nonstandard configured endpoints are normalized by suffix below.
	}

	if (normalized.endsWith("/codex/responses")) return normalized.slice(0, -"/responses".length);
	if (normalized.endsWith("/codex")) return normalized;
	if (normalized.endsWith("/backend-api") || normalized.endsWith("/api")) return `${normalized}/codex`;
	return normalized;
};

export const imageEndpoint = (action: ImagegenAction, baseUrl?: string) =>
	`${codexBaseUrl(baseUrl)}/images/${action === "edit" ? "edits" : "generations"}`;

export const buildImageRequest = async ({ prompt, action, images, cwd }: ImageRequestParams) => {
	const common = {
		prompt,
		model: IMAGE_MODEL,
		background: "auto",
		quality: "auto",
		size: "auto",
	};
	if (action === "generate") return common;
	if (images.length === 0) throw new Error("image edit requires at least one image path or URL");

	return {
		images: await Promise.all(images.map(async (image) => ({ image_url: await imageUrl(image, cwd) }))),
		...common,
	};
};

export const parseImageResponse = (text: string): ImageResponse => {
	let parsed: unknown;
	try {
		parsed = JSON.parse(text);
	} catch {
		throw new Error("OpenAI image generation returned invalid JSON");
	}

	const payload = asObject(parsed);
	if (!payload || !Array.isArray(payload.data)) {
		throw new Error("OpenAI image generation response has no image data array");
	}
	const data = payload.data.map((item) => {
		const encoded = asObject(item)?.b64_json;
		if (typeof encoded !== "string" || !encoded) {
			throw new Error("OpenAI image generation response contains invalid image data");
		}
		return { b64_json: encoded };
	});
	if (data.length === 0) throw new Error("OpenAI image generation returned no images");
	const background = optionalString(payload.background);
	const quality = optionalString(payload.quality);
	const size = optionalString(payload.size);

	return {
		data,
		...(background ? { background } : {}),
		...(quality ? { quality } : {}),
		...(size ? { size } : {}),
	};
};

const homePath = (home: string, path: string) => join("~", relative(home, path));

export const saveImages = async (
	home: string,
	response: ImageResponse,
	queue: FileMutationQueue,
): Promise<ImagegenDetails> => {
	if (response.data.length === 0) throw new Error("OpenAI image generation returned no images");
	const outputDir = join(home, IMAGE_DIR);
	const latestAbsolutePath = join(outputDir, "latest.png");

	return queue(latestAbsolutePath, async () => {
		await mkdir(outputDir, { recursive: true });
		const images: SavedImage[] = [];

		for (const [index, image] of response.data.entries()) {
			const bytes = Buffer.from(image.b64_json, "base64");
			if (bytes.length === 0) throw new Error("OpenAI image generation returned an empty image");
			const suffix = index === 0 ? "" : `_${index + 1}`;
			const absolutePath = join(outputDir, `ig_${randomUUID().replaceAll("-", "")}${suffix}.png`);
			await writeFile(absolutePath, bytes);
			if (index === 0) await writeFile(latestAbsolutePath, bytes);
			images.push({
				path: homePath(home, absolutePath),
				absolute_path: absolutePath,
				latest_path: homePath(home, latestAbsolutePath),
				latest_absolute_path: latestAbsolutePath,
			});
		}

		return {
			path: images[0].path,
			latest_path: images[0].latest_path,
			images,
			...(response.background ? { background: response.background } : {}),
			...(response.quality ? { quality: response.quality } : {}),
			...(response.size ? { size: response.size } : {}),
		};
	});
};

index.test.ts

import { describe, expect, test } from "bun:test";
import { readFile } from "node:fs/promises";
import { join } from "node:path";

describe("openai-image-generation extension", () => {
	test("saves generated images under the user's home without imposing a fixed request timeout", async () => {
		const source = await readFile(join(import.meta.dir, "index.ts"), "utf8");

		expect(source).toContain("saveImages(homedir(), response, withFileMutationQueue)");
		expect(source).not.toContain("workspaceRoot");
		expect(source).not.toContain('pi.exec("git"');
		expect(source).not.toContain("AbortSignal.timeout");
	});
});

index.ts

import { StringEnum } from "@earendil-works/pi-ai";
import { homedir } from "node:os";
import {
	type ExtensionAPI,
	type ExtensionContext,
	type ToolDefinition,
	withFileMutationQueue,
} from "@earendil-works/pi-coding-agent";
import { Container, Image, Spacer, Text } from "@earendil-works/pi-tui";
import { Type } from "typebox";
import { withStatusCard } from "../tool-status-style/style.ts";
import {
	buildImageRequest,
	imageEndpoint,
	parseImageResponse,
	saveImages,
	type ImagegenDetails,
} from "./core.ts";

const CODEX_PROVIDER = "openai-codex";
const USER_AGENT = "pi-openai-image-generation-extension";
const ACCOUNT_ID_CLAIM = "https://api.openai.com/auth";

const ImagegenParameters = Type.Object({
	prompt: Type.String({ description: "Description of the image to generate or the edit to apply" }),
	action: Type.Optional(StringEnum(["generate", "edit"] as const, { description: "Generate a new image by default, or edit supplied images" })),
	images: Type.Optional(Type.Array(Type.String({ description: "Image path, data URL, or HTTP(S) URL used for editing" }))),
});

type CodexAuth = {
	token: string;
	accountId: string;
};

const asObject = (value: unknown): Record<string, unknown> | undefined =>
	value !== null && typeof value === "object" && !Array.isArray(value)
		? value as Record<string, unknown>
		: undefined;

const accountIdFromToken = (token: string) => {
	try {
		const payload = token.split(".")[1];
		if (!payload) return undefined;
		const claims = asObject(JSON.parse(Buffer.from(payload, "base64url").toString("utf8")));
		const account = asObject(claims?.[ACCOUNT_ID_CLAIM]);
		return typeof account?.chatgpt_account_id === "string" ? account.chatgpt_account_id : undefined;
	} catch {
		return undefined;
	}
};

const headerValue = (headers: Record<string, string> | undefined, name: string) =>
	Object.entries(headers ?? {}).find(([key]) => key.toLowerCase() === name.toLowerCase())?.[1];

const resolveCodexModel = (ctx: ExtensionContext) => {
	if (ctx.model?.provider !== CODEX_PROVIDER) {
		throw new Error("imagegen requires the active Pi model to use the openai-codex provider");
	}
	return ctx.model;
};

const resolveCodexAuth = async (ctx: ExtensionContext): Promise<CodexAuth> => {
	const model = resolveCodexModel(ctx);
	const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
	if (!auth.ok) throw new Error(`OpenAI Codex auth failed: ${auth.error}`);
	if (!auth.apiKey) throw new Error("OpenAI Codex auth did not provide a bearer token");
	const accountId = headerValue(auth.headers, "ChatGPT-Account-ID") ?? accountIdFromToken(auth.apiKey);
	if (!accountId) throw new Error("OpenAI Codex auth did not provide a ChatGPT account ID");

	return {
		token: auth.apiKey,
		accountId,
	};
};

const requestHeaders = (auth: CodexAuth) => {
	const headers = new Headers();
	headers.set("accept", "application/json");
	headers.set("authorization", `Bearer ${auth.token}`);
	headers.set("content-type", "application/json");
	headers.set("originator", "codex_cli_rs");
	headers.set("user-agent", USER_AGENT);
	headers.set("version", "0.0.0");
	headers.set("chatgpt-account-id", auth.accountId);
	return headers;
};

const generateImage = async (
	ctx: ExtensionContext,
	action: "generate" | "edit",
	body: Record<string, unknown>,
	signal?: AbortSignal,
) => {
	const auth = await resolveCodexAuth(ctx);
	const baseUrl = process.env.PI_CODEX_BASE_URL ?? ctx.model?.baseUrl;
	let response: Response;
	try {
		response = await fetch(imageEndpoint(action, baseUrl), {
			method: "POST",
			headers: requestHeaders(auth),
			body: JSON.stringify(body),
			signal,
		});
	} catch (error) {
		const cause = error instanceof Error && error.cause instanceof Error ? `: ${error.cause.message}` : "";
		throw new Error(`OpenAI image generation request failed${cause}`, { cause: error });
	}
	const text = await response.text();
	if (!response.ok) {
		const detail = text.trim().slice(0, 1_000);
		throw new Error(`OpenAI image generation failed (${response.status} ${response.statusText})${detail ? `: ${detail}` : ""}`);
	}
	return parseImageResponse(text);
};

const supportsImageInput = (ctx: ExtensionContext) =>
	!Array.isArray(ctx.model?.input) || ctx.model.input.includes("image");

const resultText = (details: ImagegenDetails) =>
	[`Generated image: ${details.path}`, `Latest: ${details.latest_path}`].join("\n");

const renderResult = (
	result: { content: Array<{ type: string; text?: string; data?: string; mimeType?: string }>; details?: ImagegenDetails },
	theme: { fg(role: string, text: string): string },
	showImages: boolean,
) => {
	const text = result.content.find((item) => item.type === "text")?.text ?? "Generated image";
	const label = new Text(theme.fg("dim", text), 0, 0);
	const images = showImages
		? result.content.filter((item) => item.type === "image" && item.data && item.mimeType)
		: [];
	if (images.length === 0) return label;

	const container = new Container();
	container.addChild(label);
	for (const [index, image] of images.entries()) {
		container.addChild(new Spacer(1));
		container.addChild(new Image(
			image.data!,
			image.mimeType!,
			{ fallbackColor: (value) => theme.fg("dim", value) },
			{ maxWidthCells: 60, filename: result.details?.images[index]?.path },
		));
	}
	return container;
};

const createImagegenTool = () => withStatusCard({
	name: "imagegen",
	label: "Image Generation",
	description: "Generate an image from a description or edit existing images with OpenAI's gpt-image-2 model. Edit inputs may be local paths, data URLs, or HTTP(S) URLs.",
	promptSnippet: "Generate new images or edit existing images with OpenAI",
	promptGuidelines: [
		"Use imagegen when the user asks to create a diagram, portrait, comic, meme, or other visual.",
		"Use imagegen with action edit when the user asks to modify an existing image; supply its path or URL in images.",
		"Call imagegen directly without reconfirming an image request unless essential information is missing.",
		"After imagegen succeeds, do not narrate the image, mention downloading it, or ask a follow-up question.",
	],
	executionMode: "sequential",
	parameters: ImagegenParameters,
	async execute(_toolCallId, params, signal, onUpdate, ctx) {
		const prompt = params.prompt.trim();
		if (!prompt) throw new Error("prompt is required");
		const action = params.action ?? "generate";
		const images = params.images ?? [];
		onUpdate?.({
			content: [{ type: "text", text: action === "edit" ? "Editing image..." : "Generating image..." }],
			details: undefined,
		});

		const body = await buildImageRequest({ prompt, action, images, cwd: ctx.cwd });
		const response = await generateImage(ctx, action, body, signal);
		if (signal?.aborted) throw new Error("imagegen aborted");
		const details = await saveImages(homedir(), response, withFileMutationQueue);
		const content: Array<
			{ type: "text"; text: string }
			| { type: "image"; data: string; mimeType: string }
		> = [{ type: "text", text: resultText(details) }];
		if (supportsImageInput(ctx)) {
			content.push(...response.data.map((image) => ({
				type: "image" as const,
				data: image.b64_json,
				mimeType: "image/png",
			})));
		}
		return { content, details };
	},
	renderCall(args, theme) {
		const action = theme.fg("dim", ` (${args.action ?? "generate"})`);
		const prompt = typeof args.prompt === "string" ? ` ${theme.fg("accent", JSON.stringify(args.prompt))}` : "";
		return new Text(theme.fg("toolTitle", theme.bold("imagegen")) + action + prompt, 0, 0);
	},
	renderResult(result, _options, theme, context) {
		return renderResult(result, theme, context.showImages);
	},
} satisfies ToolDefinition<typeof ImagegenParameters, ImagegenDetails>);

export default function openaiImageGenerationExtension(pi: ExtensionAPI) {
	pi.registerTool(createImagegenTool());
}