diff --git a/docs-site/src/content/docs/guides/image-bridge.md b/docs-site/src/content/docs/guides/image-bridge.md new file mode 100644 index 000000000..86338a5cf --- /dev/null +++ b/docs-site/src/content/docs/guides/image-bridge.md @@ -0,0 +1,75 @@ +--- +title: Image Bridge +description: Route image_generation hosted-tool calls to xAI Grok Imagine when using a non-OpenAI provider. +--- + +## Overview + +When you route Codex through a non-OpenAI model (Claude, Gemini, Grok, etc.), the +`image_generation` **hosted tool** normally doesn't work — it requires OpenAI's server-side +execution environment. The Image Bridge detects these calls and transparently reroutes them to +xAI Grok Imagine, so the model you're actually chatting with can still generate images. + +## Prerequisites + +- **Enable the bridge** by setting `images.bridgeEnabled: true` in your config (it is off by + default to avoid unexpected xAI charges — see [Configuration](#configuration) below). +- An xAI provider configured in settings with `baseUrl: "https://api.x.ai/v1"` and the + `openai-chat` adapter. +- Authentication via `authMode: "oauth"` (`ocx login xai` — uses a stored, auto-refreshed + bearer token) or `authMode: "key"` (a configured API key). +- A non-OpenAI model selected as your active provider. (When the active provider is OpenAI, + the native hosted tool is used directly and the bridge is bypassed.) + +## Configuration + +Image Bridge options live under `images` in `~/.opencodex/config.json`. Bridging is +**opt-in** — you must set `bridgeEnabled: true` to enable paid xAI Grok Imagine generation: + +```json +{ + "images": { + "bridgeEnabled": true, + "bridgeModel": "grok-imagine-image-quality", + "maxRounds": 3 + } +} +``` + +| Option | Default | Description | +| --- | --- | --- | +| `bridgeEnabled` | `false` | Master switch. Set `true` to enable bridging. Off by default to avoid unexpected xAI charges. | +| `bridgeModel` | `grok-imagine-image-quality` | The xAI image model id to send prompts to. | +| `maxRounds` | `3` | Maximum number of image-generation loop iterations per turn. | +| `artifactsKeepCount` | `200` | Maximum number of files retained under `artifacts/`. When exceeded, the oldest files are deleted automatically. | + +## Artifact Retention + +Generated images are written to `~/.opencodex/artifacts/`. To prevent unbounded disk +growth in long-running sessions, the directory is pruned automatically after each image +write — the oldest files (by modification time) are deleted when the count exceeds the +configured maximum (default 200, configurable via `images.artifactsKeepCount`). + +## How It Works + +1. When Codex sends a request with `image_generation` in the tools array, OpenCodex detects it + during request preprocessing. +2. The hosted tool is replaced with a **synthetic function tool** that the routed model can call + normally — the model sees a callable tool rather than an opaque hosted tool it can't execute. +3. When the model invokes that tool, OpenCodex intercepts the call and sends the prompt to xAI's + image generation API. +4. Generated images are saved to `~/.opencodex/artifacts/` and the **local file path** is returned + to the model as the tool result. +5. The model continues the conversation with knowledge of the generated image and its location. + +From the model's perspective nothing changed — it called a tool and got a result. From the user's +perspective, image generation works with any routed provider instead of silently failing. + +## Limitations + +- **Only xAI Grok Imagine is supported.** DALL-E and other image providers may be added later. +- **Web search takes priority.** If both web search and image generation are requested in the same + turn, the web-search bridge runs and image generation is skipped for that turn. +- **xAI costs apply.** Image generation via xAI requires an active xAI subscription or API credits. +- **Streaming only.** The bridge works by intercepting the SSE response stream; requests with + `stream: false` are rejected with a 400 error. diff --git a/src/images/artifacts.ts b/src/images/artifacts.ts new file mode 100644 index 000000000..6e10b5a64 --- /dev/null +++ b/src/images/artifacts.ts @@ -0,0 +1,195 @@ +import { mkdirSync, readdirSync, statSync, unlinkSync } from "node:fs"; +import { mkdir, writeFile } from "node:fs/promises"; +import { join } from "node:path"; +import { getConfigDir } from "../config"; +import { assessUrlDestination, assertUrlResolvesPublic } from "../lib/destination-policy"; + +const MAX_DECODED_BYTES_PER_IMAGE = 50 * 1024 * 1024; +const MAX_DECODED_BYTES_PER_RESPONSE = 100 * 1024 * 1024; +const MAX_DOWNLOAD_BYTES = 50 * 1024 * 1024; // 50 MiB + +/** Default cap on files retained under artifacts/. Oldest files are pruned when exceeded. */ +export const DEFAULT_ARTIFACT_KEEP_COUNT = 200; + +// Strict alphabet check: Buffer.from(..., "base64") silently ignores invalid +// characters, so malformed payloads would otherwise decode to garbage bytes. +const BASE64_RE = /^[A-Za-z0-9+/]*={0,2}$/; + +export interface ImageBudget { + spent: number; +} + +export function createImageBudget(): ImageBudget { + return { spent: 0 }; +} + +function getArtifactsDir(): string { + return join(getConfigDir(), "artifacts"); +} + +/** + * Best-effort retention cap: when the artifact directory holds more than `maxFiles`, + * delete the oldest (by mtime) until the count is back under the limit. Synchronous + * on purpose — it runs right after each successful write and touches at most a handful + * of files. All errors are swallowed and logged so a prune failure never breaks an image write. + */ +export function pruneOldArtifacts(dir: string, maxFiles: number): void { + // A non-positive maxFiles disables pruning entirely (do not delete everything). + if (maxFiles <= 0) return; + let entries: string[]; + try { + entries = readdirSync(dir); + } catch (e) { + console.warn(`[images] prune: could not read ${dir}:`, e instanceof Error ? e.message : e); + return; + } + if (entries.length <= maxFiles) return; + + let stats: Array<{ name: string; mtime: number }>; + try { + stats = entries.map(name => { + const st = statSync(join(dir, name)); + return { name, mtime: st.mtimeMs }; + }); + } catch (e) { + console.warn(`[images] prune: could not stat files in ${dir}:`, e instanceof Error ? e.message : e); + return; + } + + // Sort oldest-first, delete the excess. + stats.sort((a, b) => a.mtime - b.mtime); + const toDelete = stats.slice(0, stats.length - maxFiles); + for (const { name } of toDelete) { + try { + unlinkSync(join(dir, name)); + } catch (e) { + console.warn(`[images] prune: could not delete ${name}:`, e instanceof Error ? e.message : e); + } + } +} + +function timestampPrefix(): string { + const now = new Date(); + return [ + now.getFullYear(), + String(now.getMonth() + 1).padStart(2, "0"), + String(now.getDate()).padStart(2, "0"), + "-", + String(now.getHours()).padStart(2, "0"), + String(now.getMinutes()).padStart(2, "0"), + String(now.getSeconds()).padStart(2, "0"), + "-", + String(now.getMilliseconds()).padStart(3, "0"), + ].join(""); +} + +export function guessExtFromMagic(bytes: Uint8Array): string { + const sig = Buffer.from(bytes.slice(0, 12)).toString("latin1"); + if (sig.startsWith("\x89PNG\r\n\x1a\n")) return "png"; + if (sig.startsWith("\xff\xd8\xff")) return "jpg"; + if (sig.startsWith("RIFF") && sig.slice(8, 12) === "WEBP") return "webp"; + if (sig.startsWith("GIF8")) return "gif"; + return "png"; +} + +export async function materializeInlineImage( + base64Data: string, + budget?: ImageBudget, + keepCount?: number, +): Promise { + const dir = getArtifactsDir(); + await mkdir(dir, { recursive: true, mode: 0o700 }); + + const normalized = base64Data.replace(/\s+/g, ""); + if (!BASE64_RE.test(normalized) || normalized.length % 4 !== 0) { + throw new Error("inline image data is not valid base64"); + } + // Validate decoded size from the base64 length *before* allocating a Buffer, so a + // malicious or broken upstream cannot force a large allocation / OOM. + const padding = normalized.endsWith("==") ? 2 : normalized.endsWith("=") ? 1 : 0; + const decodedBytes = (normalized.length / 4) * 3 - padding; + if (decodedBytes === 0) throw new Error("inline image data is empty after base64 decode"); + if (decodedBytes > MAX_DECODED_BYTES_PER_IMAGE) throw new Error(`inline image exceeds ${MAX_DECODED_BYTES_PER_IMAGE} byte per-image cap`); + if (budget && budget.spent + decodedBytes > MAX_DECODED_BYTES_PER_RESPONSE) { + throw new Error(`inline image response exceeds ${MAX_DECODED_BYTES_PER_RESPONSE} byte per-response cap`); + } + + const buf = Buffer.from(normalized, "base64"); + if (budget) budget.spent += buf.length; + + // Sniff actual format from decoded bytes rather than trusting the declared mimeType. + const ext = guessExtFromMagic(buf); + const filePath = join(dir, `img-${timestampPrefix()}-${crypto.randomUUID()}.${ext}`); + await writeFile(filePath, buf, { mode: 0o600 }); + pruneOldArtifacts(dir, keepCount ?? DEFAULT_ARTIFACT_KEEP_COUNT); + return filePath; +} + +export async function downloadImageToArtifact( + url: string, + budget?: ImageBudget, + signal?: AbortSignal, + keepCount?: number, +): Promise { + if (url.startsWith("data:")) { + const m = /^data:([^;]+);base64,(.+)$/.exec(url); + if (!m) throw new Error("data URL is not a valid base64 image"); + return materializeInlineImage(m[2], budget, keepCount); + } + + // SSRF protection: validate the provider-returned URL before fetching. + // Require HTTPS strictly — plain HTTP and all other schemes (ftp, file, …) are rejected. + let parsedUrl: URL; + try { parsedUrl = new URL(url); } catch { throw new Error("image URL is not valid"); } + if (parsedUrl.protocol !== "https:") { + throw new Error(`image URL must use HTTPS, got ${parsedUrl.protocol}`); + } + // Reject literal private/loopback/link-local/metadata addresses. + const assessment = assessUrlDestination(url); + if (assessment && assessment.kind !== "public" && assessment.kind !== "hostname") { + throw new Error(`image URL targets ${assessment.detail}`); + } + // DNS check: resolve hostname and reject if it points at private/internal space. + await assertUrlResolvesPublic(url); + const resp = await fetch(url, { signal, redirect: "error" }); + if (!resp.ok) throw new Error("image download failed: " + resp.status); + + // Stream the body with a hard byte cap so a missing/lying Content-Length or a + // compromised CDN URL cannot exhaust memory before the size check runs. + if (!resp.body) throw new Error("image download returned no body"); + const reader = resp.body.getReader(); + const chunks: Uint8Array[] = []; + let total = 0; + try { + for (;;) { + const { value, done } = await reader.read(); + if (done) break; + total += value.byteLength; + if (total > MAX_DOWNLOAD_BYTES) { + throw new Error(`image download exceeds ${MAX_DOWNLOAD_BYTES} byte cap`); + } + chunks.push(value); + } + } finally { + try { await reader.cancel(); } catch { /* ignore cancel errors */ } + reader.releaseLock(); + } + + const bytes = new Uint8Array(total); + let offset = 0; + for (const c of chunks) { bytes.set(c, offset); offset += c.byteLength; } + + if (budget && budget.spent + bytes.length > MAX_DECODED_BYTES_PER_RESPONSE) { + throw new Error(`image download exceeds ${MAX_DECODED_BYTES_PER_RESPONSE} byte per-response budget`); + } + + const ext = guessExtFromMagic(bytes); + const dir = getArtifactsDir(); + await mkdir(dir, { recursive: true, mode: 0o700 }); + if (budget) budget.spent += bytes.length; + + const filePath = join(dir, `dl-${timestampPrefix()}-${crypto.randomUUID()}.${ext}`); + await writeFile(filePath, bytes, { mode: 0o600 }); + pruneOldArtifacts(dir, keepCount ?? DEFAULT_ARTIFACT_KEEP_COUNT); + return filePath; +} diff --git a/src/images/fulfill.ts b/src/images/fulfill.ts new file mode 100644 index 000000000..48367da9a --- /dev/null +++ b/src/images/fulfill.ts @@ -0,0 +1,76 @@ +import type { ImageBridgePlan, ImageCallResult } from "./types"; +import { callXaiImages } from "./xai-client"; +import { materializeInlineImage, downloadImageToArtifact, type ImageBudget } from "./artifacts"; + +/** + * Fulfill ONE image-generation tool call end-to-end: parse args, call xAI, materialize the returned + * images to disk, and return a structured result. NEVER throws — all errors become `{ ok: false }` + * so the caller can inject the error as a tool result and let the model respond gracefully. + */ +export async function fulfillImageCall( + call: { id: string; name: string; arguments: string }, + plan: ImageBridgePlan, + budget: ImageBudget, + signal?: AbortSignal, +): Promise { + let args: unknown; + try { + args = JSON.parse(call.arguments || "{}"); + } catch { + return { ok: false, model: plan.model, prompt: "", files: [], count: 0, error: "invalid arguments JSON" }; + } + if (typeof args !== "object" || args === null) { + return { ok: false, model: plan.model, prompt: "", files: [], count: 0, error: "invalid arguments JSON" }; + } + const obj = args as Record; + + const prompt = + typeof obj.prompt === "string" ? obj.prompt : typeof obj.input === "string" ? obj.input : ""; + if (!prompt) { + return { ok: false, model: plan.model, prompt: "", files: [], count: 0, error: "missing prompt" }; + } + + const n = typeof obj.n === "number" ? Math.max(1, Math.min(4, Math.floor(obj.n))) : 1; + const imageUrl = + typeof obj.image_url === "string" ? obj.image_url : typeof obj.image === "string" ? obj.image : undefined; + const size = typeof obj.size === "string" ? obj.size : undefined; + const quality = typeof obj.quality === "string" ? obj.quality : undefined; + + let result; + try { + result = await callXaiImages({ prompt, model: plan.model, n, imageUrl, size, quality }, plan.auth, signal); + } catch (e) { + const error = e instanceof Error ? e.message : String(e); + return { ok: false, model: plan.model, prompt, files: [], count: 0, error }; + } + + const files: string[] = []; + for (const img of result.images ?? []) { + try { + if (img.b64_json) { + files.push(await materializeInlineImage(img.b64_json, budget, plan.artifactsKeepCount)); + } else if (img.url) { + files.push(await downloadImageToArtifact(img.url, budget, signal, plan.artifactsKeepCount)); + } + } catch { + // Keep warnings URL-free — error messages may embed provider CDN URLs. + // Partial success is OK — silently skip this image and continue. + console.warn("[images] failed to materialize image"); + } + } + + if (files.length === 0) { + return { ok: false, model: plan.model, prompt, files: [], count: 0, error: "image generation returned no usable images" }; + } + + const primary = files[0]; + return { + ok: true, + model: plan.model, + prompt, + path: primary, + files, + count: files.length, + markdown: `![image](${primary})`, + }; +} diff --git a/src/images/index.ts b/src/images/index.ts new file mode 100644 index 000000000..6f0adb1fd --- /dev/null +++ b/src/images/index.ts @@ -0,0 +1,4 @@ +export { planImageBridge, findXaiProvider, resolveXaiToken } from "./plan"; +export { runWithImageBridge } from "./loop"; +export type { ImageBridgePlan, ImageCallResult } from "./types"; +export { buildImageTool, extractHostedImageGeneration, IMAGE_GEN_TOOL_NAME, isImageGenName } from "./synthetic-tool"; diff --git a/src/images/loop.ts b/src/images/loop.ts new file mode 100644 index 000000000..e5b2c5786 --- /dev/null +++ b/src/images/loop.ts @@ -0,0 +1,430 @@ +/** + * Image bridge agentic loop — adapted from src/web-search/loop.ts but significantly simpler. + * + * The routed (non-OpenAI) model runs in a bounded loop. Each iteration is streamed and fully + * buffered internally. If the model calls an image-generation tool, the bridge fulfills it via + * the xAI sidecar, injects the result as a tool_result, and loops (bounded by maxRounds). When + * the model produces a real tool call or the budget is exhausted, the passthrough events are + * replayed to the bridge for final SSE output. + * + * Removed vs web-search: no sidecar backend selection, no 429 key-failover, no forced-answer + * nudge, no failed-query dedup, no describeImages/structuredOutput, no recordSidecarOutcome. + */ +import type { ProviderAdapter } from "../adapters/base"; +import { createAdapterEventQueue } from "../adapters/run-turn-queue"; +import type { AdapterEvent, OcxMessage, OcxParsedRequest, OcxThinkingContent } from "../types"; +import { namespacedToolName } from "../types"; +import { bridgeToResponsesSSE } from "../bridge"; +import { clearableDeadline } from "../lib/abort"; +import { readBoundedResponseBody } from "../lib/bounded-body"; +import { fetchWithResetRetry } from "../lib/upstream-retry"; +import { parseStreamWithProgress, RoutedModelInactivityError, WebSearchStreamProtocolError } from "../web-search/progress-stream"; +import { fulfillImageCall } from "./fulfill"; +import { createImageBudget } from "./artifacts"; +import type { ImageBridgePlan } from "./types"; + +const SSE_HEADERS = { + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", +}; + +const CONNECT_TIMEOUT_MS = 200_000; +const STALL_TIMEOUT_MS = 200_000; +const DEFAULT_MAX_ROUNDS = 3; + +interface ImageCall { + id: string; + name: string; + args: string; +} + +/** + * Split an iteration's adapter events into (a) the image-generation tool calls to intercept and + * (b) the events to pass through to Codex. An image tool-call's own start/delta/end events are + * dropped (Codex never sees the synthetic tool); every other event — text, thinking, real tool + * calls, done — is preserved in order. + */ +function scanEventsForImageCall(events: AdapterEvent[], toolNames: Set): { + calls: ImageCall[]; + passthrough: AdapterEvent[]; + hasRealToolCall: boolean; +} { + const calls: ImageCall[] = []; + const passthrough: AdapterEvent[] = []; + let hasRealToolCall = false; + let pending: { name: string; id: string; argsBuf: string; events: AdapterEvent[] } | null = null; + const flushPending = (): void => { + if (pending && !toolNames.has(pending.name)) { + passthrough.push(...pending.events); + hasRealToolCall = true; + } + pending = null; + }; + for (const e of events) { + if (e.type === "tool_call_start") { + flushPending(); + pending = { name: e.name, id: e.id, argsBuf: "", events: [e] }; + } else if (e.type === "tool_call_delta" && pending) { + pending.argsBuf += e.arguments; + pending.events.push(e); + } else if (e.type === "tool_call_end" && pending) { + pending.events.push(e); + if (toolNames.has(pending.name)) { + calls.push({ id: pending.id, name: pending.name, args: pending.argsBuf }); + } else { + passthrough.push(...pending.events); + hasRealToolCall = true; + } + pending = null; + } else { + flushPending(); + passthrough.push(e); + } + } + flushPending(); + return { calls, passthrough, hasRealToolCall }; +} + +async function* replay(events: AdapterEvent[]): AsyncGenerator { + for (const e of events) yield e; +} + +/** + * Collect the thinking block that preceded an image tool call, so the replayed assistant turn can + * carry it. Anthropic extended thinking REQUIRES the assistant message containing tool_use to start + * with its signed thinking blocks. + */ +function extractIterationThinking(events: AdapterEvent[]): OcxThinkingContent | null { + let thinking = ""; + let signature: string | undefined; + const redacted: string[] = []; + for (const e of events) { + if (e.type === "thinking_delta") thinking += e.thinking; + else if (e.type === "thinking_signature") signature = e.signature; + else if (e.type === "redacted_thinking") redacted.push(e.data); + } + if (!thinking && !signature && redacted.length === 0) return null; + return { + type: "thinking", + thinking, + ...(signature ? { signature } : {}), + ...(redacted.length > 0 ? { redacted } : {}), + }; +} + +function jsonError(status: number, message: string): Response { + return new Response(JSON.stringify({ error: { message, type: "upstream_error", code: null } }), { + status, + headers: { "Content-Type": "application/json" }, + }); +} + +/** Hard provider/parse failure inside an iteration. The eager first iteration converts it to a + * non-2xx jsonError; later (already-streaming) iterations surface it as an in-stream error event. */ +class LoopError extends Error { + constructor(readonly status: number, message: string) { + super(message); + this.name = "LoopError"; + } +} + +export interface ImageBridgeDeps { + parsed: OcxParsedRequest; + adapter: ProviderAdapter; + plan: ImageBridgePlan; + /** Headers forwarded from the original request (e.g. Codex auth). Cloned per iteration. */ + forwardHeaders?: Headers; + /** Called before each routed-model dispatch in the bridge loop, for attempt telemetry. */ + onAttemptSend?: () => void; + abortSignal?: AbortSignal; + onFirstOutput?: () => void; + /** Max image-generation rounds before forcing a final answer. Defaults to 3. */ + maxRounds?: number; +} + +/** + * Run the main (non-OpenAI) model in a small agentic loop. Each upstream iteration is streamed and + * fully buffered internally so raw byte progress is observable without leaking the synthetic tool or + * preliminary assistant output. If the model invokes image generation, run it via the xAI sidecar, + * inject the answer as a tool_result, and loop (bounded by `maxRounds`). + */ +export async function runWithImageBridge(deps: ImageBridgeDeps): Promise { + const { parsed, adapter, plan, abortSignal } = deps; + const maxRounds = Math.max(0, deps.maxRounds ?? DEFAULT_MAX_ROUNDS); + const HARD_CAP = maxRounds + 1; + + const messages: OcxMessage[] = [...parsed.context.messages]; + const allTools = parsed.context.tools ?? []; + // For the forced-final pass we drop image tools so the model MUST answer from the results already + // in `messages` (can't generate again) — this guarantees a non-empty final answer. + const toolsNoImage = allTools.filter(t => !t.imageGeneration); + const budget = createImageBudget(); + + // Link an internal AbortController to the turn signal so a client cancel of the SSE body aborts + // in-flight model fetches AND the sidecar. + const internalAbort = new AbortController(); + const linkAbort = (): void => internalAbort.abort(abortSignal?.reason); + if (abortSignal) { + if (abortSignal.aborted) linkAbort(); + else abortSignal.addEventListener("abort", linkAbort, { once: true }); + } + const signal = internalAbort.signal; + + interface IterationResponse { + response: Response; + responseAdapter: ProviderAdapter; + } + type IterationSplit = ReturnType; + + // Acquire one iteration's final response headers. The first call is drained eagerly so an initial + // connect/header/HTTP failure stays a non-2xx JSON response. + const prepareIterationEvents = async function* (forceFinal: boolean): AsyncGenerator { + const iterParsed: OcxParsedRequest = { + ...parsed, stream: true, + context: { ...parsed.context, messages, tools: forceFinal ? toolsNoImage : allTools }, + }; + + // runTurn adapters (Cursor) own all upstream communication via an emit callback. They don't + // expose buildRequest/fetchResponse/parseStream to the bridge, so collect their events through + // an AdapterEventQueue and wrap them in a pseudo-response whose parseStream replays them. + if (adapter.runTurn) { + const queue = createAdapterEventQueue({ + onBacklogExceeded: () => internalAbort.abort("runTurn backlog exceeded"), + }); + void adapter + .runTurn( + iterParsed, + { headers: deps.forwardHeaders ? new Headers(deps.forwardHeaders) : new Headers(), abortSignal: signal }, + queue.push, + ) + .then(() => queue.close()) + .catch(err => { + queue.push({ type: "error", message: err instanceof Error ? err.message : String(err) }); + queue.close(); + }); + + const events = await queue.collect(); + deps.onAttemptSend?.(); + + // runTurn adapters signal errors via {type:"error"} events, not HTTP status codes. + const errorEvent = events.find(e => e.type === "error"); + if (errorEvent && errorEvent.type === "error") { + throw new LoopError(502, errorEvent.message); + } + + // Wrap as an adapter whose parseStream replays the collected events. The pseudo-response + // carries an empty (but non-null) body so parseStreamWithProgress can acquire a reader + // without crashing; the wrapped parseStream never reads from it. + const wrappedAdapter: ProviderAdapter = { + ...adapter, + async *parseStream() { + for (const e of events) yield e; + }, + }; + return { response: new Response(new Uint8Array(0), { status: 200 }), responseAdapter: wrappedAdapter }; + } + + const headerDeadline = clearableDeadline(CONNECT_TIMEOUT_MS, signal); + try { + const request = await adapter.buildRequest(iterParsed, { + headers: deps.forwardHeaders ? new Headers(deps.forwardHeaders) : new Headers(), + abortSignal: headerDeadline.signal, + }); + deps.onAttemptSend?.(); + const response = adapter.fetchResponse + ? await adapter.fetchResponse(request, { + abortSignal: headerDeadline.signal, + timeoutMs: CONNECT_TIMEOUT_MS, + returnRawErrors: true, + stream: true, + }) + : await fetchWithResetRetry( + () => { + const h = new Headers(request.headers); + if (!h.has("accept-encoding")) h.set("accept-encoding", "identity"); + return fetch(request.url, { + method: request.method, + headers: h, + body: request.body, + signal: headerDeadline.signal, + }); + }, + { abortSignal: headerDeadline.signal, label: "image-bridge-loop" }, + ); + + // Final headers have arrived. Clear only the deadline timer before ANY body read. + headerDeadline.clear(); + if (!response.ok) { + let body: Awaited>; + try { + body = await readBoundedResponseBody(response, { signal }); + } catch { + if (signal.aborted) throw new LoopError(499, "client closed request during image-bridge"); + throw new LoopError(response.status, `Provider error ${response.status}`); + } + let formatted = ""; + if (body.displaySafe && !body.truncated && body.text.trim() && adapter.formatErrorBody) { + try { + formatted = adapter.formatErrorBody(response.status, response.headers, body.text).trim(); + } catch { /* formatter hooks are best-effort */ } + } + const suffix = formatted ? `: ${formatted.slice(0, 400)}` : ""; + throw new LoopError(response.status, `Provider error ${response.status}${suffix}`); + } + return { response, responseAdapter: adapter }; + } catch (error) { + if (headerDeadline.didExpire()) { + throw new LoopError(504, `Provider response-header timeout after ${CONNECT_TIMEOUT_MS}ms during image-bridge`); + } + if (signal.aborted) throw new LoopError(499, "client closed request during image-bridge"); + if (error instanceof LoopError) throw error; + throw new LoopError(502, `Provider unreachable: ${error instanceof Error ? error.message : String(error)}`); + } finally { + headerDeadline.clear(); + } + }; + + const prepareIterationDrained = async (forceFinal: boolean): Promise => { + const it = prepareIterationEvents(forceFinal); + let r = await it.next(); + while (!r.done) r = await it.next(); + return r.value; + }; + + // Consume and validate one successful response body. Only invisible heartbeat events escape while + // semantic output remains buffered for safe scanning. + const consumeIterationEvents = async function* (prepared: IterationResponse): AsyncGenerator { + const events: AdapterEvent[] = []; + try { + const parse = prepared.responseAdapter.parseStream.bind(prepared.responseAdapter); + for await (const event of parseStreamWithProgress(prepared.response, parse, { + signal, + inactivityTimeoutMs: STALL_TIMEOUT_MS, + })) { + if (event.type === "heartbeat") yield event; + else events.push(event); + } + } catch (error) { + if (signal.aborted) throw new LoopError(499, "client closed request during image-bridge"); + if (error instanceof RoutedModelInactivityError) throw new LoopError(504, error.message); + if (error instanceof WebSearchStreamProtocolError) throw new LoopError(502, error.message); + throw new LoopError(502, `Provider stream error: ${error instanceof Error ? error.message : String(error)}`); + } + + const terminalIndexes = events.flatMap((event, index) => + event.type === "done" || event.type === "incomplete" || event.type === "error" ? [index] : []); + if (terminalIndexes.length !== 1 || terminalIndexes[0] !== events.length - 1) { + throw new LoopError(502, `Image-bridge adapter stream protocol error: expected one final terminal event, received ${terminalIndexes.length}`); + } + const terminal = events[terminalIndexes[0]!]; + if (terminal.type === "error") throw new LoopError(502, terminal.message); + return scanEventsForImageCall(events, plan.toolNames); + }; + + // Eagerly acquire only the FIRST iteration's final headers so connect/header/HTTP failures remain + // non-2xx JSON. + let firstPrepared: IterationResponse; + try { + firstPrepared = await prepareIterationDrained(maxRounds <= 0); + } catch (e) { + if (abortSignal) abortSignal.removeEventListener("abort", linkAbort); + if (e instanceof LoopError) return jsonError(e.status, e.message); + throw e; + } + + const toolNsMap = new Map(); + const freeform = new Set(); + const toolSearch = new Set(); + for (const t of parsed.context.tools ?? []) { + if (t.namespace) toolNsMap.set(namespacedToolName(t.namespace, t.name), { namespace: t.namespace, name: t.name }); + if (t.freeform) freeform.add(t.name); + if (t.toolSearch) toolSearch.add(t.name); + } + + // Drive the remaining iterations live. Image generation runs interleaved with the real sidecar + // timing; the final answer's passthrough events come last. + async function* produce(): AsyncGenerator { + let prepared = firstPrepared; + try { + for (let i = 0; i < HARD_CAP; i++) { + const forceFinal = i >= maxRounds; + try { + // First loop turn reuses the eager HEADERS. Subsequent header acquisitions run here. + if (i > 0) { + yield { type: "heartbeat" }; + prepared = yield* prepareIterationEvents(forceFinal); + } + // Raw-byte progress heartbeats reach the bridge; semantic events remain buffered. + const split = yield* consumeIterationEvents(prepared); + + // Loop (fulfill + re-ask) ONLY when the model's actionable output is purely image_gen. A + // real tool call means this turn is terminal for Codex — finalize so those calls reach + // Codex. forceFinal also finalizes. + const shouldLoop = split.calls.length > 0 && !split.hasRealToolCall && !forceFinal; + if (!shouldLoop) { + yield* replay(split.passthrough); + return; + } + + // Fulfill each image call, then inject assistant + toolResult into messages. + const iterationThinking = extractIterationThinking(split.passthrough); + for (const [callIndex, call] of split.calls.entries()) { + yield { type: "heartbeat" }; + const result = await fulfillImageCall( + { id: call.id, name: call.name, arguments: call.args }, + plan, budget, signal, + ); + if (signal.aborted) throw new LoopError(499, "client closed request during image-bridge"); + const now = Date.now(); + let parsedArgs: Record = {}; + try { + const raw: unknown = JSON.parse(call.args || "{}"); + if (raw !== null && typeof raw === "object" && !Array.isArray(raw)) { + parsedArgs = raw as Record; + } + } catch { /* malformed args */ } + messages.push({ + role: "assistant", + content: [ + ...(callIndex === 0 && iterationThinking ? [iterationThinking] : []), + { type: "toolCall" as const, id: call.id, name: call.name, arguments: parsedArgs }, + ], + timestamp: now, + }); + messages.push({ + role: "toolResult", + toolCallId: call.id, + toolName: call.name, + content: JSON.stringify(result), + isError: !result.ok, + timestamp: now, + }); + } + } catch (e) { + yield { + type: "error", + message: e instanceof LoopError ? e.message : (e instanceof Error ? e.message : String(e)), + ...(e instanceof LoopError ? { status: e.status } : {}), + }; + return; + } + } + } finally { + if (abortSignal) abortSignal.removeEventListener("abort", linkAbort); + } + } + + const sse = bridgeToResponsesSSE( + produce(), parsed.modelId, toolNsMap, freeform, toolSearch, () => { + internalAbort.abort("client closed responses stream"); + }, undefined, + { + responseId: "", + hideThinkingSummary: parsed.options.hideThinkingSummary, + ...(deps.onFirstOutput ? { onFirstOutput: deps.onFirstOutput } : {}), + }, + ); + return new Response(sse, { headers: SSE_HEADERS }); +} diff --git a/src/images/plan.ts b/src/images/plan.ts new file mode 100644 index 000000000..b6a5b708b --- /dev/null +++ b/src/images/plan.ts @@ -0,0 +1,65 @@ +import type { OcxConfig, OcxParsedRequest, OcxProviderConfig } from "../types"; +import type { ImageBridgePlan } from "./types"; +import { getValidAccessToken } from "../oauth/index"; +import { resolveEnvValue } from "../config"; +import { getProviderRegistryEntry } from "../providers/registry"; +import { IMAGE_GEN_TOOL_NAME } from "./synthetic-tool"; + +const DEFAULT_MODEL = "grok-imagine-image-quality"; + +export function findXaiProvider(config: OcxConfig): { name: string; provider: OcxProviderConfig } | undefined { + // Primary: well-known name "xai" + const xai = config.providers["xai"]; + if (xai && xai.disabled !== true) return { name: "xai", provider: xai }; + // Fallback: hostname match for custom-named xAI configs + for (const [name, p] of Object.entries(config.providers)) { + if (p.disabled) continue; + try { + const host = new URL(p.baseUrl).hostname; + if (host === "api.x.ai" || host === "cli-chat-proxy.grok.com") return { name, provider: p }; + } catch { /* invalid baseUrl */ } + } + return undefined; +} + +export async function resolveXaiToken(providerName: string, provider: OcxProviderConfig): Promise { + const apiKey = resolveEnvValue(provider.apiKey)?.trim(); + if (apiKey) return apiKey; + // Built-in OAuth token only for the canonical "xai" provider — never for custom-named configs. + if (providerName !== "xai") return undefined; + try { + return await getValidAccessToken("xai"); + } catch { + return undefined; + } +} + +export async function planImageBridge( + config: OcxConfig, + parsed: OcxParsedRequest, + routedProvider: OcxProviderConfig, +): Promise { + if (config.images?.bridgeEnabled !== true) return undefined; + if (!parsed._imageGeneration) return undefined; + // Don't intercept for OpenAI native passthrough + const host = (() => { try { return new URL(routedProvider.baseUrl).hostname; } catch { return ""; } })(); + if (host === "api.openai.com") return undefined; + const found = findXaiProvider(config); + if (!found) return undefined; + const token = await resolveXaiToken(found.name, found.provider); + if (!token) return undefined; + // Pin the baseUrl to the registry entry, ignoring any config-level baseUrl override. + const registryEntry = getProviderRegistryEntry("xai"); + const pinnedBaseUrl = (registryEntry?.baseUrl ?? "https://api.x.ai/v1").replace(/\/+$/, ""); + // The synthetic tool injected into the conversation is named IMAGE_GEN_TOOL_NAME, + // which is what the model will actually call. Merge it with any original hosted tool names. + const toolNames = new Set(parsed._imageGeneration.toolNames); + toolNames.add(IMAGE_GEN_TOOL_NAME); + return { + provider: found.provider, + auth: { baseUrl: pinnedBaseUrl, token }, + model: config.images?.bridgeModel ?? DEFAULT_MODEL, + toolNames, + artifactsKeepCount: config.images?.artifactsKeepCount, + }; +} diff --git a/src/images/synthetic-tool.ts b/src/images/synthetic-tool.ts new file mode 100644 index 000000000..f6b1a7cbe --- /dev/null +++ b/src/images/synthetic-tool.ts @@ -0,0 +1,70 @@ +import type { OcxTool } from "../types"; + +/** The function name the chat model sees + the name the loop intercepts. */ +export const IMAGE_GEN_TOOL_NAME = "image_gen"; + +const IMAGE_GEN_NAMES = new Set([ + "image_gen", "image_generation", "imagegen", + "generate_image", "generateimage", +]); + +export function isImageGenName(name: string): boolean { + return IMAGE_GEN_NAMES.has(name.toLowerCase()); +} + +/** + * Scan a Responses request's `tools[]` for hosted image-generation entries (`{type:"image_generation"}` or + * `{type:"image_gen"}`) and function entries whose name matches `isImageGenName`. Returns the set of all + * matched tool names plus the first matched raw tool object (so its config can be replayed), or undefined + * when image generation isn't enabled. + */ +export function extractHostedImageGeneration( + tools: unknown[] | undefined, +): { toolNames: Set; originalTool?: Record } | undefined { + if (!Array.isArray(tools)) return undefined; + const toolNames = new Set(); + let originalTool: Record | undefined; + for (const t of tools) { + if (!t || typeof t !== "object") continue; + const obj = t as Record; + if (obj.type === "image_generation" || obj.type === "image_gen") { + if (!originalTool) originalTool = obj; + const name = typeof obj.name === "string" ? obj.name : (obj.type as string); + toolNames.add(name); + } else if (obj.type === "function") { + // Responses API function tools have a flat shape: {type:"function", name:"...", parameters:{...}}. + // Also handle the nested Chat Completions shape {type:"function", function:{name:"..."}} for safety. + const fnName = + typeof obj.name === "string" ? obj.name : (obj as { function?: { name?: string } }).function?.name; + if (fnName && isImageGenName(fnName)) { + if (!originalTool) originalTool = obj; + toolNames.add(fnName); + } + } + } + if (toolNames.size === 0) return undefined; + return { toolNames, originalTool }; +} + +/** + * The synthetic function tool exposed to a chat/anthropic model in place of the dropped hosted + * image_generation. The model calls it like any function; the proxy intercepts the call and runs the + * real generation via the sidecar (the call is never relayed to Codex). `imageGeneration:true` flags it. + */ +export function buildImageTool(): OcxTool { + return { + name: IMAGE_GEN_TOOL_NAME, + description: + "Generate an image from a text prompt. Returns absolute local filesystem path(s). " + + "Use when the user asks to create or draw an image.", + parameters: { + type: "object", + properties: { + prompt: { type: "string", description: "Detailed image generation prompt. Required." }, + n: { type: "integer", minimum: 1, maximum: 4 }, + }, + required: ["prompt"], + }, + imageGeneration: true, + }; +} diff --git a/src/images/types.ts b/src/images/types.ts new file mode 100644 index 000000000..830a1e4cf --- /dev/null +++ b/src/images/types.ts @@ -0,0 +1,21 @@ +import type { OcxProviderConfig } from "../types"; + +export interface ImageBridgePlan { + provider: OcxProviderConfig; + auth: { baseUrl: string; token: string }; + model: string; + toolNames: Set; + /** Max artifact files to retain (from config.images.artifactsKeepCount). Default 200. */ + artifactsKeepCount?: number; +} + +export interface ImageCallResult { + ok: boolean; + model: string; + prompt: string; + path?: string; + files: string[]; + count: number; + markdown?: string; + error?: string; +} diff --git a/src/images/xai-client.ts b/src/images/xai-client.ts new file mode 100644 index 000000000..4bb3b67b0 --- /dev/null +++ b/src/images/xai-client.ts @@ -0,0 +1,133 @@ +/** + * xAI image generation/editing client. + * + * Calls xAI's `/images/generations` or `/images/edits` endpoint, composing a + * 60 s timeout with the caller's abort signal so the deadline covers the entire + * response body read. Non-2xx responses throw with the original status code — + * no 502 compression — so callers can distinguish rate-limit / auth failures + * from transient errors. + */ + +export interface XaiImageRequest { + prompt: string; + model?: string; // default "grok-imagine-image-quality" + n?: number; // 1-4 + size?: string; + quality?: string; + imageUrl?: string; // if set → /images/edits +} + +export interface XaiImageResult { + images: Array<{ b64_json?: string; url?: string }>; +} + +const XAI_IMAGES_TIMEOUT_MS = 60_000; +const XAI_DEFAULT_MODEL = "grok-imagine-image-quality"; + +// xAI only accepts aspect_ratio + resolution, not OpenAI's size/quality. Map +// OpenAI's "WxH" size to the closest standard ratio (log-space distance) and +// OpenAI quality buckets onto xAI's 1k/2k resolution. Unknown values are +// dropped rather than forwarded, to avoid sending parameters xAI rejects. +const XAI_ASPECT_RATIOS: ReadonlyArray = [ + ["1:1", 1], + ["3:4", 0.75], + ["4:3", 4 / 3], + ["9:16", 0.5625], + ["16:9", 16 / 9], +]; + +function mapSizeToAspectRatio(size?: string): string | undefined { + if (!size) return undefined; + const m = /^(\d+)x(\d+)$/.exec(size); + if (!m) return undefined; + const ratio = parseInt(m[1], 10) / parseInt(m[2], 10); + let best = XAI_ASPECT_RATIOS[0]!; + let bestDiff = Infinity; + for (const [label, r] of XAI_ASPECT_RATIOS) { + const diff = Math.abs(Math.log(ratio / r)); + if (diff < bestDiff) { + bestDiff = diff; + best = [label, r] as const; + } + } + return best[0]; +} + +function mapQualityToResolution(quality?: string): string | undefined { + if (!quality) return undefined; + const q = quality.toLowerCase(); + if (q === "hd" || q === "high") return "2k"; + if (q === "standard" || q === "low" || q === "auto") return "1k"; + return undefined; +} + +export async function callXaiImages( + req: XaiImageRequest, + auth: { baseUrl: string; token: string }, + signal?: AbortSignal, +): Promise { + const isEdit = typeof req.imageUrl === "string" && req.imageUrl.length > 0; + const endpoint = isEdit ? "/images/edits" : "/images/generations"; + + const body: Record = { + model: req.model ?? XAI_DEFAULT_MODEL, + prompt: req.prompt, + n: req.n ?? 1, + }; + const aspectRatio = mapSizeToAspectRatio(req.size); + const resolution = mapQualityToResolution(req.quality); + if (aspectRatio) body.aspect_ratio = aspectRatio; + if (resolution) body.resolution = resolution; + if (isEdit) { + body.image = { url: req.imageUrl as string, type: "image_url" }; + } + + const timeout = AbortSignal.timeout(XAI_IMAGES_TIMEOUT_MS); + const linkedSignal = signal ? AbortSignal.any([signal, timeout]) : timeout; + + const resp = await fetch(`${auth.baseUrl}${endpoint}`, { + method: "POST", + headers: { + "Authorization": `Bearer ${auth.token}`, + "Content-Type": "application/json", + }, + body: JSON.stringify(body), + signal: linkedSignal, + }); + + if (!resp.ok) { + throw new Error("xAI images API returned " + resp.status); + } + + // Read the body as text under the linked signal, then parse. The 60 s timeout + // and caller abort cover the read. A 200 MiB hard cap on the text prevents a + // runaway response from exhausting memory before materialization caps apply. + const MAX_RESPONSE_BYTES = 200 * 1024 * 1024; + const reader = resp.body?.getReader(); + if (!reader) throw new Error("xAI images API returned no body"); + const decoder = new TextDecoder(); + let text = ""; + let totalBytes = 0; + try { + for (;;) { + const { value, done } = await reader.read(); + if (done) break; + totalBytes += value.byteLength; + if (totalBytes > MAX_RESPONSE_BYTES) throw new Error("xAI images API response exceeds size cap"); + text += decoder.decode(value, { stream: true }); + } + text += decoder.decode(); + } finally { + try { await reader.cancel(); } catch { /* ignore cancel errors */ } + reader.releaseLock(); + } + + const json = JSON.parse(text) as { data?: Array<{ b64_json?: string; url?: string }> }; + + const images = (json.data ?? []).map((entry) => ({ + b64_json: entry.b64_json, + url: entry.url, + })); + + return { images }; +} diff --git a/src/lib/destination-policy.ts b/src/lib/destination-policy.ts index 54cffcbd1..3ed0bf533 100644 --- a/src/lib/destination-policy.ts +++ b/src/lib/destination-policy.ts @@ -19,7 +19,7 @@ const BLOCKED_METADATA_IPV6 = new Set([ "fd00:ec2::254", ]); -type DestinationKind = +export type DestinationKind = | "public" | "hostname" | "localhost" @@ -79,6 +79,16 @@ function classifyIpv6(hostname: string): DestinationAssessment { if (BLOCKED_METADATA_IPV6.has(hostname)) return { kind: "metadata", detail: "blocked metadata endpoint" }; const mappedIpv4 = hostname.match(/^::ffff:(\d+\.\d+\.\d+\.\d+)$/i)?.[1]; if (mappedIpv4) return classifyIpv4(mappedIpv4); + // Decode hex IPv4-mapped IPv6: ::ffff:7f00:1 → 127.0.0.1 + // The dotted-decimal regex above only matches ::ffff:127.0.0.1; without this, + // hex form bypasses all private/loopback checks (hextet is 0 → classified "public"). + const hexMapped = hostname.match(/^::ffff:([0-9a-f]{1,4}):([0-9a-f]{1,4})$/i); + if (hexMapped) { + const hi = Number.parseInt(hexMapped[1], 16); + const lo = Number.parseInt(hexMapped[2], 16); + const ipv4 = `${(hi >> 8) & 255}.${hi & 255}.${(lo >> 8) & 255}.${lo & 255}`; + return classifyIpv4(ipv4); + } if (hostname === "::1") return { kind: "loopback", detail: "loopback address" }; if (hostname === "::") return { kind: "unspecified", detail: "unspecified address" }; const hextet = firstIpv6Hextet(hostname); @@ -178,3 +188,52 @@ export async function providerDestinationResolvedError( } return null; } + +export interface UrlDestinationAssessment { + kind: DestinationKind; + detail: string; +} + +/** + * Synchronous literal URL destination assessment — classifies the hostname + * without DNS resolution. Returns null for unparseable URLs. + */ +export function assessUrlDestination(url: string): UrlDestinationAssessment | null { + return assessDestination(url); +} + +/** + * Async DNS-resolved URL safety check. Resolves A/AAAA records and rejects + * if any address is loopback, private, link-local, unspecified, or metadata. + * Throws on unsafe destination; returns void on safe/public destination. + * DNS resolution failures are treated as unsafe (fail-closed). + */ +export async function assertUrlResolvesPublic(url: string): Promise { + let hostname: string; + try { + hostname = normalizeHostname(new URL(url.trim()).hostname); + } catch { + throw new Error("image URL is not a valid URL"); + } + if (!hostname) throw new Error("image URL has no hostname"); + const literalAssessment = assessDestination(url); + if (literalAssessment && literalAssessment.kind !== "public" && literalAssessment.kind !== "hostname") { + throw new Error(`image URL targets ${literalAssessment.detail}`); + } + // For literal IPs and localhost, the sync path already classified them. + if (isIP(hostname) !== 0 || hostname === "localhost" || hostname.endsWith(".localhost")) return; + let addresses: { address: string }[]; + try { + addresses = await lookup(hostname, { all: true, verbatim: true }); + } catch { + // If DNS fails, we can't verify — fail-closed (unlike provider config-time validation, + // this is a runtime fetch to an untrusted URL, so be conservative). + throw new Error(`image URL hostname ${hostname} could not be resolved`); + } + for (const { address } of addresses) { + const ipKind = isIP(address); + const assessment = ipKind === 4 ? classifyIpv4(address) : ipKind === 6 ? classifyIpv6(normalizeHostname(address)) : null; + if (!assessment || assessment.kind === "public") continue; + throw new Error(`image URL hostname ${hostname} resolves to ${assessment.detail} (${address})`); + } +} diff --git a/src/responses/parser.ts b/src/responses/parser.ts index 3c13b477e..c0a311c1c 100644 --- a/src/responses/parser.ts +++ b/src/responses/parser.ts @@ -16,6 +16,7 @@ import { compactionItemToText } from "./compaction"; import { previousResponseReplayPrefixLength } from "./state"; import { decodeReasoningEnvelope } from "./reasoning-envelope"; import { extractHostedWebSearch, WEB_SEARCH_TOOL_NAME } from "../web-search/synthetic-tool"; +import { extractHostedImageGeneration } from "../images/synthetic-tool"; function isObj(v: unknown): v is Record { return typeof v === "object" && v !== null && !Array.isArray(v); @@ -612,6 +613,10 @@ export function parseRequest(body: unknown): OcxParsedRequest { // gpt-mini sidecar for routed providers. buildTools still drops the hosted tool; the sidecar path // re-injects a synthetic function tool only when it will actually handle the call. const webSearch = extractHostedWebSearch(data.tools as unknown[] | undefined); + const imageGen = extractHostedImageGeneration([ + ...(data.tools as unknown[] ?? []), + ...loadedToolSpecs, + ]); // Detect structured-output mode (Responses `text.format`) so the web-search sidecar can render its // tool_result as JSON rather than prose that could corrupt the model's schema-constrained answer. const structuredOutput = detectStructuredOutput(data.text); @@ -625,6 +630,7 @@ export function parseRequest(body: unknown): OcxParsedRequest { _rawBody: body, ...(replayedInputPrefixLength > 0 ? { _replayPrefixLen: replayedInputPrefixLength } : {}), ...(webSearch ? { _webSearch: webSearch } : {}), + ...(imageGen ? { _imageGeneration: imageGen } : {}), ...(structuredOutput ? { _structuredOutput: true } : {}), ...(compactionRequest ? { _compactionRequest: true } : {}), ...(contextCompactionBoundary ? { _contextCompactionBoundary: true } : {}), diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index fbed02554..0f3aeff69 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -39,6 +39,7 @@ import { UnsupportedOAuthProviderError, } from "../../oauth"; import { buildWebSearchTool, planWebSearch, runWithWebSearch, shouldResolveOpenAiWebSearchSidecar } from "../../web-search"; +import { buildImageTool, planImageBridge, runWithImageBridge } from "../../images"; import { describeImagesInPlace, planVisionSidecar, shouldResolveOpenAiVisionSidecar, stripImagesInPlace } from "../../vision"; import { createAdapterEventQueue, preflightAdapterEvents } from "../../adapters/run-turn-queue"; import { @@ -1526,6 +1527,104 @@ export async function handleResponses( }); } + // Image bridge: check BEFORE the runTurn early-return so runTurn adapters (e.g. Cursor) + // also route through the bridge when image_generation is requested. The bridge loop + // internally supports both standard and runTurn adapter paths. + // Routed-compaction turns must NOT hit the bridge: compaction clears tools/_webSearch but + // leaves _imageGeneration, so planImageBridge would activate and return a normal Responses + // completion instead of the synthetic compaction item Codex expects (#424). + const imgPlan = !routedCompaction ? await planImageBridge(config, parsed, route.provider) : undefined; + if (imgPlan) { + // Web-search takes priority when both are eligible (design: image defers to web-search). + const wsPlan = planWebSearch(config, parsed, false, route.provider, route.modelId, openAiSidecar); + if (!wsPlan) { + // The bridge forces stream:true internally and returns SSE. Non-streaming requests can't be + // served — reject explicitly rather than returning SSE to a client expecting JSON. + if (!parsed.stream) { + return formatErrorResponse(400, "invalid_request_error", "image bridge requires stream=true"); + } + parsed.context.tools = [...(parsed.context.tools ?? []), buildImageTool()]; + const imgResponse = await runWithImageBridge({ + parsed, adapter, + plan: imgPlan, + forwardHeaders: selectedForwardHeaders, + onAttemptSend: () => noteAttemptSend(logCtx.activeAttempt, logCtx.usageLogInputTokens), + abortSignal: options.abortSignal, + ...(config.images?.maxRounds != null ? { maxRounds: config.images.maxRounds } : {}), + ...(options.onFirstOutput ? { onFirstOutput: options.onFirstOutput } : {}), + }); + if (imgResponse.body) { + const imgTurnAc = new AbortController(); + return new Response(trackStreamLifetime(imgResponse.body, imgTurnAc), { + status: imgResponse.status, + headers: imgResponse.headers, + }); + } + return imgResponse; + } + } + + // Web-search sidecar: Codex enabled web_search but this is a routed (non-OpenAI) model that can't + // run it server-side. Expose web_search as a function tool and run searches via the gpt-mini sidecar + // through the ChatGPT passthrough, looping until the model answers. Otherwise take the normal path. + // Placed BEFORE the runTurn early-return so dual-tool turns (web_search + image_generation) on + // runTurn adapters (e.g. Cursor) dispatch through the web-search sidecar instead of being swallowed + // by the runTurn branch — which would leave dual-tool turns with neither bridge. + const wsPlan = planWebSearch(config, parsed, false, route.provider, route.modelId, openAiSidecar); + if (wsPlan) { + parsed.context.tools = [...(parsed.context.tools ?? []), buildWebSearchTool()]; + noteAttemptSend(logCtx.activeAttempt, logCtx.usageLogInputTokens); + const wsResponse = await runWithWebSearch({ + parsed, adapter, + backend: wsPlan.backend, + forwardProvider: wsPlan.forwardSidecar?.provider, + anthropicSidecar: wsPlan.anthropicSidecar, + hostedTool: wsPlan.hostedTool, + selectedForwardHeaders: wsPlan.forwardSidecar?.headers ?? selectedForwardHeaders, + settings: wsPlan.settings, + maxSearches: wsPlan.maxSearches, + forceEmptyResponseId: true, + abortSignal: options.abortSignal, + ...(options.onFirstOutput ? { onFirstOutput: options.onFirstOutput } : {}), + onRequestBuilt: request => recordAdapterReasoning(logCtx, request), + onUsage: usage => { + logCtx.usageFromBridge = true; + if (usage) { + logCtx.usage = usage; + if (logCtx.activeAttempt) logCtx.activeAttempt.usage = usage; + } + }, + recordSidecarOutcome: wsPlan.forwardSidecar?.recordOutcome, + connectTimeoutMs: config.connectTimeoutMs ?? 200_000, + routedModelStallTimeoutMs: wsPlan.routedModelStallTimeoutMs, + stallTimeoutSec: wsPlan.stallTimeoutSec, + on429: retryAfter => { + const rotated = rotateProviderTransportOn429(config, route.providerName, { + retryAfter, + now: Date.now(), + attemptedKey: route.provider.apiKey, + promptCacheKey: parsed.options.promptCacheKey, + }); + if (!rotated) return null; + route.provider = rotated; + return resolveAdapter( + resolveWireProtocolOverride(route.providerName, route.modelId, route.provider), + config.cacheRetention, + ); + }, + }); + // Register the sidecar stream as an active turn so drainAndShutdown waits for (or aborts) + // in-flight web-search turns instead of skipping them during graceful shutdown. + if (wsResponse.body) { + const wsTurnAc = new AbortController(); + return new Response(trackStreamLifetime(wsResponse.body, wsTurnAc), { + status: wsResponse.status, + headers: wsResponse.headers, + }); + } + return wsResponse; + } + if (adapter.runTurn) { const runTurnAbort = new AbortController(); linkAbortSignal(runTurnAbort, options.abortSignal); @@ -1646,64 +1745,6 @@ export async function handleResponses( return new Response(JSON.stringify(json), { headers: { "Content-Type": "application/json" } }); } - // Web-search sidecar: Codex enabled web_search but this is a routed (non-OpenAI) model that can't - // run it server-side. Expose web_search as a function tool and run searches via the gpt-mini sidecar - // through the ChatGPT passthrough, looping until the model answers. Otherwise take the normal path. - const wsPlan = planWebSearch(config, parsed, false, route.provider, route.modelId, openAiSidecar); - if (wsPlan) { - parsed.context.tools = [...(parsed.context.tools ?? []), buildWebSearchTool()]; - noteAttemptSend(logCtx.activeAttempt, logCtx.usageLogInputTokens); - const wsResponse = await runWithWebSearch({ - parsed, adapter, - backend: wsPlan.backend, - forwardProvider: wsPlan.forwardSidecar?.provider, - anthropicSidecar: wsPlan.anthropicSidecar, - hostedTool: wsPlan.hostedTool, - selectedForwardHeaders: wsPlan.forwardSidecar?.headers ?? selectedForwardHeaders, - settings: wsPlan.settings, - maxSearches: wsPlan.maxSearches, - forceEmptyResponseId: true, - abortSignal: options.abortSignal, - ...(options.onFirstOutput ? { onFirstOutput: options.onFirstOutput } : {}), - onRequestBuilt: request => recordAdapterReasoning(logCtx, request), - onUsage: usage => { - logCtx.usageFromBridge = true; - if (usage) { - logCtx.usage = usage; - if (logCtx.activeAttempt) logCtx.activeAttempt.usage = usage; - } - }, - recordSidecarOutcome: wsPlan.forwardSidecar?.recordOutcome, - connectTimeoutMs: config.connectTimeoutMs ?? 200_000, - routedModelStallTimeoutMs: wsPlan.routedModelStallTimeoutMs, - stallTimeoutSec: wsPlan.stallTimeoutSec, - on429: retryAfter => { - const rotated = rotateProviderTransportOn429(config, route.providerName, { - retryAfter, - now: Date.now(), - attemptedKey: route.provider.apiKey, - promptCacheKey: parsed.options.promptCacheKey, - }); - if (!rotated) return null; - route.provider = rotated; - return resolveAdapter( - resolveWireProtocolOverride(route.providerName, route.modelId, route.provider), - config.cacheRetention, - ); - }, - }); - // Register the sidecar stream as an active turn so drainAndShutdown waits for (or aborts) - // in-flight web-search turns instead of skipping them during graceful shutdown. - if (wsResponse.body) { - const wsTurnAc = new AbortController(); - return new Response(trackStreamLifetime(wsResponse.body, wsTurnAc), { - status: wsResponse.status, - headers: wsResponse.headers, - }); - } - return wsResponse; - } - const upstream = new AbortController(); const cleanupUpstreamAbort = linkAbortSignal(upstream, options.abortSignal); const connectMs = config.connectTimeoutMs ?? 200_000; diff --git a/src/types.ts b/src/types.ts index 822ac4984..3ee03a819 100644 --- a/src/types.ts +++ b/src/types.ts @@ -31,6 +31,8 @@ export interface OcxParsedRequest { * executes searches via the gpt-5.4-mini sidecar (see src/web-search). Absent when not requested. */ _webSearch?: Record; + /** Hosted image_generation tool config stashed for the image bridge sidecar (see src/images). */ + _imageGeneration?: { toolNames: Set; originalTool?: Record }; /** * True when Codex requested structured output (`text.format` = json_schema/json_object). The * web-search tool_result is then rendered as compact JSON instead of markdown prose, so its @@ -152,6 +154,8 @@ export interface OcxTool { loadedFromToolSearch?: boolean; /** Synthetic web_search tool: the model's call is executed by the gpt-5.4-mini sidecar, not relayed to Codex. */ webSearch?: boolean; + /** Synthetic image_gen tool: the model's call is executed by the xAI image bridge sidecar, not relayed to Codex. */ + imageGeneration?: boolean; } /** @@ -716,6 +720,14 @@ export interface OcxImagesConfig { provider?: string; /** Upstream timeout (ms) for one /v1/images relay. Default 300000 — generation is slow. */ timeoutMs?: number; + /** Master switch for the image bridge. Default false — set true to enable paid xAI Grok Imagine generation. */ + bridgeEnabled?: boolean; + /** xAI image model id. Default "grok-imagine-image-quality" (see DEFAULT_MODEL in images/plan.ts). */ + bridgeModel?: string; + /** Max image-generation loop iterations before forced-final. Default 3 (see DEFAULT_MAX_ROUNDS in images/loop.ts). */ + maxRounds?: number; + /** Max files retained under artifacts/. Oldest deleted when exceeded. Default 200. */ + artifactsKeepCount?: number; } export interface OcxSearchConfig { diff --git a/tests/images/artifacts-prune.test.ts b/tests/images/artifacts-prune.test.ts new file mode 100644 index 000000000..e42ae72e2 --- /dev/null +++ b/tests/images/artifacts-prune.test.ts @@ -0,0 +1,95 @@ +import { afterAll, beforeAll, describe, expect, test } from "bun:test"; +import { tmpdir } from "node:os"; +import { basename, dirname, join } from "node:path"; +import { randomUUID } from "node:crypto"; +import { mkdirSync, mkdtempSync, readdirSync, utimesSync, writeFileSync } from "node:fs"; +import { rm } from "node:fs/promises"; + +const PREV_HOME = process.env.OPENCODEX_HOME; +beforeAll(() => { process.env.OPENCODEX_HOME = join(tmpdir(), "ocx-prune-" + randomUUID()); }); +afterAll(() => { if (PREV_HOME === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = PREV_HOME; }); + +const { pruneOldArtifacts } = await import("../../src/images/artifacts"); + +function touch(path: string, content: string = "x", ageMs = 0): void { + writeFileSync(path, content); + if (ageMs > 0) { + const t = (Date.now() - ageMs) / 1000; + utimesSync(path, t, t); + } +} + +describe("pruneOldArtifacts", () => { + test("count <= maxFiles → no deletion", () => { + const dir = mkdtempSync(join(tmpdir(), "prune-keep-")); + for (let i = 0; i < 5; i++) touch(join(dir, `f${i}.png`)); + pruneOldArtifacts(dir, 10); + expect(readdirSync(dir).length).toBe(5); + }); + + test("count > maxFiles → oldest deleted until under limit", () => { + const dir = mkdtempSync(join(tmpdir(), "prune-trim-")); + // Create 10 files with staggered mtimes — f0 is oldest (10s ago), f9 is newest (1s ago). + for (let i = 0; i < 10; i++) { + touch(join(dir, `f${i}.png`), "data", (10 - i) * 1000); + } + pruneOldArtifacts(dir, 5); + const remaining = readdirSync(dir); + expect(remaining.length).toBe(5); + // The 5 newest (f5–f9) should survive; f0–f4 (oldest) deleted. + for (let i = 0; i < 5; i++) { + expect(remaining).not.toContain(`f${i}.png`); + } + for (let i = 5; i < 10; i++) { + expect(remaining).toContain(`f${i}.png`); + } + }); + + test("nonexistent dir → logs warn, no throw", () => { + expect(() => pruneOldArtifacts(join(tmpdir(), "does-not-exist-" + randomUUID()), 10)).not.toThrow(); + }); + + test("default keep count is 200", async () => { + const { DEFAULT_ARTIFACT_KEEP_COUNT } = await import("../../src/images/artifacts"); + expect(DEFAULT_ARTIFACT_KEEP_COUNT).toBe(200); + }); + + test("maxFiles <= 0 disables pruning (does not delete everything)", () => { + const { pruneOldArtifacts } = require("../../src/images/artifacts"); + const dir = mkdtempSync(join(tmpdir(), "ocx-prune-zero-")); + for (let i = 0; i < 5; i++) writeFileSync(join(dir, `f${i}.png`), "x"); + pruneOldArtifacts(dir, 0); + expect(readdirSync(dir).length).toBe(5); + pruneOldArtifacts(dir, -1); + expect(readdirSync(dir).length).toBe(5); + }); +}); + +describe("pruneOldArtifacts: integration with materializeInlineImage", () => { + test("writing >keepCount images triggers prune down to limit", async () => { + const { materializeInlineImage } = await import("../../src/images/artifacts"); + // Use a small keepCount to keep the test fast. + const KEEP = 3; + const TOTAL = 6; + // Minimal valid base64 PNG (1x1 transparent). Decodes to 1 byte — but we need + // >0 bytes so the empty-check passes. Use "AA==" which decodes to one 0x00 byte. + const b64 = "AA=="; + const written: string[] = []; + for (let i = 0; i < TOTAL; i++) { + const path = await materializeInlineImage(b64, undefined, KEEP); + written.push(path); + // Small delay so mtimes are distinct. + await new Promise(r => setTimeout(r, 5)); + } + const dir = dirname(written[0]!); + const remaining = readdirSync(dir); + expect(remaining.length).toBe(KEEP); + // The newest KEEP files should be the last ones written. + for (let i = 0; i < TOTAL; i++) { + const fname = basename(written[i]!); + const shouldExist = i >= TOTAL - KEEP; + const exists = remaining.includes(fname); + expect(exists).toBe(shouldExist); + } + }); +}); diff --git a/tests/images/artifacts-ssrf.test.ts b/tests/images/artifacts-ssrf.test.ts new file mode 100644 index 000000000..18039f584 --- /dev/null +++ b/tests/images/artifacts-ssrf.test.ts @@ -0,0 +1,130 @@ +import { rm } from "node:fs/promises"; +import { describe, expect, mock, test } from "bun:test"; + +// Mock DNS before importing destination-policy — it binds `lookup` at load time. +const lookupMock = mock(async (_hostname: string, _opts: unknown): Promise<{ address: string; family: number }[]> => []); +mock.module("node:dns/promises", () => ({ lookup: lookupMock })); + +const { assessUrlDestination, assertUrlResolvesPublic } = await import("../../src/lib/destination-policy"); +const { downloadImageToArtifact } = await import("../../src/images/artifacts"); + +describe("SSRF: assessUrlDestination", () => { + test("loopback IPv4 → loopback", () => { + expect(assessUrlDestination("http://127.0.0.1/test")?.kind).toBe("loopback"); + }); + test("link-local → link-local", () => { + expect(assessUrlDestination("http://169.254.1.1/latest")?.kind).toBe("link-local"); + }); + test("private 10.x → private", () => { + expect(assessUrlDestination("http://10.0.0.1/test")?.kind).toBe("private"); + }); + test("private 192.168 → private", () => { + expect(assessUrlDestination("http://192.168.1.1/test")?.kind).toBe("private"); + }); + test("private 172.16 → private", () => { + expect(assessUrlDestination("http://172.16.0.1/test")?.kind).toBe("private"); + }); + test("metadata endpoint → metadata", () => { + expect(assessUrlDestination("http://169.254.170.2/test")?.kind).toBe("metadata"); + }); + test("localhost → localhost", () => { + expect(assessUrlDestination("http://localhost/test")?.kind).toBe("localhost"); + }); + test("public HTTPS → hostname or public", () => { + const kind = assessUrlDestination("https://example.com/image.png")?.kind; + expect(kind === "hostname" || kind === "public").toBe(true); + }); + test("public IP → public", () => { + expect(assessUrlDestination("https://8.8.8.8/image.png")?.kind).toBe("public"); + }); + test("IPv4-mapped IPv6 dotted-decimal [::ffff:127.0.0.1] → loopback", () => { + expect(assessUrlDestination("https://[::ffff:127.0.0.1]/image.png")?.kind).toBe("loopback"); + }); + test("IPv4-mapped IPv6 hex [::ffff:7f00:1] → loopback", () => { + expect(assessUrlDestination("https://[::ffff:7f00:1]/image.png")?.kind).toBe("loopback"); + }); + test("IPv4-mapped IPv6 hex private [::ffff:0a00:1] (10.0.0.1) → private", () => { + expect(assessUrlDestination("https://[::ffff:0a00:1]/image.png")?.kind).toBe("private"); + }); + test("invalid URL → null", () => { + expect(assessUrlDestination("not a url")).toBeNull(); + }); +}); + +describe("SSRF: assertUrlResolvesPublic", () => { + test("loopback IP → throws", async () => { + await expect(assertUrlResolvesPublic("http://127.0.0.1/x")).rejects.toThrow(); + }); + test("metadata endpoint → throws", async () => { + await expect(assertUrlResolvesPublic("http://169.254.169.254/x")).rejects.toThrow(); + }); + test("private 10.x → throws", async () => { + await expect(assertUrlResolvesPublic("http://10.0.0.1/x")).rejects.toThrow(); + }); + test("invalid URL → throws", async () => { + await expect(assertUrlResolvesPublic("not-a-url")).rejects.toThrow(); + }); +}); + +describe("SSRF: downloadImageToArtifact scheme enforcement", () => { + test("http:// → rejects (non-HTTPS)", async () => { + await expect(downloadImageToArtifact("http://public-host/path")).rejects.toThrow(/HTTPS/); + }); + + test("ftp:// → rejects", async () => { + await expect(downloadImageToArtifact("ftp://host/path")).rejects.toThrow(/HTTPS/); + }); + + test("gopher:// → rejects (non-HTTPS)", async () => { + await expect(downloadImageToArtifact("gopher://host/path")).rejects.toThrow(/HTTPS/); + }); + + test("private 10.x via download helper → rejects", async () => { + await expect(downloadImageToArtifact("https://10.0.0.1/img.png")).rejects.toThrow(); + }); + + test("IPv4-mapped IPv6 [::ffff:127.0.0.1] via download helper → rejects", async () => { + await expect(downloadImageToArtifact("https://[::ffff:127.0.0.1]/image.png")).rejects.toThrow(); + }); + + test("IPv4-mapped IPv6 hex [::ffff:7f00:1] via download helper → rejects", async () => { + await expect(downloadImageToArtifact("https://[::ffff:7f00:1]/image.png")).rejects.toThrow(); + }); + + test(`3xx redirect response → rejects (redirect: 'error')`, async () => { + lookupMock.mockResolvedValue([{ address: "93.184.216.34", family: 4 }]); + const originalFetch = globalThis.fetch; + try { + globalThis.fetch = (async (_input: RequestInfo | URL, _init?: RequestInit) => { + return new Response("", { status: 301, headers: { Location: "https://evil.example/redirect" } }); + }) as typeof fetch; + await expect(downloadImageToArtifact("https://public-host/redirect-img")).rejects.toThrow(); + } finally { + globalThis.fetch = originalFetch; + lookupMock.mockClear(); + } + }); + + test("https:// public host → succeeds with mocked fetch", async () => { + // Stub DNS so public-host resolves to a public address. + lookupMock.mockResolvedValue([{ address: "93.184.216.34", family: 4 }]); + const originalFetch = globalThis.fetch; + let downloadedPath: string | undefined; + try { + globalThis.fetch = (async (input: RequestInfo | URL, _init?: RequestInit) => { + const raw = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; + if (!raw.startsWith("https://")) throw new Error("fetch must only be called over HTTPS"); + // Minimal PNG signature so guessExtFromMagic returns "png". + const pngBytes = new Uint8Array([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]); + return new Response(pngBytes, { status: 200 }); + }) as typeof fetch; + + downloadedPath = await downloadImageToArtifact("https://public-host/valid-image"); + expect(downloadedPath).toMatch(/dl-.*\.png$/); + } finally { + globalThis.fetch = originalFetch; + lookupMock.mockClear(); + if (downloadedPath) await rm(downloadedPath).catch(() => {}); + } + }); +}); diff --git a/tests/images/fulfill.test.ts b/tests/images/fulfill.test.ts new file mode 100644 index 000000000..67e06b62b --- /dev/null +++ b/tests/images/fulfill.test.ts @@ -0,0 +1,146 @@ +import { afterAll, beforeAll, describe, expect, mock, test } from "bun:test"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { randomUUID } from "node:crypto"; +import type { ImageBridgePlan } from "../../src/images/types"; +import type { XaiImageRequest } from "../../src/images/xai-client"; + +const PREV_HOME = process.env.OPENCODEX_HOME; +beforeAll(() => { process.env.OPENCODEX_HOME = join(tmpdir(), "ocx-test-" + randomUUID()); }); +afterAll(() => { if (PREV_HOME === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = PREV_HOME; }); + +// --- Mutable mock state (reset() restores defaults before each test) --- +let xaiResult: { images: Array<{ b64_json?: string; url?: string }> } = { images: [{ b64_json: "dGVzdA==" }] }; +let xaiError: Error | null = null; +const xaiCalls: XaiImageRequest[] = []; +let matIdx = 0; +let dlIdx = 0; +let materializeFn: (i: number) => Promise = async (i) => `/test/img-${i}.png`; +let downloadFn: (i: number) => Promise = async (i) => `/test/dl-${i}.png`; + +mock.module("../../src/images/xai-client", () => ({ + callXaiImages: async (req: XaiImageRequest) => { xaiCalls.push(req); if (xaiError) throw xaiError; return xaiResult; }, +})); +mock.module("../../src/images/artifacts", () => ({ + createImageBudget: () => ({ spent: 0 }), + materializeInlineImage: async () => materializeFn(matIdx++), + downloadImageToArtifact: async () => downloadFn(dlIdx++), +})); + +const { fulfillImageCall } = await import("../../src/images/fulfill"); + +const plan = { + provider: {} as never, + auth: { baseUrl: "https://api.x.ai", token: "test-token" }, + model: "grok-imagine-image-quality", + toolNames: new Set(["image_gen"]), +} as ImageBridgePlan; + +function reset(): void { + xaiResult = { images: [{ b64_json: "dGVzdA==" }] }; + xaiError = null; + xaiCalls.length = 0; + matIdx = 0; + dlIdx = 0; + materializeFn = async (i) => `/test/img-${i}.png`; + downloadFn = async (i) => `/test/dl-${i}.png`; +} + +describe("fulfillImageCall", () => { + test("valid args → ok:true with file", async () => { + reset(); + const r = await fulfillImageCall( + { id: "c1", name: "image_gen", arguments: JSON.stringify({ prompt: "a cat", n: 2 }) }, + plan, { spent: 0 }, + ); + expect(r.ok).toBe(true); + expect(r.files.length).toBe(1); + }); + + test("missing prompt → ok:false 'missing prompt'", async () => { + reset(); + const r = await fulfillImageCall({ id: "c1", name: "image_gen", arguments: "{}" }, plan, { spent: 0 }); + expect(r.ok).toBe(false); + expect(r.error).toBe("missing prompt"); + }); + + test("invalid JSON args → ok:false 'invalid arguments JSON'", async () => { + reset(); + const r = await fulfillImageCall({ id: "c1", name: "image_gen", arguments: "{bad" }, plan, { spent: 0 }); + expect(r.ok).toBe(false); + expect(r.error).toBe("invalid arguments JSON"); + }); + + test("xAI throws → ok:false with error message", async () => { + reset(); + xaiError = new Error("xAI images API returned 500"); + const r = await fulfillImageCall( + { id: "c1", name: "image_gen", arguments: JSON.stringify({ prompt: "x" }) }, plan, { spent: 0 }, + ); + expect(r.ok).toBe(false); + expect(r.error).toContain("500"); + }); + + test("b64_json result → materialized via materializeInlineImage", async () => { + reset(); + xaiResult = { images: [{ b64_json: "dGVzdA==" }] }; + await fulfillImageCall({ id: "c1", name: "image_gen", arguments: `{"prompt":"x"}` }, plan, { spent: 0 }); + expect(matIdx).toBe(1); + expect(dlIdx).toBe(0); + }); + + test("URL result → materialized via downloadImageToArtifact", async () => { + reset(); + xaiResult = { images: [{ url: "https://cdn.example.com/i.png" }] }; + await fulfillImageCall({ id: "c1", name: "image_gen", arguments: `{"prompt":"x"}` }, plan, { spent: 0 }); + expect(dlIdx).toBe(1); + expect(matIdx).toBe(0); + }); + + test("all images fail → ok:false", async () => { + reset(); + materializeFn = async () => { throw new Error("disk full"); }; + const r = await fulfillImageCall({ id: "c1", name: "image_gen", arguments: `{"prompt":"x"}` }, plan, { spent: 0 }); + expect(r.ok).toBe(false); + expect(r.error).toContain("no usable images"); + }); + + test("one of two images fails → ok:true with 1 file", async () => { + reset(); + xaiResult = { images: [{ b64_json: "AAA=" }, { b64_json: "QkI=" }] }; + materializeFn = async (i) => { if (i === 1) throw new Error("partial fail"); return `/test/img-${i}.png`; }; + const r = await fulfillImageCall({ id: "c1", name: "image_gen", arguments: `{"prompt":"x"}` }, plan, { spent: 0 }); + expect(r.ok).toBe(true); + expect(r.files.length).toBe(1); + }); + + test("forwards prompt, model, and n to callXaiImages", async () => { + reset(); + await fulfillImageCall( + { id: "c1", name: "image_gen", arguments: JSON.stringify({ prompt: "a cat", n: 2 }) }, + plan, { spent: 0 }, + ); + expect(xaiCalls.length).toBe(1); + expect(xaiCalls[0]!.prompt).toBe("a cat"); + expect(xaiCalls[0]!.model).toBe(plan.model); + expect(xaiCalls[0]!.n).toBe(2); + }); + + test("clamps n > 4 down to 4", async () => { + reset(); + await fulfillImageCall( + { id: "c1", name: "image_gen", arguments: JSON.stringify({ prompt: "x", n: 100 }) }, + plan, { spent: 0 }, + ); + expect(xaiCalls[0]!.n).toBe(4); + }); + + test("forwards imageUrl from image_url arg", async () => { + reset(); + await fulfillImageCall( + { id: "c1", name: "image_gen", arguments: JSON.stringify({ prompt: "x", image_url: "https://example.com/i.png" }) }, + plan, { spent: 0 }, + ); + expect(xaiCalls[0]!.imageUrl).toBe("https://example.com/i.png"); + }); +}); diff --git a/tests/images/handler-activation.test.ts b/tests/images/handler-activation.test.ts new file mode 100644 index 000000000..64152f404 --- /dev/null +++ b/tests/images/handler-activation.test.ts @@ -0,0 +1,200 @@ +import { afterAll, beforeAll, describe, expect, mock, test } from "bun:test"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { randomUUID } from "node:crypto"; +import type { OcxConfig, OcxProviderConfig } from "../../src/types"; +import type { ProviderAdapter } from "../../src/adapters/base"; + +/** + * Dispatch-priority regression test for the image bridge (PR #424). + * + * The image bridge and the web-search sidecar are both opt-in dispatch paths in + * handleResponses(). The design contract is "image defers to web-search": when a + * request is eligible for BOTH, the web-search sidecar wins and the image bridge + * must NOT activate. This was previously broken because planImageBridge ran and + * returned before planWebSearch was ever consulted. + * + * These tests drive handleResponses() end-to-end (real parser + real routing + + * real planImageBridge) with only the adapter, the runners, and the web-search + * planner stubbed, so they exercise the actual dispatch ordering in + * src/server/responses/core.ts. + * + * NOTE: Full server-level integration testing of every adapter path is out of + * scope here — the focus is the dispatch priority ordering at the planImageBridge + * / planWebSearch fork (core.ts ~L1516). + */ + +const PREV_HOME = process.env.OPENCODEX_HOME; +beforeAll(() => { process.env.OPENCODEX_HOME = join(tmpdir(), "ocx-test-" + randomUUID()); }); +afterAll(() => { if (PREV_HOME === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = PREV_HOME; }); + +// --- Activation spies, flipped by the stubbed runners --- +let imageBridgeRun = false; +let webSearchRun = false; +/** Whether the stubbed adapter should expose runTurn (simulates Cursor-style adapters). */ +let useRunTurnAdapter = false; +/** Spy: flipped when the stubbed runTurn is actually invoked. */ +let runTurnCalled = false; +/** Controlled return value for the stubbed planWebSearch (truthy ⇒ web-search plan active). */ +let mockWsPlan: unknown = undefined; + +// --- Stub adapter-resolve: inject a minimal adapter so no real upstream is hit --- +const actualResolver = await import("../../src/server/adapter-resolve"); +mock.module("../../src/server/adapter-resolve", () => ({ + ...actualResolver, + resolveAdapter(provider: OcxProviderConfig) { + const base = { + name: "test", + buildRequest: async () => ({ url: provider.baseUrl, method: "POST", headers: {}, body: "" }), + async fetchResponse() { + return new Response("data: {\"type\":\"done\"}\n\n", { + status: 200, headers: { "content-type": "text/event-stream" }, + }); + }, + async *parseStream() { yield { type: "done" as const }; }, + }; + if (useRunTurnAdapter) { + return { + ...base, + async runTurn(_parsed: unknown, _incoming: unknown, emit: (event: { type: string }) => void) { + runTurnCalled = true; + emit({ type: "done" }); + }, + } as ProviderAdapter; + } + return base as ProviderAdapter; + }, +})); + +// --- Stub the image bridge runner: detect activation without hitting the real loop --- +mock.module("../../src/images/loop", () => ({ + runWithImageBridge: async () => { + imageBridgeRun = true; + return new Response("data: {\"type\":\"done\"}\n\n", { + status: 200, headers: { "content-type": "text/event-stream" }, + }); + }, +})); + +// --- Stub the web-search planner + runner: control eligibility and detect activation --- +mock.module("../../src/web-search/index", () => ({ + // Re-export the symbols parser.ts imports statically (deep path → no mock there). + buildWebSearchTool: () => ({ name: "web_search", parameters: { type: "object", properties: {} } }), + WEB_SEARCH_TOOL_NAME: "web_search", + extractHostedWebSearch: (tools: unknown[]) => { + if (!Array.isArray(tools)) return undefined; + for (const t of tools) { + if (t && typeof t === "object" && (t as Record).type === "web_search") { + return { search_context_size: "medium" }; + } + } + return undefined; + }, + runWithWebSearch: async () => { + webSearchRun = true; + return new Response("data: {\"type\":\"done\"}\n\n", { + status: 200, headers: { "content-type": "text/event-stream" }, + }); + }, + planWebSearch: () => mockWsPlan, + shouldResolveOpenAiWebSearchSidecar: () => false, +})); + +const { handleResponses } = await import("../../src/server/responses"); + +/** Routed (non-OpenAI) keyed provider + an xAI provider with an API key so the real planImageBridge returns a plan. */ +function makeConfig(): OcxConfig { + return { + port: 0, + defaultProvider: "fixture", + providers: { + fixture: { adapter: "openai-chat", baseUrl: "https://fixture.test/v1", authMode: "key", apiKey: "fixture-key" }, + xai: { adapter: "openai", baseUrl: "https://api.x.ai", apiKey: "xai-test-token" }, + }, + images: { bridgeEnabled: true }, + } as OcxConfig; +} + +function post(stream: boolean, tools: unknown[]): Promise { + return handleResponses( + new Request("http://localhost/v1/responses", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ model: "fixture/model", input: "hello", stream, tools }), + }), + makeConfig(), + { model: "", provider: "" } as never, + {}, + ); +} + +describe("image bridge dispatch priority (handler activation)", () => { + test("stream=true + image_generation tool → image bridge activates and returns SSE", async () => { + imageBridgeRun = false; webSearchRun = false; mockWsPlan = undefined; + const res = await post(true, [{ type: "image_generation" }]); + expect(imageBridgeRun).toBe(true); + expect(res.headers.get("content-type")).toBe("text/event-stream"); + }); + + test("stream=false + image_generation tool → 400 (bridge requires stream=true)", async () => { + imageBridgeRun = false; webSearchRun = false; mockWsPlan = undefined; + const res = await post(false, [{ type: "image_generation" }]); + expect(res.status).toBe(400); + // The runner must not execute when the request is rejected upfront. + expect(imageBridgeRun).toBe(false); + expect((await res.text())).toContain("image bridge requires stream=true"); + }); + + test("dual-tool (image_generation + web_search), both eligible → web-search wins, image bridge deferred", async () => { + imageBridgeRun = false; webSearchRun = false; + // A truthy plan ⇒ web-search is eligible; the image bridge must defer to it. + mockWsPlan = { backend: "openai" }; + const res = await post(true, [{ type: "web_search" }, { type: "image_generation" }]); + expect(webSearchRun).toBe(true); + expect(imageBridgeRun).toBe(false); + expect(res.headers.get("content-type")).toBe("text/event-stream"); + }); + + test("routed compaction with image_generation tool → image bridge does NOT hijack compaction (#424)", async () => { + imageBridgeRun = false; webSearchRun = false; mockWsPlan = undefined; + // A routed-compaction request carries both _compactionRequest and _imageGeneration: + // compaction clears tools/_webSearch but leaves _imageGeneration, so without the + // routedCompaction guard planImageBridge would activate and return a normal Responses + // completion instead of the synthetic compaction item Codex expects. + const res = await handleResponses( + new Request("http://localhost/v1/responses", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + model: "fixture/model", + input: [{ type: "compaction_trigger" }], + stream: true, + tools: [{ type: "image_generation" }], + }), + }), + makeConfig(), + { model: "", provider: "" } as never, + {}, + ); + expect(imageBridgeRun).toBe(false); + expect(res.headers.get("content-type")).toBe("text/event-stream"); + }); + + test("dual-tool on a runTurn adapter → web-search path wins, runTurn does not eat the request (#424)", async () => { + imageBridgeRun = false; webSearchRun = false; runTurnCalled = false; + useRunTurnAdapter = true; + mockWsPlan = { backend: "openai" }; + try { + const res = await post(true, [{ type: "web_search" }, { type: "image_generation" }]); + // Web-search sidecar must handle the turn. + expect(webSearchRun).toBe(true); + // Image bridge must not activate (design: image defers to web-search). + expect(imageBridgeRun).toBe(false); + // runTurn must NOT be reached — the web-search dispatch now runs before the runTurn early-return. + expect(runTurnCalled).toBe(false); + expect(res.headers.get("content-type")).toBe("text/event-stream"); + } finally { + useRunTurnAdapter = false; + } + }); +}); diff --git a/tests/images/loop.test.ts b/tests/images/loop.test.ts new file mode 100644 index 000000000..82d8e7b56 --- /dev/null +++ b/tests/images/loop.test.ts @@ -0,0 +1,183 @@ +import { afterAll, beforeAll, describe, expect, mock, test } from "bun:test"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { randomUUID } from "node:crypto"; +import type { ProviderAdapter, IncomingMeta } from "../../src/adapters/base"; +import type { AdapterEvent, OcxParsedRequest } from "../../src/types"; +import type { ImageBridgePlan, ImageCallResult } from "../../src/images/types"; + +const PREV_HOME = process.env.OPENCODEX_HOME; +beforeAll(() => { process.env.OPENCODEX_HOME = join(tmpdir(), "ocx-test-" + randomUUID()); }); +afterAll(() => { if (PREV_HOME === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = PREV_HOME; }); + +// --- Mock parseStreamWithProgress: simplify to direct delegation --- +mock.module("../../src/web-search/progress-stream", () => ({ + parseStreamWithProgress: async function* (_resp: Response, parse: (r: Response) => AsyncGenerator, _opts: unknown) { + for await (const e of parse(_resp)) yield e; + }, + RoutedModelInactivityError: class extends Error { readonly timeoutMs = 0; }, + WebSearchStreamProtocolError: class extends Error { /* */ }, +})); + +// --- Mock fulfillImageCall --- +let fulfillResult: ImageCallResult = { + ok: true, model: "grok-imagine-image-quality", prompt: "a cat", + files: ["/test/img.png"], count: 1, markdown: "![image](/test/img.png)", +}; +mock.module("../../src/images/fulfill", () => ({ + fulfillImageCall: async (): Promise => fulfillResult, +})); + +const { runWithImageBridge } = await import("../../src/images/loop"); + +// --- Mock adapter: yields canned events per iteration from a queue --- +let streamQueue: AdapterEvent[][] = []; +let buildRequestCalls = 0; +const mockAdapter: ProviderAdapter = { + name: "test", + buildRequest: async () => { buildRequestCalls++; return { url: "https://test/v1/chat", method: "POST", headers: {}, body: "{}" }; }, + fetchResponse: async () => new Response("{}", { status: 200, headers: { "content-type": "application/json" } }), + parseStream: async function* (): AsyncGenerator { + const events = streamQueue.shift(); + if (events) for (const e of events) yield e; + }, +}; + +const plan = { + provider: {} as never, + auth: { baseUrl: "https://api.x.ai", token: "test-token" }, + model: "grok-imagine-image-quality", + toolNames: new Set(["image_gen"]), +} as ImageBridgePlan; + +function makeParsed(): OcxParsedRequest { + return { modelId: "test-model", context: { messages: [], tools: [] }, stream: true, options: {} } as OcxParsedRequest; +} + +const imageCallEvents: AdapterEvent[] = [ + { type: "tool_call_start", id: "call_1", name: "image_gen" }, + { type: "tool_call_delta", arguments: '{"prompt":"a cat"}' }, + { type: "tool_call_end" }, + { type: "done" }, +]; + +async function runAndGetSSE(streams: AdapterEvent[][], fulfill?: ImageCallResult): Promise { + streamQueue = streams.map(s => [...s]); + if (fulfill) fulfillResult = fulfill; + const response = await runWithImageBridge({ parsed: makeParsed(), adapter: mockAdapter, plan }); + return await response.text(); +} + +describe("runWithImageBridge", () => { + test("no image tool call → passthrough text + done", async () => { + const sse = await runAndGetSSE([ + [{ type: "text_delta", text: "hello world" }, { type: "done" }], + ]); + expect(sse).toContain("hello world"); + }); + + test("single image call → fulfilled, second iteration yields text", async () => { + const sse = await runAndGetSSE( + [imageCallEvents, [{ type: "text_delta", text: "Here is your image" }, { type: "done" }]], + { ok: true, model: "grok-imagine-image-quality", prompt: "a cat", files: ["/test/img.png"], count: 1, markdown: "![image](/test/img.png)" }, + ); + expect(sse).toContain("Here is your image"); + }); + + test("fulfillImageCall error → model responds about failure", async () => { + const sse = await runAndGetSSE( + [imageCallEvents, [{ type: "text_delta", text: "Sorry, image generation failed" }, { type: "done" }]], + { ok: false, model: "grok-imagine-image-quality", prompt: "a cat", files: [], count: 0, error: "xAI unreachable" }, + ); + expect(sse).toContain("Sorry, image generation failed"); + }); + + test("image_gen tool call is intercepted — not visible in client SSE", async () => { + const sse = await runAndGetSSE( + [imageCallEvents, [{ type: "text_delta", text: "done" }, { type: "done" }]], + ); + // The tool_call_start event for image_gen should NOT appear in client-facing SSE + expect(sse).not.toContain("image_gen"); + expect(sse).not.toContain("tool_call_start"); + }); + + test("maxRounds: 1 bounds upstream requests and forces final after limit", async () => { + buildRequestCalls = 0; + // Round 0: model calls image_gen (within limit → fulfill + loop) + // Round 1: forced-final pass (forceFinal, image tools stripped from request) + streamQueue = [ + [...imageCallEvents], + [{ type: "text_delta" as const, text: "final answer" }, { type: "done" as const }], + ]; + const response = await runWithImageBridge({ parsed: makeParsed(), adapter: mockAdapter, plan, maxRounds: 1 }); + const sse = await response.text(); + expect(sse).toContain("final answer"); + // Exactly 2 upstream requests: round 0 (image call) + round 1 (forced final) + expect(buildRequestCalls).toBe(2); + }); + + test("maxRounds: 0 forces final immediately — no image tool offered", async () => { + buildRequestCalls = 0; + streamQueue = [ + [{ type: "text_delta" as const, text: "direct answer" }, { type: "done" as const }], + ]; + const response = await runWithImageBridge({ parsed: makeParsed(), adapter: mockAdapter, plan, maxRounds: 0 }); + const sse = await response.text(); + expect(sse).toContain("direct answer"); + // Single upstream request — first iteration is already forced-final + expect(buildRequestCalls).toBe(1); + }); +}); + +// --------------------------------------------------------------------------- +// runTurn adapter path (Cursor) — events arrive via an emit callback, not +// buildRequest/fetchResponse/parseStream. +// --------------------------------------------------------------------------- + +describe("runWithImageBridge — runTurn adapter", () => { + let runTurnEventQueue: AdapterEvent[][] = []; + const runTurnAdapter: ProviderAdapter = { + ...mockAdapter, + runTurn: async (_parsed: OcxParsedRequest, _incoming: IncomingMeta, emit: (e: AdapterEvent) => void) => { + const events = runTurnEventQueue.shift(); + if (events) for (const e of events) emit(e); + }, + }; + + test("runTurn adapter → image call intercepted and fulfilled", async () => { + runTurnEventQueue = [ + [...imageCallEvents], + [{ type: "text_delta", text: "Here is your image" }, { type: "done" }], + ]; + fulfillResult = { + ok: true, model: "grok-imagine-image-quality", prompt: "a cat", + files: ["/test/img.png"], count: 1, markdown: "![image](/test/img.png)", + }; + const response = await runWithImageBridge({ + parsed: makeParsed(), adapter: runTurnAdapter, plan, maxRounds: 1, + }); + const sse = await response.text(); + expect(sse).toContain("Here is your image"); + // The synthetic image_gen tool call must NOT leak to the client + expect(sse).not.toContain("image_gen"); + expect(sse).not.toContain("tool_call_start"); + }); + + test("runTurn adapter → text passthrough (no image call)", async () => { + runTurnEventQueue = [ + [{ type: "text_delta", text: "hello from runTurn" }, { type: "done" }], + ]; + const response = await runWithImageBridge({ parsed: makeParsed(), adapter: runTurnAdapter, plan }); + const sse = await response.text(); + expect(sse).toContain("hello from runTurn"); + }); + + test("runTurn adapter → error event surfaces as upstream failure", async () => { + runTurnEventQueue = [ + [{ type: "error", message: "cursor blew up" }], + ]; + const response = await runWithImageBridge({ parsed: makeParsed(), adapter: runTurnAdapter, plan }); + const sse = await response.text(); + expect(sse).toContain("cursor blew up"); + }); +}); diff --git a/tests/images/plan.test.ts b/tests/images/plan.test.ts new file mode 100644 index 000000000..85a536ee6 --- /dev/null +++ b/tests/images/plan.test.ts @@ -0,0 +1,143 @@ +import { afterAll, beforeAll, describe, expect, mock, test } from "bun:test"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { randomUUID } from "node:crypto"; +import type { OcxConfig, OcxProviderConfig, OcxParsedRequest } from "../../src/types"; + +const PREV_HOME = process.env.OPENCODEX_HOME; +beforeAll(() => { process.env.OPENCODEX_HOME = join(tmpdir(), "ocx-test-" + randomUUID()); }); +afterAll(() => { if (PREV_HOME === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = PREV_HOME; }); + +/** Mutable token that the mocked getValidAccessToken resolves to. */ +let tokenResult: string | null = null; +mock.module("../../src/oauth/index", () => ({ + getValidAccessToken: async () => tokenResult, +})); + +const { planImageBridge } = await import("../../src/images/plan"); + +function makeConfig( + providers: Record>, + images?: { bridgeEnabled?: boolean; bridgeModel?: string }, +): OcxConfig { + return { + port: 0, + defaultProvider: "test", + providers: Object.fromEntries( + Object.entries(providers).map(([k, v]) => [k, { adapter: "openai", baseUrl: "https://api.test.com", ...v }]), + ), + ...(images ? { images } : {}), + } as OcxConfig; +} + +function makeParsed(withImageGen: boolean): OcxParsedRequest { + return { + modelId: "test-model", + context: { messages: [], tools: [] }, + stream: true, + options: {}, + ...(withImageGen ? { _imageGeneration: { toolNames: new Set(["image_gen"]) } } : {}), + } as OcxParsedRequest; +} + +const routed = { adapter: "openai", baseUrl: "https://api.anthropic.com" } as OcxProviderConfig; +const openaiRouted = { adapter: "openai", baseUrl: "https://api.openai.com" } as OcxProviderConfig; + +describe("planImageBridge", () => { + test("bridgeEnabled false → undefined", async () => { + expect(await planImageBridge(makeConfig({ test: routed }, { bridgeEnabled: false }), makeParsed(true), routed)).toBeUndefined(); + }); + + test("bridgeEnabled not set → undefined (opt-in required)", async () => { + // xAI provider configured but images.bridgeEnabled is absent — must not bridge. + const cfg = makeConfig({ xai: { baseUrl: "https://api.x.ai", apiKey: "test-token" } }); + expect(await planImageBridge(cfg, makeParsed(true), routed)).toBeUndefined(); + }); + + test("_imageGeneration not set → undefined", async () => { + expect(await planImageBridge(makeConfig({ test: routed }, { bridgeEnabled: true }), makeParsed(false), routed)).toBeUndefined(); + }); + + test("routedProvider is api.openai.com → undefined", async () => { + const cfg = makeConfig({ xai: { baseUrl: "https://api.x.ai", apiKey: "test-token" } }, { bridgeEnabled: true }); + expect(await planImageBridge(cfg, makeParsed(true), openaiRouted)).toBeUndefined(); + }); + + test("no xAI provider → undefined", async () => { + expect(await planImageBridge(makeConfig({ test: routed }, { bridgeEnabled: true }), makeParsed(true), routed)).toBeUndefined(); + }); + + test("xAI provider but apiKey empty and no OAuth → undefined", async () => { + tokenResult = null; + const cfg = makeConfig({ xai: { baseUrl: "https://api.x.ai", apiKey: "" } }, { bridgeEnabled: true }); + expect(await planImageBridge(cfg, makeParsed(true), routed)).toBeUndefined(); + }); + + test("xAI provider with API key → returns plan with correct model", async () => { + const cfg = makeConfig({ xai: { baseUrl: "https://api.x.ai", apiKey: "test-token" } }, { bridgeEnabled: true }); + const plan = await planImageBridge(cfg, makeParsed(true), routed); + expect(plan).toBeDefined(); + expect(plan!.model).toBe("grok-imagine-image-quality"); + expect(plan!.auth.token).toBe("test-token"); + expect(plan!.auth.baseUrl).toBe("https://api.x.ai/v1"); + }); + + test("xAI provider with OAuth (getValidAccessToken) → returns plan", async () => { + tokenResult = "fake-oauth-123"; + const cfg = makeConfig({ xai: { baseUrl: "https://api.x.ai" } }, { bridgeEnabled: true }); + const plan = await planImageBridge(cfg, makeParsed(true), routed); + expect(plan).toBeDefined(); + expect(plan!.auth.token).toBe("fake-oauth-123"); + tokenResult = null; + }); + + test("custom-named provider with api.x.ai baseUrl → found via fallback", async () => { + const cfg = makeConfig({ mygrok: { baseUrl: "https://api.x.ai", apiKey: "test-token" } }, { bridgeEnabled: true }); + const plan = await planImageBridge(cfg, makeParsed(true), routed); + expect(plan).toBeDefined(); + expect(plan!.provider).toBe(cfg.providers.mygrok); + }); + + test("custom bridgeModel is honored", async () => { + const cfg = makeConfig( + { xai: { baseUrl: "https://api.x.ai", apiKey: "test-token" } }, + { bridgeEnabled: true, bridgeModel: "custom-img-model" }, + ); + expect((await planImageBridge(cfg, makeParsed(true), routed))!.model).toBe("custom-img-model"); + }); + + test("toolNames includes IMAGE_GEN_TOOL_NAME so the loop can intercept synthetic calls", async () => { + const { IMAGE_GEN_TOOL_NAME } = await import("../../src/images/synthetic-tool"); + const cfg = makeConfig({ xai: { baseUrl: "https://api.x.ai", apiKey: "test-token" } }, { bridgeEnabled: true }); + const plan = await planImageBridge(cfg, makeParsed(true), routed); + expect(plan).toBeDefined(); + // The plan always merges in IMAGE_GEN_TOOL_NAME, even if _imageGeneration.toolNames + // only contained the original hosted tool name. + expect(plan!.toolNames.has(IMAGE_GEN_TOOL_NAME)).toBe(true); + }); + + test("baseUrl is pinned to registry regardless of config override", async () => { + // config 里 xai provider 的 baseUrl 被改成恶意 host + const cfg = makeConfig( + { xai: { adapter: "openai-chat", baseUrl: "https://evil.example.com/v1", apiKey: "test-key" } }, + { bridgeEnabled: true }, + ); + const plan = await planImageBridge(cfg, makeParsed(true), routed); + expect(plan).toBeDefined(); + // auth.baseUrl 必须是 registry pin 的地址,不是 config 里的恶意地址 + expect(plan!.auth.baseUrl).toBe("https://api.x.ai/v1"); + }); + + test("custom-named provider with api.x.ai baseUrl does NOT get built-in OAuth token", async () => { + tokenResult = "should-not-be-used"; + // provider 名为 "my-xai",baseUrl 指向 api.x.ai,没有 apiKey + const cfg = makeConfig( + { "my-xai": { adapter: "openai-chat", baseUrl: "https://api.x.ai/v1" } }, + { bridgeEnabled: true }, + ); + const plan = await planImageBridge(cfg, makeParsed(true), routed); + // 没有 apiKey 也没有 "xai" 的 OAuth → 没有 token → 没有 plan + expect(plan).toBeUndefined(); + tokenResult = null; + }); +}); diff --git a/tests/images/synthetic-tool.test.ts b/tests/images/synthetic-tool.test.ts new file mode 100644 index 000000000..7572c94c1 --- /dev/null +++ b/tests/images/synthetic-tool.test.ts @@ -0,0 +1,62 @@ +import { describe, expect, test } from "bun:test"; +import { isImageGenName, extractHostedImageGeneration, buildImageTool } from "../../src/images/synthetic-tool"; + +describe("isImageGenName", () => { + test("'image_gen' → true", () => { + expect(isImageGenName("image_gen")).toBe(true); + }); + + test("'IMAGE_GENERATION' → true (case insensitive)", () => { + expect(isImageGenName("IMAGE_GENERATION")).toBe(true); + }); + + test("'imagegen' → true", () => { + expect(isImageGenName("imagegen")).toBe(true); + }); + + test("'not_image' → false", () => { + expect(isImageGenName("not_image")).toBe(false); + }); +}); + +describe("extractHostedImageGeneration", () => { + test("type 'image_generation' → returns toolNames with that name", () => { + const result = extractHostedImageGeneration([{ type: "image_generation" }]); + expect(result).toBeDefined(); + expect(result!.toolNames.has("image_generation")).toBe(true); + }); + + test("flat Responses function tool with 'image_gen' name → returns with that name", () => { + const result = extractHostedImageGeneration([ + { type: "function", name: "image_gen", parameters: { type: "object" } }, + ]); + expect(result).toBeDefined(); + expect(result!.toolNames.has("image_gen")).toBe(true); + }); + + test("nested Chat Completions function tool with 'image_gen' name → returns with that name", () => { + const result = extractHostedImageGeneration([ + { type: "function", function: { name: "image_gen" } }, + ]); + expect(result).toBeDefined(); + expect(result!.toolNames.has("image_gen")).toBe(true); + }); + + test("no matching tools → undefined", () => { + expect( + extractHostedImageGeneration([{ type: "function", function: { name: "shell" } }]), + ).toBeUndefined(); + }); + + test("undefined → undefined", () => { + expect(extractHostedImageGeneration(undefined)).toBeUndefined(); + }); +}); + +describe("buildImageTool", () => { + test("has name 'image_gen' and imageGeneration flag", () => { + const tool = buildImageTool(); + expect(tool.name).toBe("image_gen"); + expect(tool.imageGeneration).toBe(true); + }); +}); diff --git a/tests/images/xai-client.test.ts b/tests/images/xai-client.test.ts new file mode 100644 index 000000000..6c8a6cdc7 --- /dev/null +++ b/tests/images/xai-client.test.ts @@ -0,0 +1,114 @@ +import { afterAll, afterEach, beforeAll, describe, expect, test } from "bun:test"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { randomUUID } from "node:crypto"; +import { callXaiImages } from "../../src/images/xai-client"; + +const PREV_HOME = process.env.OPENCODEX_HOME; +beforeAll(() => { process.env.OPENCODEX_HOME = join(tmpdir(), "ocx-test-" + randomUUID()); }); +afterAll(() => { if (PREV_HOME === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = PREV_HOME; }); + +const AUTH = { baseUrl: "https://api.x.ai", token: "test-token" }; +const originalFetch = globalThis.fetch; +afterEach(() => { globalThis.fetch = originalFetch; }); + +/** Replace globalThis.fetch with a stub that captures the request and returns a canned response. */ +function stubFetch(status: number, body: unknown): { url: string; init?: RequestInit }[] { + const calls: { url: string; init?: RequestInit }[] = []; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + calls.push({ url: input.toString(), init }); + return new Response(JSON.stringify(body), { + status, + headers: { "content-type": "application/json" }, + }); + }) as typeof fetch; + return calls; +} + +describe("callXaiImages", () => { + test("no imageUrl → POST /images/generations", async () => { + const calls = stubFetch(200, { data: [{ b64_json: "dGVzdA==" }] }); + await callXaiImages({ prompt: "a cat" }, AUTH); + expect(calls[0]!.url).toContain("/images/generations"); + expect(calls[0]!.init?.method).toBe("POST"); + }); + + test("with imageUrl → POST /images/edits", async () => { + const calls = stubFetch(200, { data: [{ b64_json: "dGVzdA==" }] }); + await callXaiImages({ prompt: "edit this", imageUrl: "https://example.com/img.png" }, AUTH); + expect(calls[0]!.url).toContain("/images/edits"); + }); + + test("request body has correct model, prompt, n", async () => { + const calls = stubFetch(200, { data: [{ b64_json: "dGVzdA==" }] }); + await callXaiImages({ prompt: "a dog", model: "grok-imagine-fast", n: 3 }, AUTH); + const body = JSON.parse((calls[0]!.init?.body as string) ?? "{}"); + expect(body.model).toBe("grok-imagine-fast"); + expect(body.prompt).toBe("a dog"); + expect(body.n).toBe(3); + }); + + test("non-2xx → throws Error containing status code", async () => { + stubFetch(429, { error: "rate limited" }); + await expect(callXaiImages({ prompt: "x" }, AUTH)).rejects.toThrow("429"); + }); + + test("2xx with b64_json → returns normalized XaiImageResult", async () => { + stubFetch(200, { data: [{ b64_json: "dGVzdA==" }] }); + const result = await callXaiImages({ prompt: "x" }, AUTH); + expect(result.images.length).toBe(1); + expect(result.images[0]!.b64_json).toBe("dGVzdA=="); + }); + + test("2xx with url → returns images[0].url", async () => { + stubFetch(200, { data: [{ url: "https://cdn.example.com/img.png" }] }); + const result = await callXaiImages({ prompt: "x" }, AUTH); + expect(result.images[0]!.url).toBe("https://cdn.example.com/img.png"); + }); + + test("caller abort propagates into the composed signal", async () => { + const calls = stubFetch(200, { data: [{ b64_json: "dGVzdA==" }] }); + const controller = new AbortController(); + await callXaiImages({ prompt: "x" }, AUTH, controller.signal); + const passed = calls[0]!.init?.signal as AbortSignal; + expect(passed.aborted).toBe(false); + controller.abort("client gone"); + expect(passed.aborted).toBe(true); + }); + + test("size/quality mapped to aspect_ratio/resolution, no passthrough", async () => { + const calls = stubFetch(200, { data: [{ b64_json: "dGVzdA==" }] }); + await callXaiImages({ prompt: "x", size: "1024x1792", quality: "hd" }, AUTH); + const body = JSON.parse((calls[0]!.init?.body as string) ?? "{}"); + expect(body.aspect_ratio).toBe("9:16"); + expect(body.resolution).toBe("2k"); + expect(body).not.toHaveProperty("size"); + expect(body).not.toHaveProperty("quality"); + }); + + test("square size → 1:1", async () => { + const calls = stubFetch(200, { data: [{ b64_json: "dGVzdA==" }] }); + await callXaiImages({ prompt: "x", size: "1024x1024" }, AUTH); + const body = JSON.parse((calls[0]!.init?.body as string) ?? "{}"); + expect(body.aspect_ratio).toBe("1:1"); + expect(body).not.toHaveProperty("resolution"); + }); + + test("quality: standard → 1k", async () => { + const calls = stubFetch(200, { data: [{ b64_json: "dGVzdA==" }] }); + await callXaiImages({ prompt: "x", quality: "standard" }, AUTH); + const body = JSON.parse((calls[0]!.init?.body as string) ?? "{}"); + expect(body.resolution).toBe("1k"); + expect(body).not.toHaveProperty("aspect_ratio"); + }); + + test("unknown size/quality dropped", async () => { + const calls = stubFetch(200, { data: [{ b64_json: "dGVzdA==" }] }); + await callXaiImages({ prompt: "x", size: "weird", quality: "ultra" }, AUTH); + const body = JSON.parse((calls[0]!.init?.body as string) ?? "{}"); + expect(body).not.toHaveProperty("aspect_ratio"); + expect(body).not.toHaveProperty("resolution"); + expect(body).not.toHaveProperty("size"); + expect(body).not.toHaveProperty("quality"); + }); +}); diff --git a/tests/responses-parser.test.ts b/tests/responses-parser.test.ts index fce5c005e..7f49ce435 100644 --- a/tests/responses-parser.test.ts +++ b/tests/responses-parser.test.ts @@ -283,6 +283,20 @@ describe("codex-rs compat surface (260707)", () => { expect(parsed.options.reasoning).toBeUndefined(); }); + test("detects image_generation hosted tool arriving via additional_tools (responses_lite WS shape)", () => { + // Codex Desktop responses_websockets lite path: NO body.tools; the hosted tool spec rides + // inside an input item {type:"additional_tools", tools:[...]}. extractHostedImageGeneration + // must still see it so the image bridge activates. + const parsed = parseRequest({ + model: "p/m", + input: [ + { type: "additional_tools", tools: [{ type: "image_generation" }] }, + { type: "message", role: "user", content: [{ type: "input_text", text: "draw a cat" }] }, + ], + }); + expect(parsed._imageGeneration?.toolNames.has("image_generation")).toBe(true); + }); + test("current parser ignores null empty and unknown string efforts", () => { expect(parseRequest({ model: "p/m", input: "hi", reasoning: null }).options.reasoning).toBeUndefined(); expect(parseRequest({ model: "p/m", input: "hi", reasoning: { effort: "" } }).options.reasoning).toBeUndefined();