From 536d37cbc7f688751f57888793b15a8d2003b5f9 Mon Sep 17 00:00:00 2001 From: Steve Date: Sun, 19 Jul 2026 06:25:54 +0000 Subject: [PATCH] Add Codex image generation tool --- packages/rig/sources/providers/codex.test.ts | 49 +++++++++++ packages/rig/sources/providers/codex.ts | 84 +++++++++++++++++++ .../providers/routeProviderThroughGym.ts | 12 ++- packages/rig/sources/providers/types.ts | 8 ++ .../createCodingAssistantAgent.test.ts | 3 + .../runtime/createCodingAssistantAgent.ts | 8 +- .../rig/sources/tools/codex/image_gen.test.ts | 45 ++++++++++ packages/rig/sources/tools/codex/image_gen.ts | 51 +++++++++++ packages/rig/sources/tools/codex/index.ts | 1 + 9 files changed, 256 insertions(+), 5 deletions(-) create mode 100644 packages/rig/sources/tools/codex/image_gen.test.ts create mode 100644 packages/rig/sources/tools/codex/image_gen.ts diff --git a/packages/rig/sources/providers/codex.test.ts b/packages/rig/sources/providers/codex.test.ts index f759af88..5a9240c5 100644 --- a/packages/rig/sources/providers/codex.test.ts +++ b/packages/rig/sources/providers/codex.test.ts @@ -16,6 +16,55 @@ afterEach(() => { }); describe("codex provider", () => { + it("generates images through the Codex backend with the upstream request shape", async () => { + let requestUrl = ""; + let requestBody: unknown; + vi.stubGlobal( + "fetch", + vi.fn().mockImplementation(async (input, init) => { + requestUrl = String(input); + requestBody = parseRequestBody(init); + return Response.json({ + data: [{ b64_json: validPng32Base64, revised_prompt: "A precise diagram" }], + }); + }), + ); + const provider = createCodexProvider({ apiKey: "codex-token" }); + + await expect(provider.generateImage?.("Draw a diagram")).resolves.toEqual({ + data: validPng32Base64, + mediaType: "image/png", + revisedPrompt: "A precise diagram", + }); + expect(requestUrl).toBe("https://chatgpt.com/backend-api/codex/images/generations"); + expect(requestBody).toEqual({ + background: "auto", + model: "gpt-image-2", + prompt: "Draw a diagram", + quality: "auto", + size: "auto", + }); + }); + + it("reports image API failures and empty output without fabricating an image", async () => { + const provider = createCodexProvider({ apiKey: "codex-token" }); + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue(new Response("denied", { status: 403 })), + ); + await expect(provider.generateImage?.("Draw it")).rejects.toThrow( + "Codex image generation failed (403): denied", + ); + + vi.stubGlobal( + "fetch", + vi.fn().mockResolvedValue(Response.json({ data: [] })), + ); + await expect(provider.generateImage?.("Draw it")).rejects.toThrow( + "Codex image generation returned no image data", + ); + }); + it("loads local authentication from CODEX_HOME", async () => { const codexHome = await mkdtemp(join(tmpdir(), "rig-codex-home-")); const accessToken = diff --git a/packages/rig/sources/providers/codex.ts b/packages/rig/sources/providers/codex.ts index b96f1ac7..1dd5d95f 100644 --- a/packages/rig/sources/providers/codex.ts +++ b/packages/rig/sources/providers/codex.ts @@ -30,8 +30,10 @@ import { createProviderQuotaCache } from "./createProviderQuotaCache.js"; import { fetchCodexProviderQuota } from "./fetchCodexProviderQuota.js"; import { getCodexAuthPath } from "./getCodexAuthPath.js"; import { unavailableProviderQuota } from "./unavailableProviderQuota.js"; +import { readCodexQuotaAuth, type CodexQuotaAuth } from "./readCodexQuotaAuth.js"; const CODEX_PROVIDER_ID = "openai-codex"; +const DEFAULT_CODEX_BASE_URL = "https://chatgpt.com/backend-api"; function toPiCodexModelId(id: string): string { return id.startsWith("openai/") ? id.slice("openai/".length) : id; @@ -75,6 +77,7 @@ export function createCodexProvider(options: CodexProviderOptions = {}): Provide } } const resolveApiKey = buildApiKeyResolver(options, authPath); + const resolveImageAuth = buildImageAuthResolver(options, authPath); const quota = createProviderQuotaCache(() => options.apiKey !== undefined || options.resolveApiKey !== undefined || @@ -93,6 +96,13 @@ export function createCodexProvider(options: CodexProviderOptions = {}): Provide models: codexModels, serviceTiers: ["fast"], quota: (quotaOptions) => quota.get(quotaOptions), + generateImage: (prompt, imageOptions) => + generateCodexImage({ + auth: resolveImageAuth(), + prompt, + ...(options.baseUrl === undefined ? {} : { baseUrl: options.baseUrl }), + ...(imageOptions?.signal === undefined ? {} : { signal: imageOptions.signal }), + }), stream(model, context, streamOptions) { const piModel = piModelById.get(toPiCodexModelId(model.id)); if (!piModel) { @@ -124,6 +134,80 @@ export function createCodexProvider(options: CodexProviderOptions = {}): Provide }); } +function buildImageAuthResolver( + options: CodexProviderOptions, + authPath: string, +): () => CodexQuotaAuth | undefined { + if (options.apiKey !== undefined) return () => ({ accessToken: options.apiKey! }); + if (options.resolveApiKey !== undefined) { + return () => { + const accessToken = options.resolveApiKey!(); + return accessToken === undefined ? undefined : { accessToken }; + }; + } + if (options.useLocalCodexAuth === false) return () => undefined; + return () => { + if (!existsSync(authPath)) return undefined; + try { + return readCodexQuotaAuth(readFileSync(authPath, "utf8")); + } catch { + return undefined; + } + }; +} + +async function generateCodexImage(options: { + auth: CodexQuotaAuth | undefined; + baseUrl?: string; + prompt: string; + signal?: AbortSignal; +}) { + if (options.auth === undefined) { + throw new Error("Codex image generation requires a usable access token."); + } + const headers = new Headers({ + authorization: `Bearer ${options.auth.accessToken}`, + "content-type": "application/json", + }); + if (options.auth.accountId !== undefined) { + headers.set("chatgpt-account-id", options.auth.accountId); + } + const configuredBase = (options.baseUrl ?? DEFAULT_CODEX_BASE_URL).replace(/\/+$/, ""); + const baseUrl = configuredBase.endsWith("/codex") ? configuredBase : `${configuredBase}/codex`; + const response = await fetch(`${baseUrl}/images/generations`, { + body: JSON.stringify({ + background: "auto", + model: "gpt-image-2", + prompt: options.prompt, + quality: "auto", + size: "auto", + }), + headers, + method: "POST", + ...(options.signal === undefined ? {} : { signal: options.signal }), + }); + if (!response.ok) { + const detail = (await response.text()).trim(); + throw new Error( + `Codex image generation failed (${response.status})${detail ? `: ${detail}` : "."}`, + ); + } + const body = (await response.json()) as { + data?: readonly { b64_json?: unknown; revised_prompt?: unknown }[]; + }; + const first = body.data?.[0]; + if (typeof first?.b64_json !== "string" || first.b64_json.length === 0) { + throw new Error("Codex image generation returned no image data."); + } + return { + data: first.b64_json, + mediaType: "image/png" as const, + ...(typeof first.revised_prompt === "string" + ? { revisedPrompt: first.revised_prompt } + : {}), + }; +} + function buildApiKeyResolver( options: CodexProviderOptions, authPath: string, diff --git a/packages/rig/sources/providers/routeProviderThroughGym.ts b/packages/rig/sources/providers/routeProviderThroughGym.ts index 18baf589..5e9107f1 100644 --- a/packages/rig/sources/providers/routeProviderThroughGym.ts +++ b/packages/rig/sources/providers/routeProviderThroughGym.ts @@ -26,7 +26,13 @@ export function routeProviderThroughGym(provider: Provider, env: NodeJS.ProcessE ...(provider.serviceTiers === undefined ? {} : { serviceTiers: provider.serviceTiers }), ...(env.RIG_GYM_TOKEN === undefined ? {} : { token: env.RIG_GYM_TOKEN }), }); - return provider.quota === undefined - ? gymProvider - : { ...gymProvider, quota: (options) => provider.quota!(options) }; + return { + ...gymProvider, + ...(provider.quota === undefined + ? {} + : { quota: (options?: { fresh?: boolean }) => provider.quota!(options) }), + ...(provider.generateImage === undefined + ? {} + : { generateImage: provider.generateImage.bind(provider) }), + }; } diff --git a/packages/rig/sources/providers/types.ts b/packages/rig/sources/providers/types.ts index d384bdf6..b5751c72 100644 --- a/packages/rig/sources/providers/types.ts +++ b/packages/rig/sources/providers/types.ts @@ -164,6 +164,7 @@ export interface Provider { imageProfile(model: Model): ProviderImageProfile; toolProfile(model: Model): ProviderToolProfile; quota?(options?: { fresh?: boolean }): Promise; + generateImage?(prompt: string, options?: { signal?: AbortSignal }): Promise; stream( model: Model, context: Context, @@ -171,6 +172,12 @@ export interface Provider { ): InferenceStream; } +export interface GeneratedImage { + data: string; + mediaType: "image/png"; + revisedPrompt?: string; +} + export type InferProviderModels = T["models"]; export type InferModel = TModels[number]; @@ -197,6 +204,7 @@ export function defineProvider(provider: { imageProfile?: (model: Model) => ProviderImageProfile; toolProfile?: (model: Model) => ProviderToolProfile; quota?: (options?: { fresh?: boolean }) => Promise; + generateImage?: (prompt: string, options?: { signal?: AbortSignal }) => Promise; stream( model: Model, context: Context, diff --git a/packages/rig/sources/runtime/createCodingAssistantAgent.test.ts b/packages/rig/sources/runtime/createCodingAssistantAgent.test.ts index 2589b1ec..f6e34b2b 100644 --- a/packages/rig/sources/runtime/createCodingAssistantAgent.test.ts +++ b/packages/rig/sources/runtime/createCodingAssistantAgent.test.ts @@ -30,6 +30,7 @@ describe("createCodingAssistantAgent", () => { expect(runtime.context.bash.cwd).toBe(cwd); expect(runtime.agent.snapshot().instructions).toContain(cwd); expect(runtime.agent.snapshot().effort).toBe("medium"); + expect(runtime.agent.tools.map((tool) => tool.name)).toContain("image_gen"); }); it("creates a Claude SDK agent for Anthropic models", () => { @@ -297,6 +298,7 @@ describe("createCodingAssistantAgent", () => { "view_image", "update_plan", "request_user_input", + "image_gen", "workflow", "wait_for_workflow", "workflow_status", @@ -405,6 +407,7 @@ describe("createCodingAssistantAgent", () => { "update_plan", "request_user_input", ]); + expect(runtime.agent.tools.map((tool) => tool.name)).not.toContain("image_gen"); }); it("uses provider-neutral tools for Bedrock Kimi and GLM models", () => { diff --git a/packages/rig/sources/runtime/createCodingAssistantAgent.ts b/packages/rig/sources/runtime/createCodingAssistantAgent.ts index 1d4e3e37..e8b73c5f 100644 --- a/packages/rig/sources/runtime/createCodingAssistantAgent.ts +++ b/packages/rig/sources/runtime/createCodingAssistantAgent.ts @@ -26,7 +26,7 @@ import { modelMoonshotKimiK3, modelOpenaiGpt56Sol } from "../providers/models.js import type { ServiceTier } from "../providers/types.js"; import { routeProviderThroughGym } from "../providers/routeProviderThroughGym.js"; import { claudeCollaborationTools } from "../tools/claude/index.js"; -import { codexCollaborationTools } from "../tools/codex/index.js"; +import { codexCollaborationTools, createCodexImageGenerationTool } from "../tools/codex/index.js"; import { grokCollaborationTools } from "../tools/grok/index.js"; import { agentTool } from "../tools/Agent.js"; import { goalTools } from "../tools/goals/index.js"; @@ -156,7 +156,11 @@ export function createCodingAssistantAgent( const usesCodexTools = toolProfile === "codex"; const usesGrokTools = toolProfile === "grok"; const usesKimiTools = toolProfile === "kimi"; - const baseTools = selectToolsForModel({ model, provider }); + const selectedBaseTools = selectToolsForModel({ model, provider }); + const baseTools = + usesCodexTools && provider.generateImage !== undefined + ? [...selectedBaseTools, createCodexImageGenerationTool(provider.generateImage)] + : selectedBaseTools; const collaborationTools = ( usesCodexTools ? codexCollaborationTools diff --git a/packages/rig/sources/tools/codex/image_gen.test.ts b/packages/rig/sources/tools/codex/image_gen.test.ts new file mode 100644 index 00000000..5a6b9a8d --- /dev/null +++ b/packages/rig/sources/tools/codex/image_gen.test.ts @@ -0,0 +1,45 @@ +import { describe, expect, it, vi } from "vitest"; + +import { createJustBashToolHarness } from "../testing/createJustBashToolHarness.js"; +import { validPng32Base64 } from "../testing/validImageFixtures.js"; +import { createCodexImageGenerationTool } from "./image_gen.js"; + +describe("codex image generation tool", () => { + it("persists and returns generated image content to the model", async () => { + const harness = createJustBashToolHarness(); + const generateImage = vi.fn(async () => ({ + data: validPng32Base64, + mediaType: "image/png" as const, + revisedPrompt: "A revised prompt", + })); + const tool = createCodexImageGenerationTool(generateImage); + + const result = await tool.execute({ prompt: "A small diagram" }, harness.context, { + toolCallId: "call/1", + }); + + expect(generateImage).toHaveBeenCalledWith("A small diagram", {}); + expect(result.path).toBe("/workspace/.rig/generated-images/call_1.png"); + expect(await harness.context.fs.readFileBuffer(result.path)).toEqual( + Buffer.from(validPng32Base64, "base64"), + ); + expect(tool.toLLM(result)).toEqual([ + { type: "text", text: `Generated image saved to ${result.path}` }, + { type: "image", data: validPng32Base64, mediaType: "image/png" }, + ]); + }); + + it("propagates generation failures without writing an output artifact", async () => { + const harness = createJustBashToolHarness(); + const tool = createCodexImageGenerationTool(async () => { + throw new Error("image service unavailable"); + }); + + await expect( + tool.execute({ prompt: "A small diagram" }, harness.context, { toolCallId: "failed" }), + ).rejects.toThrow("image service unavailable"); + await expect( + harness.context.fs.exists("/workspace/.rig/generated-images/failed.png"), + ).resolves.toBe(false); + }); +}); diff --git a/packages/rig/sources/tools/codex/image_gen.ts b/packages/rig/sources/tools/codex/image_gen.ts new file mode 100644 index 00000000..31fd56a9 --- /dev/null +++ b/packages/rig/sources/tools/codex/image_gen.ts @@ -0,0 +1,51 @@ +import { join } from "node:path"; + +import { Type } from "@sinclair/typebox"; + +import { defineTool } from "../../agent/types.js"; +import type { Provider } from "../../providers/types.js"; + +const DESCRIPTION = `Generate an image from a detailed text description. Use this when the user requests a diagram, portrait, comic, meme, or any other visual. Directly generate the image without reconfirmation unless essential details are missing.`; + +export function createCodexImageGenerationTool( + generateImage: NonNullable, +) { + return defineTool({ + name: "image_gen", + label: "image_gen", + description: DESCRIPTION, + arguments: Type.Object( + { + prompt: Type.String({ + description: "Detailed description of the image to generate.", + }), + }, + { additionalProperties: false }, + ), + returnType: Type.Object({ + data: Type.String(), + mediaType: Type.Literal("image/png"), + path: Type.String(), + revisedPrompt: Type.Optional(Type.String()), + }), + execute: async ({ prompt }, context, options) => { + const image = await generateImage( + prompt, + options.signal === undefined ? {} : { signal: options.signal }, + ); + const callId = options.toolCallId?.replaceAll(/[^a-zA-Z0-9_-]/g, "_") ?? "image"; + const directory = join(context.fs.cwd, ".rig", "generated-images"); + const path = join(directory, `${callId}.png`); + await context.fs.mkdir(directory, { recursive: true }); + await context.fs.writeFile(path, Buffer.from(image.data, "base64")); + return { ...image, path }; + }, + toLLM: ({ data, mediaType, path }) => [ + { type: "text", text: `Generated image saved to ${path}` }, + { type: "image", data, mediaType }, + ], + toUI: ({ path }) => `Generated image ${path}`, + shouldReviewInAutoMode: () => true, + locks: ["codex-image-generation"], + }); +} diff --git a/packages/rig/sources/tools/codex/index.ts b/packages/rig/sources/tools/codex/index.ts index 049d3e2b..72a13bfb 100644 --- a/packages/rig/sources/tools/codex/index.ts +++ b/packages/rig/sources/tools/codex/index.ts @@ -10,6 +10,7 @@ export { codexInterruptAgentTool } from "./interrupt_agent.js"; export { codexResumeAgentTool } from "./resume_agent.js"; export { codexListAgentsTool } from "./list_agents.js"; export { codexWaitAgentTool } from "./wait_agent.js"; +export { createCodexImageGenerationTool } from "./image_gen.js"; export { unifiedExecOutputSchema } from "./unifiedExecOutput.js"; import { codexApplyPatchTool } from "./apply_patch.js";