diff --git a/docs-site/src/content/docs/guides/codex-integration.md b/docs-site/src/content/docs/guides/codex-integration.md index cd4fe810a..44eaec610 100644 --- a/docs-site/src/content/docs/guides/codex-integration.md +++ b/docs-site/src/content/docs/guides/codex-integration.md @@ -36,7 +36,11 @@ The proxy listens on port `10100` by default and serves `POST /v1/responses`, Codex's built-in `image_gen` tool does not go through `/v1/responses` — the codex-rs extension POSTs `{base_url}/images/generations` (or `/images/edits` when reference images are attached) directly, with the same ChatGPT bearer auth it uses for chat. Because the injected `base_url` -points at opencodex, the proxy relays those calls to the OpenAI upstream: +points at opencodex, the proxy relays those calls to the OpenAI upstream. + +This is separate from the [Image Bridge](/guides/image-bridge/), which only activates when a +**Responses** turn lists the hosted `image_generation` tool while a non-OpenAI model is selected. +Standalone `/images/generations` calls never enter that bridge. - **One mode-aware forward candidate:** Pool selects an eligible main/added account; Direct uses the caller OAuth bearer. The configured mode applies consistently to the image request. 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..748b3e22f --- /dev/null +++ b/docs-site/src/content/docs/guides/image-bridge.md @@ -0,0 +1,88 @@ +--- +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 entry with credentials. The bridge pins fulfillment to the registry xAI + endpoint (`https://api.x.ai/v1`); any configured `baseUrl` override is ignored for image calls. + A minimal key-mode entry is enough: + + ```json + { + "providers": { + "xai": { "adapter": "openai-chat", "apiKey": "xai-…" } + } + } + ``` + +- Authentication via `authMode: "oauth"` (`ocx login xai` — uses a stored, auto-refreshed + bearer token) or `authMode: "key"` (a configured API key). See + [Providers](/guides/providers/) for the shared auth modes. +- 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, + "timeoutMs": 60000 + } +} +``` + +| 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 image-generation loop iterations per turn. Floored to an integer and clamped to `[0, 10]`; non-finite values fall back to `3`. | +| `timeoutMs` | `60000` | Per-call xAI deadline in milliseconds. Finite positive values are floored and passed to the xAI request. | + +## How It Works + +The Image Bridge activates only on **Responses** turns that include the hosted +`image_generation` tool in the `/v1/responses` tools array while a **non-OpenAI** +model is selected. It does **not** intercept Codex's built-in `image_gen` tool, +which POSTs directly to `/v1/images/generations` (or `/images/edits`) — that path +is covered separately in [Codex Integration](/guides/codex-integration/#built-in-image-generation-image_gen). + +1. When a Responses request lists `image_generation` in `tools`, 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** on adapters that support the web-search sidecar loop. If both web + search and image generation are requested in the same turn, web-search runs and image + generation is skipped. Cursor/`runTurn` adapters cannot use that sidecar today, so the image + bridge may still run for those dual-tool turns. +- **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..f6ba898d8 --- /dev/null +++ b/src/images/artifacts.ts @@ -0,0 +1,325 @@ +import { mkdir, writeFile } from "node:fs/promises"; +import type { IncomingMessage } from "node:http"; +import https from "node:https"; +import type { RequestOptions } from "node:https"; +import { join } from "node:path"; +import { getConfigDir } from "../config"; +import { assessUrlDestination, resolvePublicAddresses } from "../lib/destination-policy"; + +const MAX_DECODED_BYTES_PER_IMAGE = 50 * 1024 * 1024; +const MAX_DECODED_BYTES_PER_RESPONSE = 100 * 1024 * 1024; +/** Hard cap for remote image downloads (also enforced inside pinnedHttpsGet). */ +export const MAX_DOWNLOAD_BYTES = 50 * 1024 * 1024; // 50 MiB +/** Idle timeout for pinned HTTPS connect/headers/body when no AbortSignal is provided. */ +export const DOWNLOAD_IDLE_TIMEOUT_MS = 60_000; + +// 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 type PinnedAddress = { address: string; family: number }; + +/** Test seam / custom transport: must connect to `pinned`, not re-resolve `url`'s hostname. */ +export type PinnedDownloadFn = ( + url: string, + pinned: PinnedAddress, + signal?: AbortSignal, +) => Promise; + +export function createImageBudget(): ImageBudget { + return { spent: 0 }; +} + +function getArtifactsDir(): string { + return join(getConfigDir(), "artifacts"); +} + +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(""); +} + +/** Sniff a recognized image extension, or null when the payload is empty/non-image. */ +export function sniffImageExtension(bytes: Uint8Array): "png" | "jpg" | "webp" | "gif" | null { + if (bytes.byteLength === 0) return null; + const sig = Buffer.from(bytes.slice(0, 12)).toString("latin1"); + if (sig.startsWith("\x89PNG")) 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 null; +} + +export function guessExtFromMagic(bytes: Uint8Array): string { + return sniffImageExtension(bytes) ?? "png"; +} + +export async function materializeInlineImage( + base64Data: string, + budget?: ImageBudget, +): 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 = sniffImageExtension(buf); + if (!ext) throw new Error("inline image data is not a recognized image"); + const filePath = join(dir, `img-${timestampPrefix()}-${crypto.randomUUID()}.${ext}`); + await writeFile(filePath, buf, { mode: 0o600 }); + return filePath; +} + +/** + * HTTPS GET that connects to a previously validated address while keeping the + * original hostname for SNI / Host. The custom `lookup` never asks the OS + * resolver again, so a rebinding answer cannot redirect the TCP peer. + * + * Returns a streaming Response so callers can enforce byte caps while reading; + * the transport also destroys the request if `maxBytes` is exceeded mid-stream. + */ +export function pinnedHttpsGet( + url: string, + pinned: PinnedAddress, + signal?: AbortSignal, + options?: { + maxBytes?: number; + idleTimeoutMs?: number; + rejectUnauthorized?: boolean; + }, +): Promise { + const parsed = new URL(url); + if (parsed.protocol !== "https:") { + throw new Error(`image URL must use HTTPS, got ${parsed.protocol}`); + } + const maxBytes = options?.maxBytes ?? MAX_DOWNLOAD_BYTES; + const idleTimeoutMs = options?.idleTimeoutMs ?? DOWNLOAD_IDLE_TIMEOUT_MS; + + return new Promise((resolve, reject) => { + if (signal?.aborted) { + reject(signal.reason instanceof Error ? signal.reason : new Error("aborted")); + return; + } + + let settled = false; + const fail = (err: unknown) => { + try { req.destroy(); } catch { /* ignore */ } + if (settled) return; + settled = true; + reject(err instanceof Error ? err : new Error(String(err))); + }; + + const optionsHttps: RequestOptions & { servername?: string } = { + protocol: "https:", + hostname: parsed.hostname, + servername: parsed.hostname, + port: parsed.port || 443, + path: `${parsed.pathname}${parsed.search}`, + method: "GET", + headers: { Host: parsed.host }, + rejectUnauthorized: options?.rejectUnauthorized, + lookup(_hostname, lookupOptions, callback) { + const opts = typeof lookupOptions === "function" ? undefined : lookupOptions; + const cb = typeof lookupOptions === "function" ? lookupOptions : callback; + if (!cb) return; + // Pin the validated peer — do not call dns.lookup again. + // Honor `{ all: true }` array shape used by some Node/Bun https paths. + if (opts && typeof opts === "object" && "all" in opts && opts.all) { + (cb as (err: NodeJS.ErrnoException | null, addresses: PinnedAddress[]) => void)( + null, + [{ address: pinned.address, family: pinned.family }], + ); + return; + } + (cb as (err: NodeJS.ErrnoException | null, address: string, family: 4 | 6) => void)( + null, + pinned.address, + pinned.family as 4 | 6, + ); + }, + }; + + const req = https.request(optionsHttps, (res: IncomingMessage) => { + const status = res.statusCode ?? 0; + // Any non-2xx must destroy immediately. Returning a streaming Response for + // 4xx/5xx (or 3xx) lets callers that only check `Response.ok` abandon an + // unread body while the peer keeps sending — a failed-response socket leak. + if (status < 200 || status >= 300) { + try { res.destroy(); } catch { /* ignore */ } + fail(new Error("image download failed: " + status)); + return; + } + const headers = new Headers(); + for (const [key, value] of Object.entries(res.headers)) { + if (value === undefined || value === null) continue; + if (Array.isArray(value)) { + for (const item of value) headers.append(key, String(item)); + } else { + headers.set(key, String(value)); + } + } + + let received = 0; + const stream = new ReadableStream({ + start(controller) { + res.setTimeout(idleTimeoutMs, () => { + fail(new Error("image download stalled")); + try { controller.error(new Error("image download stalled")); } catch { /* closed */ } + }); + res.on("data", (chunk: Buffer | string) => { + const buf = typeof chunk === "string" ? Buffer.from(chunk) : chunk; + received += buf.byteLength; + if (received > maxBytes) { + const err = new Error(`image download exceeds ${maxBytes} byte cap`); + fail(err); + try { controller.error(err); } catch { /* closed */ } + return; + } + try { controller.enqueue(buf); } catch { /* closed */ } + }); + res.on("end", () => { + try { controller.close(); } catch { /* closed */ } + }); + res.on("error", (err: Error) => { + fail(err); + try { controller.error(err); } catch { /* closed */ } + }); + }, + cancel() { + req.destroy(); + }, + }); + + if (settled) return; + settled = true; + resolve(new Response(stream, { status, headers })); + }); + + const onAbort = () => { + fail(signal?.reason instanceof Error ? signal.reason : new Error("aborted")); + }; + signal?.addEventListener("abort", onAbort, { once: true }); + req.setTimeout(idleTimeoutMs, () => { + fail(new Error("image download timed out")); + }); + req.on("error", (err) => { + signal?.removeEventListener("abort", onAbort); + fail(err); + }); + req.on("close", () => signal?.removeEventListener("abort", onAbort)); + req.end(); + }); +} + +function pickPinnedAddress(addresses: PinnedAddress[]): PinnedAddress { + return addresses.find(a => a.family === 4) ?? addresses[0]!; +} + +export async function downloadImageToArtifact( + url: string, + budget?: ImageBudget, + signal?: AbortSignal, + options?: { pinnedDownload?: PinnedDownloadFn }, +): 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); + } + + // SSRF protection: validate the provider-returned URL before fetching. + // Require HTTPS strictly — plain HTTP and all other schemes (ftp, file, …) are rejected. + // Resolve DNS once, then pin that public address for the HTTPS connect (SNI/Host keep + // the original hostname) so a rebinding answer cannot retarget the TCP peer. + 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}`); + } + const resolved = await resolvePublicAddresses(url); + const pinned = pickPinnedAddress(resolved.addresses); + const download = options?.pinnedDownload ?? pinnedHttpsGet; + const resp = await download(url, pinned, signal); + if (!resp.ok) { + // Custom `pinnedDownload` seams may still return a failed Response with a + // live body; cancel it so unread error payloads cannot keep the socket warm. + try { await resp.body?.cancel(); } catch { /* ignore */ } + 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 (bytes.byteLength === 0) throw new Error("image download returned empty body"); + const ext = sniffImageExtension(bytes); + if (!ext) throw new Error("image download did not contain a recognized image"); + + 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 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 }); + return filePath; +} diff --git a/src/images/fulfill.ts b/src/images/fulfill.ts new file mode 100644 index 000000000..62d214d2f --- /dev/null +++ b/src/images/fulfill.ts @@ -0,0 +1,81 @@ +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 : plan.defaultSize; + const quality = typeof obj.quality === "string" ? obj.quality : plan.defaultQuality; + + let result; + try { + result = await callXaiImages( + { prompt, model: plan.model, n, imageUrl, size, quality }, + plan.auth, + signal, + plan.timeoutMs, + ); + } 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)); + } else if (img.url) { + files.push(await downloadImageToArtifact(img.url, budget, signal)); + } + } 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..9a402706b --- /dev/null +++ b/src/images/index.ts @@ -0,0 +1,4 @@ +export { planImageBridge, findXaiProvider, resolveXaiToken } from "./plan"; +export { runWithImageBridge, clampImageMaxRounds, DEFAULT_MAX_ROUNDS, MAX_ROUNDS_HARD_LIMIT } 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..9a3561831 --- /dev/null +++ b/src/images/loop.ts @@ -0,0 +1,615 @@ +/** + * 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 forced-answer nudge, no failed-query + * dedup, no describeImages/structuredOutput, no recordSidecarOutcome. + */ +import type { AdapterRequest, ProviderAdapter } from "../adapters/base"; +import { createAdapterEventQueue } from "../adapters/run-turn-queue"; +import type { AdapterEvent, OcxMessage, OcxParsedRequest, OcxProviderContinuationState, OcxThinkingContent, OcxUsage } from "../types"; +import { namespacedToolName } from "../types"; +import { bridgeToResponsesSSE } from "../bridge"; +import { clearableDeadline, idleDeadline } 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; +export const DEFAULT_MAX_ROUNDS = 3; +/** Absolute ceiling so a hand-edited `images.maxRounds: 10000` cannot unbound paid xAI calls. */ +export const MAX_ROUNDS_HARD_LIMIT = 10; +/** Cap paid xAI fulfillments per turn (parallel calls in one round count separately). */ +export const MAX_IMAGE_CALLS_PER_TURN = 10; + +/** + * Clamp a configured maxRounds value to a safe integer in [0, MAX_ROUNDS_HARD_LIMIT]. + * Non-finite / non-number inputs fall back to DEFAULT_MAX_ROUNDS. + */ +export function clampImageMaxRounds(value: unknown): number { + if (typeof value !== "number" || !Number.isFinite(value)) return DEFAULT_MAX_ROUNDS; + return Math.max(0, Math.min(MAX_ROUNDS_HARD_LIMIT, Math.floor(value))); +} + +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) return; + if (toolNames.has(pending.name)) { + // Unterminated image call still carries buffered args — fulfill so malformed JSON + // becomes a normal tool_result error instead of silently vanishing. + calls.push({ id: pending.id, name: pending.name, args: pending.argsBuf }); + } else { + 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; + /** Called after each upstream request is built (parity with web-search / normal path). */ + onRequestBuilt?: (request: AdapterRequest) => void; + abortSignal?: AbortSignal; + onFirstOutput?: () => void; + /** Max image-generation rounds before forcing a final answer. Defaults to 3; clamped to [0, 10]. */ + maxRounds?: number; + /** Connect / response-header budget for non-runTurn iterations. */ + connectTimeoutMs?: number; + /** Stall budget (seconds) forwarded to bridgeToResponsesSSE; also bounds runTurn collect. */ + stallTimeoutSec?: number; + /** Provider-specific fetch (e.g. xAI transport wrapper). Falls back to global fetch. */ + fetchImpl?: typeof globalThis.fetch; + /** Raw adapter usage at the terminal event, pre wire-normalization (see bridgeToResponsesSSE onUsage). */ + onUsage?: (usage: OcxUsage | undefined) => void; + /** + * Optional 429 key-failover for the routed (non-xAI) model. Return a rebuilt adapter for the + * rotated key, or null when the pool is exhausted. + */ + on429?: (retryAfterHeader: string | null) => ProviderAdapter | null; + /** Called when the bridged Responses stream completes (parity with runTurn / routed paths). */ + onCompletedResponse?: (response: Record, providerState?: OcxProviderContinuationState) => void; +} + +/** + * 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, plan, abortSignal } = deps; + let adapter = deps.adapter; + const maxRounds = clampImageMaxRounds(deps.maxRounds ?? DEFAULT_MAX_ROUNDS); + const HARD_CAP = maxRounds + 1; + const connectTimeoutMs = typeof deps.connectTimeoutMs === "number" && Number.isFinite(deps.connectTimeoutMs) && deps.connectTimeoutMs > 0 + ? Math.floor(deps.connectTimeoutMs) + : CONNECT_TIMEOUT_MS; + const stallTimeoutMs = typeof deps.stallTimeoutSec === "number" && Number.isFinite(deps.stallTimeoutSec) && deps.stallTimeoutSec > 0 + ? Math.floor(deps.stallTimeoutSec * 1000) + : STALL_TIMEOUT_MS; + const fetchImpl = deps.fetchImpl ?? globalThis.fetch; + let paidImageCalls = 0; + let hiddenUsage: OcxUsage | undefined; + + const addUsage = (a: OcxUsage | undefined, b: OcxUsage | undefined): OcxUsage | undefined => { + if (!a) return b; + if (!b) return a; + return { + inputTokens: a.inputTokens + b.inputTokens, + outputTokens: a.outputTokens + b.outputTokens, + ...(a.contextTotalTokens !== undefined || b.contextTotalTokens !== undefined + ? { contextTotalTokens: Math.max(a.contextTotalTokens ?? 0, b.contextTotalTokens ?? 0) } + : {}), + ...(a.cachedInputTokens !== undefined || b.cachedInputTokens !== undefined + ? { cachedInputTokens: (a.cachedInputTokens ?? 0) + (b.cachedInputTokens ?? 0) } + : {}), + ...(a.cacheReadInputTokens !== undefined || b.cacheReadInputTokens !== undefined + ? { cacheReadInputTokens: (a.cacheReadInputTokens ?? 0) + (b.cacheReadInputTokens ?? 0) } + : {}), + ...(a.cacheCreationInputTokens !== undefined || b.cacheCreationInputTokens !== undefined + ? { cacheCreationInputTokens: (a.cacheCreationInputTokens ?? 0) + (b.cacheCreationInputTokens ?? 0) } + : {}), + ...(a.reasoningOutputTokens !== undefined || b.reasoningOutputTokens !== undefined + ? { reasoningOutputTokens: (a.reasoningOutputTokens ?? 0) + (b.reasoningOutputTokens ?? 0) } + : {}), + ...(a.estimated || b.estimated ? { estimated: true } : {}), + }; + }; + const takeUsageFrom = (events: AdapterEvent[]): void => { + for (const e of events) { + if ((e.type === "done" || e.type === "incomplete") && e.usage) { + hiddenUsage = addUsage(hiddenUsage, e.usage); + } + } + }; + + const messages: OcxMessage[] = [...parsed.context.messages]; + const allTools = parsed.context.tools ?? []; + // Forced-final must strip every image-generation alias the plan knows about — not only tools + // flagged `imageGeneration:true`. Hosted `image_generation` / function aliases would otherwise + // remain callable; scanEventsForImageCall would strip the call while forceFinal blocks fulfillment, + // leaving the client an empty completion. + const toolsNoImage = allTools.filter(t => { + if (t.imageGeneration) return false; + if (plan.toolNames.has(t.name)) return false; + if (t.namespace && plan.toolNames.has(namespacedToolName(t.namespace, t.name))) return false; + return true; + }); + 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 — except for runTurn adapters, which + // have no HTTP status surface and must not block SSE headers behind queue.collect(). + 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"), + }); + // Attempt telemetry must fire at dispatch time (parity with fetchOnce), not after collect. + deps.onAttemptSend?.(); + 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(); + }); + + // Bound collect with a real *idle* deadline that resets on each emitted event. + // A fixed wall-clock race would abort legitimate long Cursor turns that keep + // producing tokens. Do NOT manufacture adapter heartbeats here — bridgeToResponsesSSE + // treats those as upstream activity and would defeat the stall guard. SSE keepalives + // come from the bridge heartbeat interval instead. + // + // On idle expiry: abort the runTurn signal AND close the queue so the consumer + // unblocks even when adapter.runTurn ignores cancellation and never settles. + let timedOut = false; + const idle = idleDeadline(stallTimeoutMs, () => { + timedOut = true; + // Cancel the fire-and-forget runTurn so a well-behaved adapter can stop. + internalAbort.abort(`runTurn inactivity timeout after ${stallTimeoutMs}ms`); + // Independently unblock queue.stream() — do not wait for runTurn to observe abort. + queue.close(); + }); + const events: AdapterEvent[] = []; + try { + idle.reset(); + for await (const event of queue.stream()) { + if (timedOut) break; + idle.reset(); + events.push(event); + } + } finally { + idle.cancel(); + } + if (timedOut) { + throw new LoopError(504, `runTurn inactivity timeout after ${stallTimeoutMs}ms during image-bridge`); + } + + // Preserve Cursor conversation continuity across image-loop iterations. runTurn mutates + // iterParsed (shallow copy); copy the id back onto the shared parsed request. + if (iterParsed._cursorConversationId) { + parsed._cursorConversationId = iterParsed._cursorConversationId; + } + + // 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); + } + + 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(connectTimeoutMs, signal); + try { + const fetchOnce = async (requestAdapter: ProviderAdapter): Promise => { + const request = await requestAdapter.buildRequest(iterParsed, { + headers: deps.forwardHeaders ? new Headers(deps.forwardHeaders) : new Headers(), + abortSignal: headerDeadline.signal, + }); + try { deps.onRequestBuilt?.(request); } catch { /* diagnostics are best-effort */ } + deps.onAttemptSend?.(); + const response = requestAdapter.fetchResponse + ? await requestAdapter.fetchResponse(request, { + abortSignal: headerDeadline.signal, + timeoutMs: connectTimeoutMs, + returnRawErrors: true, + stream: true, + }) + : await fetchWithResetRetry( + () => { + const h = new Headers(request.headers); + if (!h.has("accept-encoding")) h.set("accept-encoding", "identity"); + return fetchImpl(request.url, { + method: request.method, + headers: h, + body: request.body, + signal: headerDeadline.signal, + }); + }, + { abortSignal: headerDeadline.signal, label: "image-bridge-loop" }, + ); + return { response, responseAdapter: requestAdapter }; + }; + + let prepared = await fetchOnce(adapter); + // 429 key-failover parity with web-search / normal routed path. + while (prepared.response.status === 429 && deps.on429) { + const rotated = deps.on429(prepared.response.headers.get("retry-after")); + if (!rotated) break; + try { void prepared.response.body?.cancel().catch(() => {}); } catch { /* already closed */ } + adapter = rotated; + yield { type: "heartbeat" }; + prepared = await fetchOnce(adapter); + } + + // Final headers have arrived. Clear only the deadline timer before ANY body read. + headerDeadline.clear(); + if (!prepared.response.ok) { + let body: Awaited>; + try { + body = await readBoundedResponseBody(prepared.response, { signal }); + } catch { + if (signal.aborted) throw new LoopError(499, "client closed request during image-bridge"); + throw new LoopError(prepared.response.status, `Provider error ${prepared.response.status}`); + } + let formatted = ""; + if (body.displaySafe && !body.truncated && body.text.trim() && prepared.responseAdapter.formatErrorBody) { + try { + formatted = prepared.responseAdapter.formatErrorBody( + prepared.response.status, + prepared.response.headers, + body.text, + ).trim(); + } catch { /* formatter hooks are best-effort */ } + } + const suffix = formatted ? `: ${formatted.slice(0, 400)}` : ""; + throw new LoopError(prepared.response.status, `Provider error ${prepared.response.status}${suffix}`); + } + return prepared; + } catch (error) { + if (headerDeadline.didExpire()) { + throw new LoopError(504, `Provider response-header timeout after ${connectTimeoutMs}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: stallTimeoutMs, + })) { + 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. Skip for runTurn adapters: their "headers" are synthetic, and awaiting + // queue.collect() before returning SSE starves clients of headers/heartbeats on slow first turns. + const skipEagerDrain = !!adapter.runTurn; + let firstPrepared: IterationResponse | undefined; + if (!skipEagerDrain) { + 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 when present. runTurn (and later iterations) + // acquire headers inside the live SSE stream so clients already have the response open. + if (!prepared || 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); + prepared = undefined; + + // 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) { + if (hiddenUsage) { + for (let i = split.passthrough.length - 1; i >= 0; i--) { + const e = split.passthrough[i]; + if (e?.type === "done" || e?.type === "incomplete") { + split.passthrough[i] = { ...e, usage: addUsage(hiddenUsage, e.usage) }; + break; + } + } + } + yield* replay(split.passthrough); + return; + } + + // Discarded iteration still contributed tokens — accumulate for the final onUsage. + takeUsageFrom(split.passthrough); + + // Fulfill each image call, then inject ONE assistant turn (thinking once + all tool + // calls) so Anthropic extended-thinking continuations stay valid across parallel calls. + const iterationThinking = extractIterationThinking(split.passthrough); + const fulfilled: Array<{ call: ImageCall; result: Awaited>; args: Record }> = []; + for (const call of split.calls) { + yield { type: "heartbeat" }; + let result: Awaited>; + if (paidImageCalls >= MAX_IMAGE_CALLS_PER_TURN) { + result = { + ok: false, + model: plan.model, + prompt: "", + files: [], + count: 0, + error: `image call budget exhausted (max ${MAX_IMAGE_CALLS_PER_TURN} per turn)`, + }; + } else { + paidImageCalls += 1; + 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"); + 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 */ } + fulfilled.push({ call, result, args: parsedArgs }); + } + const now = Date.now(); + messages.push({ + role: "assistant", + content: [ + ...(iterationThinking ? [iterationThinking] : []), + ...fulfilled.map(({ call, args }) => ({ + type: "toolCall" as const, + id: call.id, + name: call.name, + arguments: args, + })), + ], + timestamp: now, + }); + for (const { call, result } of fulfilled) { + 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"); + }, 2_000, + { + responseId: "", + hideThinkingSummary: parsed.options.hideThinkingSummary, + stallTimeoutSec: deps.stallTimeoutSec, + ...(deps.onFirstOutput ? { onFirstOutput: deps.onFirstOutput } : {}), + ...(deps.onUsage ? { + // Terminal done/incomplete already includes hiddenUsage (merged above). Do not + // add it again here or request logs double-count multi-iteration image turns. + onUsage: (usage: OcxUsage | undefined) => deps.onUsage?.(usage), + } : {}), + ...(deps.onCompletedResponse ? { onCompletedResponse: deps.onCompletedResponse } : {}), + }, + ); + 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..9d9c960c0 --- /dev/null +++ b/src/images/plan.ts @@ -0,0 +1,89 @@ +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"; +/** Absolute ceiling for `images.timeoutMs` (matches /v1/images relay budget). */ +export const MAX_IMAGE_TIMEOUT_MS = 300_000; + +function clampImageTimeoutMs(value: unknown): number | undefined { + if (typeof value !== "number" || !Number.isFinite(value) || value <= 0) return undefined; + return Math.max(1, Math.min(MAX_IMAGE_TIMEOUT_MS, Math.floor(value))); +} + +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 { + // Honor the selected auth mode (same rules as resolveModelsAuthToken): oauth must not + // prefer a stale apiKey, and key mode must not silently fall back to stored OAuth. + if (provider.authMode === "oauth") { + if (providerName !== "xai") return undefined; + try { + return await getValidAccessToken("xai"); + } catch { + return undefined; + } + } + const apiKey = resolveEnvValue(provider.apiKey)?.trim(); + if (apiKey) return apiKey; + // Legacy / unset authMode: key-first, then built-in OAuth only for the canonical "xai" name. + if (providerName !== "xai") return undefined; + if (provider.authMode === "key" || provider.authMode === "forward") 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); + const original = parsed._imageGeneration.originalTool; + const hostedSize = typeof original?.size === "string" ? original.size : undefined; + const hostedQuality = typeof original?.quality === "string" ? original.quality : undefined; + const timeoutMs = clampImageTimeoutMs(config.images?.timeoutMs); + return { + provider: found.provider, + auth: { baseUrl: pinnedBaseUrl, token }, + model: config.images?.bridgeModel ?? DEFAULT_MODEL, + toolNames, + ...(hostedSize ? { defaultSize: hostedSize } : {}), + ...(hostedQuality ? { defaultQuality: hostedQuality } : {}), + ...(timeoutMs !== undefined ? { timeoutMs } : {}), + }; +} diff --git a/src/images/synthetic-tool.ts b/src/images/synthetic-tool.ts new file mode 100644 index 000000000..6a1dc1f0c --- /dev/null +++ b/src/images/synthetic-tool.ts @@ -0,0 +1,84 @@ +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"; + +/** Aliases that only activate when a hosted image_generation/image_gen tool is also present. */ +const IMAGE_GEN_ALIASES = new Set(["imagegen", "generate_image", "generateimage"]); + +export function isImageGenName(name: string): boolean { + const lower = name.toLowerCase(); + return lower === IMAGE_GEN_TOOL_NAME || lower === "image_generation"; +} + +function isHostedImageGenFunctionName(name: string, hasHostedEntry: boolean): boolean { + const lower = name.toLowerCase(); + if (lower === IMAGE_GEN_TOOL_NAME || lower === "image_generation") return true; + return hasHostedEntry && IMAGE_GEN_ALIASES.has(lower); +} + +/** + * 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; + let hasHostedEntry = false; + 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") { + hasHostedEntry = true; + break; + } + } + 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 && isHostedImageGenFunctionName(fnName, hasHostedEntry)) { + 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..d503828ab --- /dev/null +++ b/src/images/types.ts @@ -0,0 +1,24 @@ +import type { OcxProviderConfig } from "../types"; + +export interface ImageBridgePlan { + provider: OcxProviderConfig; + auth: { baseUrl: string; token: string }; + model: string; + toolNames: Set; + /** Per-call xAI deadline (ms). Defaults inside callXaiImages when omitted. */ + timeoutMs?: number; + /** Defaults from the hosted image_generation tool when the model omits size/quality. */ + defaultSize?: string; + defaultQuality?: string; +} + +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..2b7f01de5 --- /dev/null +++ b/src/images/xai-client.ts @@ -0,0 +1,140 @@ +/** + * 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, + timeoutMs: number = XAI_IMAGES_TIMEOUT_MS, +): 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 deadlineMs = Number.isFinite(timeoutMs) && timeoutMs > 0 + ? Math.floor(timeoutMs) + : XAI_IMAGES_TIMEOUT_MS; + const timeout = AbortSignal.timeout(deadlineMs); + const linkedSignal = signal ? AbortSignal.any([signal, timeout]) : timeout; + + const baseUrl = auth.baseUrl.replace(/\/+$/, ""); + const resp = await fetch(`${baseUrl}${endpoint}`, { + method: "POST", + headers: { + "Authorization": `Bearer ${auth.token}`, + "Content-Type": "application/json", + }, + body: JSON.stringify(body), + signal: linkedSignal, + }); + + if (!resp.ok) { + const err = new Error("xAI images API returned " + resp.status) as Error & { status: number }; + err.status = resp.status; + throw err; + } + + // 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..357dd7ff9 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,13 +79,33 @@ 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); - if (hextet === null) return { kind: "public", detail: "public IP" }; - if (hextet >= 0xfc00 && hextet <= 0xfdff) return { kind: "private", detail: "private-network address" }; + if (hextet === null) return { kind: "private", detail: "non-global address" }; + // Multicast ff00::/8, deprecated site-local fec0::/10, ULA fc00::/7, link-local fe80::/10. + if (hextet >= 0xff00) return { kind: "private", detail: "multicast address" }; if (hextet >= 0xfe80 && hextet <= 0xfebf) return { kind: "link-local", detail: "link-local address" }; - return { kind: "public", detail: "public IP" }; + if (hextet >= 0xfec0 && hextet <= 0xfeff) return { kind: "private", detail: "site-local address" }; + if (hextet >= 0xfc00 && hextet <= 0xfdff) return { kind: "private", detail: "private-network address" }; + // Documentation 2001:db8::/32 (inside global-unicast 2000::/3). + if (hextet === 0x2001) { + const second = Number.parseInt(hostname.split(":")[1] || "0", 16); + if (second === 0xdb8) return { kind: "private", detail: "documentation address" }; + } + // Only global unicast 2000::/3 is treated as a public image/CDN peer. + if (hextet >= 0x2000 && hextet <= 0x3fff) return { kind: "public", detail: "public IP" }; + return { kind: "private", detail: "non-global address" }; } function assessDestination(baseUrl: string): DestinationAssessment | null { @@ -178,3 +198,74 @@ 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 the validated public addresses on success + * so callers can pin the connect peer and avoid a second, rebindable resolution. + * DNS resolution failures are treated as unsafe (fail-closed). + */ +export async function resolvePublicAddresses(url: string): Promise<{ + hostname: string; + addresses: { address: string; family: number }[]; +}> { + 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}`); + } + // Literal public IPs: no DNS round-trip; pin the literal itself. + const literalKind = isIP(hostname); + if (literalKind !== 0) { + return { hostname, addresses: [{ address: hostname, family: literalKind }] }; + } + let addresses: { address: string; family: number }[]; + 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`); + } + if (addresses.length === 0) { + throw new Error(`image URL hostname ${hostname} could not be resolved`); + } + const publicAddresses: { address: string; family: number }[] = []; + for (const { address, family } of addresses) { + const ipKind = family === 4 || family === 6 ? family : isIP(address); + const assessment = ipKind === 4 ? classifyIpv4(address) : ipKind === 6 ? classifyIpv6(normalizeHostname(address)) : null; + if (!assessment || assessment.kind !== "public") { + throw new Error(`image URL hostname ${hostname} resolves to ${assessment?.detail ?? "an unsafe address"} (${address})`); + } + publicAddresses.push({ address, family: ipKind === 4 || ipKind === 6 ? ipKind : (family || 4) }); + } + return { hostname, addresses: publicAddresses }; +} + +/** + * Void assertion wrapper around {@link resolvePublicAddresses} for call sites + * that only need the safety check. + */ +export async function assertUrlResolvesPublic(url: string): Promise { + await resolvePublicAddresses(url); +} diff --git a/src/responses/parser.ts b/src/responses/parser.ts index 3c13b477e..b59ad238b 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, IMAGE_GEN_TOOL_NAME } from "../images/synthetic-tool"; function isObj(v: unknown): v is Record { return typeof v === "object" && v !== null && !Array.isArray(v); @@ -124,6 +125,7 @@ function allowedToolName(tool: unknown): string | undefined { if (!isObj(tool)) return undefined; if (typeof tool.name === "string" && tool.name.length > 0) return tool.name; if (tool.type === "web_search" || tool.type === "web_search_preview") return WEB_SEARCH_TOOL_NAME; + if (tool.type === "image_generation" || tool.type === "image_gen") return IMAGE_GEN_TOOL_NAME; if (tool.type === "tool_search") return "tool_search"; return undefined; } @@ -612,6 +614,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 +631,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..5f31af77c 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, clampImageMaxRounds, IMAGE_GEN_TOOL_NAME } from "../../images"; import { describeImagesInPlace, planVisionSidecar, shouldResolveOpenAiVisionSidecar, stripImagesInPlace } from "../../vision"; import { createAdapterEventQueue, preflightAdapterEvents } from "../../adapters/run-turn-queue"; import { @@ -1526,6 +1527,167 @@ export async function handleResponses( }); } + // Image / web-search sidecars: plan once, then dispatch with runTurn-aware priority. + // Routed-compaction turns must NOT hit the image 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). + // + // Web-search's loop only supports buildRequest/fetch/parseStream — NOT adapter.runTurn. Sending + // Cursor/runTurn requests into runWithWebSearch produces empty HTTP failures. So: + // - non-runTurn: web-search wins over image when both eligible (documented priority) + // - runTurn: image bridge may run (it supports runTurn); web-search is skipped so runTurn + // can proceed for web-search-only turns + const wsPlan = !routedCompaction + ? planWebSearch(config, parsed, false, route.provider, route.modelId, openAiSidecar) + : undefined; + const imgPlan = !routedCompaction ? await planImageBridge(config, parsed, route.provider) : undefined; + const canRunWebSearch = !!wsPlan && !adapter.runTurn; + if (imgPlan && (!wsPlan || adapter.runTurn)) { + // 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"); + } + // Replace any pre-existing image_gen alias instead of appending a duplicate wire name. + const priorTools = parsed.context.tools ?? []; + parsed.context.tools = [ + ...priorTools.filter(t => { + if (t.imageGeneration) return false; + if (imgPlan.toolNames.has(t.name)) return false; + if (t.namespace && imgPlan.toolNames.has(namespacedToolName(t.namespace, t.name))) return false; + return true; + }), + buildImageTool(), + ]; + // Hosted image_generation tool_choice / allowed_tools must target the synthetic function name. + const tc = parsed.options.toolChoice; + if (tc && typeof tc === "object" && "allowedTools" in tc && Array.isArray(tc.allowedTools)) { + const mapped = tc.allowedTools.map(name => + name === "image_generation" || name === "image_gen" || imgPlan.toolNames.has(name) + ? IMAGE_GEN_TOOL_NAME + : name, + ); + parsed.options.toolChoice = { ...tc, allowedTools: [...new Set(mapped)] }; + } else if (tc && typeof tc === "object" && "name" in tc && typeof tc.name === "string" + && (tc.name === "image_generation" || imgPlan.toolNames.has(tc.name))) { + parsed.options.toolChoice = { ...tc, name: IMAGE_GEN_TOOL_NAME }; + } + const imgResponse = await runWithImageBridge({ + parsed, adapter, + plan: imgPlan, + forwardHeaders: selectedForwardHeaders, + onAttemptSend: () => noteAttemptSend(logCtx.activeAttempt, logCtx.usageLogInputTokens), + abortSignal: options.abortSignal, + maxRounds: clampImageMaxRounds(config.images?.maxRounds), + connectTimeoutMs: config.connectTimeoutMs ?? 200_000, + stallTimeoutSec: config.stallTimeoutSec, + fetchImpl: providerFetch(route.provider), + onRequestBuilt: request => recordAdapterReasoning(logCtx, request), + onUsage: usage => { + // Cursor may assign _cursorConversationId inside the image loop's first runTurn; + // backfill so Logs can filter/total that opening request (parity with the normal + // runTurn branch). + if (!logCtx.conversationId && parsed._cursorConversationId) { + logCtx.conversationId = normalizeLogConversationId(parsed._cursorConversationId); + } + logCtx.usageFromBridge = true; + if (usage) { + logCtx.usage = usage; + if (logCtx.activeAttempt) logCtx.activeAttempt.usage = usage; + } + }, + 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, + ); + }, + ...(options.onFirstOutput ? { onFirstOutput: options.onFirstOutput } : {}), + onCompletedResponse: (response, providerState) => + rememberResponseState( + parsed._rawBody, + response, + continuationStateForResponse(providerState), + adapterNeedsForcedContinuation(adapter.name) ? { force: true } : undefined, + ), + }); + 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 for non-runTurn adapters so dual-tool turns dispatch + // through web-search instead of being swallowed. runTurn adapters never enter this branch. + if (canRunWebSearch && 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 +1808,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 fb0e244d3..0f906998e 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; } /** @@ -739,8 +743,14 @@ export interface OcxTokenGuardianConfig { export interface OcxImagesConfig { /** Optional custom API-key provider for /v1/images relays. Built-in OpenAI tiers remain automatic. */ provider?: string; - /** Upstream timeout (ms) for one /v1/images relay. Default 300000 — generation is slow. */ + /** Upstream timeout (ms) for one image generation/edit call (bridge xAI + /v1/images relay). Default 60000 for the bridge; relay may use a higher default (300000). */ 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; clamped to [0, 10]. */ + maxRounds?: number; } export interface OcxSearchConfig { diff --git a/tests/api-storage-policy.test.ts b/tests/api-storage-policy.test.ts index ef8724766..8324efdbf 100644 --- a/tests/api-storage-policy.test.ts +++ b/tests/api-storage-policy.test.ts @@ -283,7 +283,10 @@ describe("storage cleanup policy API", () => { }, { timeout: 30_000 }); test("blocked worker completion preserves concurrent policy PUT edits", async () => { - setStorageCleanupPolicyJobTestHooks({ blockMs: 1_200 }); + // Long hold so a slow Windows CI worker spawn can load the enabled snapshot + // before the concurrent PUT lands inside holdAfterLoadMs. + const blockMs = 1_500; + setStorageCleanupPolicyJobTestHooks({ blockMs }); seedArchived(isolatedCodexHome!.path); const server = startServer(0); try { @@ -307,17 +310,21 @@ describe("storage cleanup policy API", () => { expect(runStart.started).toBe(true); expect(runStart.job?.status).toBe("running"); - // Wait until the job is visibly running, then edit policy while the worker holds. - const editDeadline = Date.now() + 2_000; + // Status flips to running before the worker loads policy — wait for that marker, + // then allow spawn+load margin before editing during the hold window. + const editDeadline = Date.now() + 5_000; + let sawRunning = false; while (Date.now() < editDeadline) { const peek = await fetch(new URL("/api/storage/cleanup-policy", server.url)); const peekBody = await peek.json() as { job?: { status?: string } }; - if (peekBody.job?.status === "running") break; + if (peekBody.job?.status === "running") { + sawRunning = true; + break; + } await Bun.sleep(20); } - - // Let the worker load the start-of-job snapshot, then edit during the hold window. - await Bun.sleep(450); + expect(sawRunning).toBe(true); + await Bun.sleep(800); const put = await fetch(new URL("/api/storage/cleanup-policy", server.url), { method: "PUT", @@ -337,6 +344,7 @@ describe("storage cleanup policy API", () => { const done = await waitForJobIdle(server.url, runStart.job!.startedAt); expect(done.job.lastOutcome?.ok).toBe(true); + expect(done.job.lastOutcome?.skipped).toBeUndefined(); expect(done.job.lastOutcome?.removed).toBe(1); expect(done.enabled).toBe(false); expect(done.lastRun?.removed).toBe(1); diff --git a/tests/destination-policy-resolved.test.ts b/tests/destination-policy-resolved.test.ts index e700c3a57..d47df9416 100644 --- a/tests/destination-policy-resolved.test.ts +++ b/tests/destination-policy-resolved.test.ts @@ -28,6 +28,11 @@ describe("providerDestinationConfigError — reserved IPv4 ranges (review findin test("still passes ordinary public literals", () => { expect(providerDestinationConfigError("custom", provider("https://93.184.216.34/v1"))).toBeNull(); }); + + test("rejects IPv6 site-local and multicast literals", () => { + expect(providerDestinationConfigError("custom", provider("http://[fec0::1]/v1"))).toContain("allowPrivateNetwork"); + expect(providerDestinationConfigError("custom", provider("http://[ff02::1]/v1"))).toContain("allowPrivateNetwork"); + }); }); describe("providerDestinationResolvedError — DNS-resolved SSRF check (activation)", () => { @@ -58,6 +63,18 @@ describe("providerDestinationResolvedError — DNS-resolved SSRF check (activati expect(error).toContain("private-network address (fd00::1)"); }); + test("blocks a hostname resolving to IPv6 site-local space", async () => { + lookupMock.mockResolvedValueOnce([{ address: "fec0::1", family: 6 }]); + const error = await providerDestinationResolvedError("custom", provider("https://v6-site.example.com/v1")); + expect(error).toMatch(/site-local address \(fec0::1\)/); + }); + + test("blocks a hostname resolving to IPv6 multicast space", async () => { + lookupMock.mockResolvedValueOnce([{ address: "ff02::1", family: 6 }]); + const error = await providerDestinationResolvedError("custom", provider("https://v6-mcast.example.com/v1")); + expect(error).toMatch(/multicast address \(ff02::1\)/); + }); + test("passes a hostname resolving only to public addresses", async () => { lookupMock.mockResolvedValueOnce([{ address: "93.184.216.34", family: 4 }]); expect(await providerDestinationResolvedError("custom", provider("https://api.example.com/v1"))).toBeNull(); diff --git a/tests/images/artifacts-ssrf.test.ts b/tests/images/artifacts-ssrf.test.ts new file mode 100644 index 000000000..703633398 --- /dev/null +++ b/tests/images/artifacts-ssrf.test.ts @@ -0,0 +1,241 @@ +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, resolvePublicAddresses } = await import("../../src/lib/destination-policy"); +const { downloadImageToArtifact } = await import("../../src/images/artifacts"); + +const MIN_PNG = new Uint8Array([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]); + +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("IPv6 site-local [fec0::1] → private", () => { + expect(assessUrlDestination("https://[fec0::1]/image.png")?.kind).toBe("private"); + expect(assessUrlDestination("https://[fec0::1]/image.png")?.detail).toContain("site-local"); + }); + test("IPv6 multicast [ff02::1] → private", () => { + expect(assessUrlDestination("https://[ff02::1]/image.png")?.kind).toBe("private"); + expect(assessUrlDestination("https://[ff02::1]/image.png")?.detail).toContain("multicast"); + }); + test("IPv6 documentation [2001:db8::1] → private", () => { + expect(assessUrlDestination("https://[2001:db8::1]/image.png")?.kind).toBe("private"); + expect(assessUrlDestination("https://[2001:db8::1]/image.png")?.detail).toContain("documentation"); + }); + test("IPv6 global unicast [2001:4860:4860::8888] → public", () => { + expect(assessUrlDestination("https://[2001:4860:4860::8888]/image.png")?.kind).toBe("public"); + }); + 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: resolvePublicAddresses", () => { + test("hostname with public A record → returns that address", async () => { + lookupMock.mockResolvedValue([{ address: "93.184.216.34", family: 4 }]); + try { + const resolved = await resolvePublicAddresses("https://public-host/img.png"); + expect(resolved.hostname).toBe("public-host"); + expect(resolved.addresses).toEqual([{ address: "93.184.216.34", family: 4 }]); + } finally { + lookupMock.mockClear(); + } + }); + + test("hostname that also resolves private → throws (fail closed on any unsafe answer)", async () => { + lookupMock.mockResolvedValue([ + { address: "93.184.216.34", family: 4 }, + { address: "127.0.0.1", family: 4 }, + ]); + try { + await expect(resolvePublicAddresses("https://mixed-host/img.png")).rejects.toThrow(/loopback|127\.0\.0\.1/); + } finally { + lookupMock.mockClear(); + } + }); + + test("hostname resolving to IPv6 site-local → throws", async () => { + lookupMock.mockResolvedValue([{ address: "fec0::1", family: 6 }]); + try { + await expect(resolvePublicAddresses("https://v6-site.example/img.png")).rejects.toThrow(/site-local|fec0/); + } finally { + lookupMock.mockClear(); + } + }); + + test("hostname resolving to IPv6 multicast → throws", async () => { + lookupMock.mockResolvedValue([{ address: "ff02::1", family: 6 }]); + try { + await expect(resolvePublicAddresses("https://v6-mcast.example/img.png")).rejects.toThrow(/multicast|ff02/); + } finally { + lookupMock.mockClear(); + } + }); +}); + +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 without following`, async () => { + lookupMock.mockResolvedValue([{ address: "93.184.216.34", family: 4 }]); + try { + await expect(downloadImageToArtifact("https://public-host/redirect-img", undefined, undefined, { + pinnedDownload: async () => new Response("", { + status: 301, + headers: { Location: "https://evil.example/redirect" }, + }), + })).rejects.toThrow(/301|failed/); + } finally { + lookupMock.mockClear(); + } + }); + + test("https:// public host → succeeds with pinned download", async () => { + lookupMock.mockResolvedValue([{ address: "93.184.216.34", family: 4 }]); + let downloadedPath: string | undefined; + let seenPinned: { address: string; family: number } | undefined; + try { + downloadedPath = await downloadImageToArtifact("https://public-host/valid-image", undefined, undefined, { + pinnedDownload: async (_url, pinned) => { + seenPinned = pinned; + return new Response(MIN_PNG, { status: 200 }); + }, + }); + expect(downloadedPath).toMatch(/dl-.*\.png$/); + expect(seenPinned).toEqual({ address: "93.184.216.34", family: 4 }); + } finally { + lookupMock.mockClear(); + if (downloadedPath) await rm(downloadedPath).catch(() => {}); + } + }); + + test("empty 200 body → rejects", async () => { + lookupMock.mockResolvedValue([{ address: "93.184.216.34", family: 4 }]); + try { + await expect(downloadImageToArtifact("https://public-host/empty", undefined, undefined, { + pinnedDownload: async () => new Response(new Uint8Array(0), { status: 200 }), + })).rejects.toThrow(/empty/); + } finally { + lookupMock.mockClear(); + } + }); + + test("non-image 200 body (HTML/JSON) → rejects", async () => { + lookupMock.mockResolvedValue([{ address: "93.184.216.34", family: 4 }]); + try { + await expect(downloadImageToArtifact("https://public-host/not-image", undefined, undefined, { + pinnedDownload: async () => new Response("error", { status: 200 }), + })).rejects.toThrow(/recognized image/); + } finally { + lookupMock.mockClear(); + } + }); + + test("DNS rebinding: connection uses the validated public address, not a later private resolve", async () => { + // Validation lookup returns public; any subsequent OS resolve would return loopback. + // The download must pin the first answer and must not call dns.lookup again. + // (Transport lookup-shape + byte-cap coverage lives in pinned-https-get.test.ts.) + let lookups = 0; + lookupMock.mockImplementation(async () => { + lookups += 1; + if (lookups === 1) return [{ address: "93.184.216.34", family: 4 }]; + return [{ address: "127.0.0.1", family: 4 }]; + }); + + let downloadedPath: string | undefined; + let seenPinned: { address: string; family: number } | undefined; + try { + downloadedPath = await downloadImageToArtifact("https://rebind.example/img.png", undefined, undefined, { + pinnedDownload: async (_url, pinned) => { + // Still only one DNS resolve at connect time — pin came from validation. + expect(lookups).toBe(1); + seenPinned = pinned; + expect(pinned.address).toBe("93.184.216.34"); + expect(pinned.address).not.toBe("127.0.0.1"); + return new Response(MIN_PNG, { status: 200 }); + }, + }); + expect(downloadedPath).toMatch(/dl-.*\.png$/); + expect(seenPinned?.address).toBe("93.184.216.34"); + expect(lookups).toBe(1); + } finally { + lookupMock.mockClear(); + if (downloadedPath) await rm(downloadedPath).catch(() => {}); + } + }); +}); diff --git a/tests/images/loop.test.ts b/tests/images/loop.test.ts new file mode 100644 index 000000000..8315effbc --- /dev/null +++ b/tests/images/loop.test.ts @@ -0,0 +1,514 @@ +import { afterAll, beforeAll, beforeEach, 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; +let runWithImageBridge: typeof import("../../src/images/loop")["runWithImageBridge"]; +let clampImageMaxRounds: typeof import("../../src/images/loop")["clampImageMaxRounds"]; +let DEFAULT_MAX_ROUNDS: typeof import("../../src/images/loop")["DEFAULT_MAX_ROUNDS"]; +let MAX_ROUNDS_HARD_LIMIT: typeof import("../../src/images/loop")["MAX_ROUNDS_HARD_LIMIT"]; + +let fulfillResult: ImageCallResult = { + ok: true, model: "grok-imagine-image-quality", prompt: "a cat", + files: ["/test/img.png"], count: 1, markdown: "![image](/test/img.png)", +}; + +beforeAll(async () => { + process.env.OPENCODEX_HOME = join(tmpdir(), "ocx-test-" + randomUUID()); + mock.restore(); + 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.module("../../src/images/fulfill", () => ({ + fulfillImageCall: async (): Promise => fulfillResult, + })); + ({ + runWithImageBridge, + clampImageMaxRounds, + DEFAULT_MAX_ROUNDS, + MAX_ROUNDS_HARD_LIMIT, + } = await import("../../src/images/loop")); +}); +afterAll(() => { if (PREV_HOME === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = PREV_HOME; mock.restore(); }); + +// --- Mock adapter: yields canned events per iteration from a queue --- +let streamQueue: AdapterEvent[][] = []; +let buildRequestCalls = 0; + +const defaultFulfillResult: ImageCallResult = { + ok: true, model: "grok-imagine-image-quality", prompt: "a cat", + files: ["/test/img.png"], count: 1, markdown: "![image](/test/img.png)", +}; +beforeEach(() => { + fulfillResult = { ...defaultFulfillResult, files: [...defaultFulfillResult.files] }; + buildRequestCalls = 0; + streamQueue = []; +}); + +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); + }); + + test("clampImageMaxRounds bounds hand-edited / fractional values", () => { + expect(clampImageMaxRounds(10000)).toBe(MAX_ROUNDS_HARD_LIMIT); + expect(clampImageMaxRounds(2.9)).toBe(2); + expect(clampImageMaxRounds(-1)).toBe(0); + expect(clampImageMaxRounds(Number.NaN)).toBe(DEFAULT_MAX_ROUNDS); + expect(clampImageMaxRounds(undefined)).toBe(DEFAULT_MAX_ROUNDS); + }); + + test("maxRounds: 10000 is clamped — hits hard limit when every round calls image_gen", async () => { + buildRequestCalls = 0; + streamQueue = []; + for (let i = 0; i < MAX_ROUNDS_HARD_LIMIT; i++) { + streamQueue.push([...imageCallEvents]); + } + // Forced-final pass after the hard cap. + streamQueue.push([{ type: "text_delta" as const, text: "clamped final" }, { type: "done" as const }]); + const response = await runWithImageBridge({ + parsed: makeParsed(), adapter: mockAdapter, plan, maxRounds: 10000, + }); + const sse = await response.text(); + expect(sse).toContain("clamped final"); + // Clamped to 10 → HARD_CAP = 11 upstream requests (10 image rounds + 1 forced final). + expect(buildRequestCalls).toBe(MAX_ROUNDS_HARD_LIMIT + 1); + }); + + test("forced-final strips image aliases from plan.toolNames, not only imageGeneration flag", async () => { + const seenTools: Array = []; + const capturingAdapter: ProviderAdapter = { + ...mockAdapter, + buildRequest: async (parsed) => { + buildRequestCalls++; + seenTools.push(parsed.context.tools?.map(t => t.name)); + return { url: "https://test/v1/chat", method: "POST", headers: {}, body: "{}" }; + }, + }; + streamQueue = [ + [ + { type: "tool_call_start", id: "call_1", name: "image_generation" }, + { type: "tool_call_delta", arguments: '{"prompt":"a cat"}' }, + { type: "tool_call_end" }, + { type: "done" }, + ], + [{ type: "text_delta", text: "done after strip" }, { type: "done" }], + ]; + const aliasPlan = { + ...plan, + toolNames: new Set(["image_generation", "image_gen"]), + } as ImageBridgePlan; + const parsed = makeParsed(); + parsed.context.tools = [ + { name: "image_generation", parameters: {}, description: "hosted" }, + { name: "Bash", parameters: {}, description: "shell" }, + ]; + const response = await runWithImageBridge({ + parsed, adapter: capturingAdapter, plan: aliasPlan, maxRounds: 1, + }); + await response.text(); + // Second request is forceFinal — image_generation must be gone, Bash remains. + expect(seenTools[1]).toEqual(["Bash"]); + }); + + test("parallel image calls share one assistant turn with thinking attached once", async () => { + const seenMessages: unknown[] = []; + const capturingAdapter: ProviderAdapter = { + ...mockAdapter, + buildRequest: async (parsed) => { + buildRequestCalls++; + seenMessages.push(parsed.context.messages.map(m => ({ + role: m.role, + contentTypes: Array.isArray(m.content) + ? m.content.map((c: { type?: string }) => c.type) + : typeof m.content, + }))); + return { url: "https://test/v1/chat", method: "POST", headers: {}, body: "{}" }; + }, + }; + streamQueue = [ + [ + { type: "thinking_delta", thinking: "planning" }, + { type: "thinking_signature", signature: "sig" }, + { type: "tool_call_start", id: "call_a", name: "image_gen" }, + { type: "tool_call_delta", arguments: '{"prompt":"a"}' }, + { type: "tool_call_end" }, + { type: "tool_call_start", id: "call_b", name: "image_gen" }, + { type: "tool_call_delta", arguments: '{"prompt":"b"}' }, + { type: "tool_call_end" }, + { type: "done" }, + ], + [{ type: "text_delta", text: "both images ready" }, { type: "done" }], + ]; + const response = await runWithImageBridge({ + parsed: makeParsed(), adapter: capturingAdapter, plan, maxRounds: 1, + }); + await response.text(); + // Second iteration messages should include exactly one assistant turn with thinking + 2 toolCalls. + const second = seenMessages[1] as Array<{ role: string; contentTypes: string[] }>; + const assistants = second.filter(m => m.role === "assistant"); + expect(assistants.length).toBe(1); + expect(assistants[0]!.contentTypes).toEqual(["thinking", "toolCall", "toolCall"]); + }); + + test("onUsage is forwarded from bridge terminal events", async () => { + let seen: unknown = "unset"; + streamQueue = [ + [{ type: "text_delta", text: "hi" }, { type: "done", usage: { inputTokens: 1, outputTokens: 2 } }], + ]; + const response = await runWithImageBridge({ + parsed: makeParsed(), + adapter: mockAdapter, + plan, + onUsage: usage => { seen = usage; }, + }); + await response.text(); + expect(seen).toEqual({ inputTokens: 1, outputTokens: 2 }); + }); + + test("onUsage does not double-count hiddenUsage across image iterations", async () => { + let seen: unknown = "unset"; + streamQueue = [ + [ + { type: "tool_call_start", id: "call_1", name: "image_gen" }, + { type: "tool_call_delta", arguments: '{"prompt":"a cat"}' }, + { type: "tool_call_end" }, + { type: "done", usage: { inputTokens: 10, outputTokens: 4 } }, + ], + [{ type: "text_delta", text: "ready" }, { type: "done", usage: { inputTokens: 3, outputTokens: 2 } }], + ]; + const response = await runWithImageBridge({ + parsed: makeParsed(), + adapter: mockAdapter, + plan, + maxRounds: 1, + onUsage: usage => { seen = usage; }, + }); + await response.text(); + // Hidden iter (10/4) + final (3/2) once — not 2*hidden + final. + expect(seen).toEqual({ inputTokens: 13, outputTokens: 6 }); + }); + + test("429 key rotation rebuilds the adapter and retries the iteration", async () => { + let fetchCalls = 0; + let rotations = 0; + let activeAdapter: ProviderAdapter | undefined; + const makeRotatingAdapter = (label: string): ProviderAdapter => ({ + name: label, + buildRequest: async () => ({ url: "https://test/v1/chat", method: "POST", headers: {}, body: "{}" }), + fetchResponse: async () => { + fetchCalls++; + if (fetchCalls === 1) return new Response("rate limited", { status: 429, headers: { "retry-after": "1" } }); + streamQueue = [[{ type: "text_delta", text: "after rotate" }, { type: "done" }]]; + return new Response("{}", { status: 200 }); + }, + parseStream: async function* (): AsyncGenerator { + const events = streamQueue.shift(); + if (events) for (const e of events) yield e; + }, + }); + const firstAdapter = makeRotatingAdapter("before-rotate"); + const secondAdapter = makeRotatingAdapter("after-rotate"); + activeAdapter = firstAdapter; + const response = await runWithImageBridge({ + parsed: makeParsed(), + adapter: firstAdapter, + plan, + on429: () => { + rotations++; + activeAdapter = secondAdapter; + return secondAdapter; + }, + }); + const sse = await response.text(); + expect(rotations).toBe(1); + expect(fetchCalls).toBe(2); + expect(activeAdapter).toBe(secondAdapter); + expect(sse).toContain("after rotate"); + }); +}); + +// --------------------------------------------------------------------------- +// 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"); + }); + + test("runTurn adapter → SSE headers return before slow collect completes", async () => { + let release!: () => void; + const gate = new Promise(resolve => { release = resolve; }); + const slowAdapter: ProviderAdapter = { + ...mockAdapter, + runTurn: async (_parsed, _incoming, emit) => { + await gate; + emit({ type: "text_delta", text: "slow ok" }); + emit({ type: "done" }); + }, + }; + const responsePromise = runWithImageBridge({ parsed: makeParsed(), adapter: slowAdapter, plan }); + // Headers must resolve without waiting for runTurn to finish. + const response = await responsePromise; + expect(response.headers.get("content-type")).toBe("text/event-stream"); + release(); + const sse = await response.text(); + expect(sse).toContain("slow ok"); + }); + + test("runTurn adapter → stall deadline aborts the in-flight runTurn signal", async () => { + let seenSignal: AbortSignal | undefined; + let aborted = false; + const hangingAdapter: ProviderAdapter = { + ...mockAdapter, + runTurn: async (_parsed, incoming) => { + seenSignal = incoming.abortSignal; + await new Promise((resolve, reject) => { + if (!incoming.abortSignal) { + reject(new Error("missing abortSignal")); + return; + } + if (incoming.abortSignal.aborted) { + aborted = true; + resolve(); + return; + } + incoming.abortSignal.addEventListener("abort", () => { + aborted = true; + resolve(); + }, { once: true }); + }); + }, + }; + const response = await runWithImageBridge({ + parsed: makeParsed(), + adapter: hangingAdapter, + plan, + // Short stall so the collect deadline wins without waiting on real upstream silence. + stallTimeoutSec: 0.05, + }); + const sse = await response.text(); + expect(seenSignal).toBeDefined(); + expect(aborted).toBe(true); + expect(sse).toContain("runTurn inactivity timeout"); + }); + + test("runTurn adapter → idle timeout completes when runTurn ignores abort and never settles", async () => { + // Regression: aborting internalAbort alone is not enough — if runTurn never observes + // cancellation and never closes the queue, queue.stream() would hang forever. The idle + // handler must close the queue to unblock the consumer independently. + const neverSettlingAdapter: ProviderAdapter = { + ...mockAdapter, + runTurn: async () => { + await new Promise(() => { /* never settles; ignores abort entirely */ }); + }, + }; + const started = Date.now(); + const response = await runWithImageBridge({ + parsed: makeParsed(), + adapter: neverSettlingAdapter, + plan, + stallTimeoutSec: 0.05, + }); + const sse = await response.text(); + const elapsedMs = Date.now() - started; + const timeoutFrames = sse.split("\n\n").filter(frame => frame.includes("runTurn inactivity timeout")); + expect(timeoutFrames).toHaveLength(1); + expect(timeoutFrames[0]).toContain("response.failed"); + // Must finish near the idle deadline, not hang on the never-settling runTurn. + expect(elapsedMs).toBeLessThan(2_000); + }); + + test("runTurn adapter → continuous progress resets the idle stall deadline", async () => { + // Wall-clock for the whole turn exceeds stallTimeoutSec, but each idle gap is shorter. + const progressingAdapter: ProviderAdapter = { + ...mockAdapter, + runTurn: async (_parsed, incoming, emit) => { + for (let i = 0; i < 6; i++) { + if (incoming.abortSignal?.aborted) return; + emit({ type: "text_delta", text: `chunk${i}` }); + await new Promise(r => setTimeout(r, 40)); + } + if (!incoming.abortSignal?.aborted) emit({ type: "done" }); + }, + }; + const response = await runWithImageBridge({ + parsed: makeParsed(), + adapter: progressingAdapter, + plan, + stallTimeoutSec: 0.1, // 100ms idle; total emit span ~240ms + }); + const sse = await response.text(); + expect(sse).toContain("chunk5"); + expect(sse).not.toContain("inactivity timeout"); + }); + + test("runTurn adapter → preserves _cursorConversationId across iterations", async () => { + const seenIds: Array = []; + const cursorAdapter: ProviderAdapter = { + ...mockAdapter, + runTurn: async (parsed, _incoming, emit) => { + seenIds.push(parsed._cursorConversationId); + if (!parsed._cursorConversationId) { + parsed._cursorConversationId = "conv-from-first-turn"; + } + const events = runTurnEventQueue.shift(); + if (events) for (const e of events) emit(e); + }, + }; + runTurnEventQueue = [ + [...imageCallEvents], + [{ type: "text_delta", text: "second turn" }, { type: "done" }], + ]; + const parsed = makeParsed(); + const response = await runWithImageBridge({ + parsed, adapter: cursorAdapter, plan, maxRounds: 1, + }); + await response.text(); + expect(seenIds[0]).toBeUndefined(); + expect(seenIds[1]).toBe("conv-from-first-turn"); + expect(parsed._cursorConversationId).toBe("conv-from-first-turn"); + }); +}); diff --git a/tests/images/pinned-https-get.test.ts b/tests/images/pinned-https-get.test.ts new file mode 100644 index 000000000..e3d1f8017 --- /dev/null +++ b/tests/images/pinned-https-get.test.ts @@ -0,0 +1,321 @@ +import { EventEmitter } from "node:events"; +import { describe, expect, mock, test } from "bun:test"; + +type LookupCb = + | ((err: Error | null, address: string, family: number) => void) + | ((err: Error | null, addresses: { address: string; family: number }[]) => void); + +/** + * Capture the custom lookup + response-stream path used by pinnedHttpsGet without + * opening a real TLS socket (Windows CI friendly). + */ +function installHttpsMock(bodyChunks: Buffer[], statusCode = 200) { + const requestMock = mock(( + _options: unknown, + onResponse?: (res: EventEmitter & { statusCode: number; headers: Record; setTimeout: Function; resume: Function }) => void, + ) => { + const req = new EventEmitter() as EventEmitter & { + setTimeout: Function; + end: Function; + destroy: Function; + destroyed: boolean; + }; + req.destroyed = false; + req.setTimeout = mock(() => {}); + req.destroy = mock(() => { req.destroyed = true; }); + req.end = mock(() => { + const res = new EventEmitter() as EventEmitter & { + statusCode: number; + headers: Record; + setTimeout: Function; + resume: Function; + }; + res.statusCode = statusCode; + res.headers = { "content-type": "image/png" }; + res.setTimeout = mock(() => {}); + res.resume = mock(() => {}); + queueMicrotask(() => { + onResponse?.(res); + queueMicrotask(() => { + for (const chunk of bodyChunks) res.emit("data", chunk); + res.emit("end"); + }); + }); + }); + return req; + }); + + mock.module("node:https", () => ({ + default: { request: requestMock }, + request: requestMock, + })); + + return requestMock; +} + +describe("pinnedHttpsGet transport", () => { + test("lookup honors scalar and { all: true } callback shapes", async () => { + let capturedLookup: ((hostname: string, opts: unknown, cb?: LookupCb) => void) | undefined; + const requestMock = mock((options: { lookup?: typeof capturedLookup }, onResponse?: Function) => { + capturedLookup = options.lookup; + const req = new EventEmitter() as EventEmitter & { setTimeout: Function; end: Function; destroy: Function }; + req.setTimeout = () => {}; + req.destroy = () => {}; + req.end = () => { + const res = new EventEmitter() as EventEmitter & { + statusCode: number; + headers: Record; + setTimeout: Function; + resume: Function; + }; + res.statusCode = 200; + res.headers = {}; + res.setTimeout = () => {}; + res.resume = () => {}; + queueMicrotask(() => { + onResponse?.(res); + queueMicrotask(() => res.emit("end")); + }); + }; + return req; + }); + mock.module("node:https", () => ({ default: { request: requestMock }, request: requestMock })); + + const { pinnedHttpsGet } = await import("../../src/images/artifacts"); + const pinned = { address: "93.184.216.34", family: 4 }; + const respPromise = pinnedHttpsGet("https://cdn.example/img.png", pinned); + // Give request() a tick to store lookup. + await Promise.resolve(); + expect(capturedLookup).toBeTypeOf("function"); + + let scalar: { address?: string; family?: number } = {}; + capturedLookup!("cdn.example", {}, ((err, address, family) => { + expect(err).toBeNull(); + scalar = { address: address as string, family: family as number }; + }) as LookupCb); + expect(scalar).toEqual({ address: "93.184.216.34", family: 4 }); + + let allAddrs: { address: string; family: number }[] | undefined; + capturedLookup!("cdn.example", { all: true }, ((err, addresses) => { + expect(err).toBeNull(); + allAddrs = addresses as { address: string; family: number }[]; + }) as LookupCb); + expect(allAddrs).toEqual([{ address: "93.184.216.34", family: 4 }]); + + const resp = await respPromise; + expect(resp.ok).toBe(true); + await resp.arrayBuffer(); // drain stream + }); + + test("exceeding maxBytes aborts mid-stream without buffering the full body", async () => { + const small = Buffer.alloc(1024, 1); + const chunks = [small, small, small]; // 3 KiB total + installHttpsMock(chunks); + + const { pinnedHttpsGet } = await import("../../src/images/artifacts"); + const maxBytes = 1500; // trip on the second chunk + const resp = await pinnedHttpsGet( + "https://cdn.example/big.png", + { address: "93.184.216.34", family: 4 }, + undefined, + { maxBytes }, + ); + expect(resp.body).toBeTruthy(); + const reader = resp.body!.getReader(); + let sawError = false; + let received = 0; + try { + for (;;) { + const { value, done } = await reader.read(); + if (done) break; + received += value.byteLength; + } + } catch { + sawError = true; + } + expect(sawError).toBe(true); + // Must fail before absorbing all three chunks (3 KiB). + expect(received).toBeLessThan(chunks.reduce((n, c) => n + c.byteLength, 0)); + expect(received).toBeLessThanOrEqual(maxBytes + small.byteLength); + }); + + test("3xx redirect status rejects without following", async () => { + let resDestroyed = false; + const requestMock = mock(( + _options: unknown, + onResponse?: (res: EventEmitter & { + statusCode: number; + headers: Record; + destroy: () => void; + }) => void, + ) => { + const req = new EventEmitter() as EventEmitter & { setTimeout: () => void; end: () => void; destroy: () => void }; + req.setTimeout = () => {}; + req.destroy = () => {}; + req.end = () => { + const res = new EventEmitter() as EventEmitter & { + statusCode: number; + headers: Record; + destroy: () => void; + }; + res.statusCode = 301; + res.headers = { location: "https://evil.example/redirect" }; + res.destroy = () => { resDestroyed = true; }; + queueMicrotask(() => onResponse?.(res)); + }; + return req; + }); + mock.module("node:https", () => ({ default: { request: requestMock }, request: requestMock })); + + const { pinnedHttpsGet } = await import("../../src/images/artifacts"); + await expect(pinnedHttpsGet( + "https://cdn.example/redirect.png", + { address: "93.184.216.34", family: 4 }, + )).rejects.toThrow(/301/); + expect(resDestroyed).toBe(true); + }); + + test("non-2xx destroys the transport immediately without buffering body chunks", async () => { + // Regression: a 500 that keeps emitting must not resolve a streaming Response + // whose body nobody will read — destroy on status and reject before any data + // listener is attached. + let dataListeners = 0; + let reqDestroyed = false; + let resDestroyed = false; + const requestMock = mock(( + _options: unknown, + onResponse?: (res: EventEmitter & { + statusCode: number; + headers: Record; + setTimeout: Function; + resume: Function; + destroy: Function; + on: Function; + }) => void, + ) => { + const req = new EventEmitter() as EventEmitter & { + setTimeout: Function; + end: Function; + destroy: Function; + }; + req.setTimeout = mock(() => {}); + req.destroy = mock(() => { reqDestroyed = true; }); + req.end = mock(() => { + const res = new EventEmitter() as EventEmitter & { + statusCode: number; + headers: Record; + setTimeout: Function; + resume: Function; + destroy: Function; + }; + res.statusCode = 500; + res.headers = { "content-type": "text/plain" }; + res.setTimeout = mock(() => {}); + res.resume = mock(() => {}); + res.destroy = mock(() => { resDestroyed = true; }); + const originalOn = res.on.bind(res); + res.on = ((event: string | symbol, listener: (...args: unknown[]) => void) => { + if (event === "data") dataListeners += 1; + return originalOn(event, listener); + }) as typeof res.on; + queueMicrotask(() => { + onResponse?.(res); + // Keep dumping body after headers — must not be buffered by pinnedHttpsGet. + queueMicrotask(() => { + for (let i = 0; i < 32; i++) res.emit("data", Buffer.alloc(64 * 1024, 7)); + res.emit("end"); + }); + }); + }); + return req; + }); + mock.module("node:https", () => ({ default: { request: requestMock }, request: requestMock })); + + const { pinnedHttpsGet } = await import("../../src/images/artifacts"); + await expect(pinnedHttpsGet( + "https://cdn.example/fail.png", + { address: "93.184.216.34", family: 4 }, + )).rejects.toThrow(/image download failed: 500/); + + expect(resDestroyed).toBe(true); + expect(reqDestroyed).toBe(true); + expect(dataListeners).toBe(0); + }); + + test("idle timeout fires when no AbortSignal is supplied", async () => { + const requestMock = mock(( + _options: unknown, + _onResponse?: Function, + ) => { + const req = new EventEmitter() as EventEmitter & { + setTimeout: (ms: number, cb: () => void) => void; + end: Function; + destroy: Function; + }; + req.destroy = mock(() => {}); + req.setTimeout = (_ms, cb) => { queueMicrotask(cb); }; + req.end = mock(() => { /* never respond */ }); + return req; + }); + mock.module("node:https", () => ({ default: { request: requestMock }, request: requestMock })); + + const { pinnedHttpsGet } = await import("../../src/images/artifacts"); + await expect(pinnedHttpsGet( + "https://cdn.example/hang.png", + { address: "93.184.216.34", family: 4 }, + undefined, + { idleTimeoutMs: 1 }, + )).rejects.toThrow(/timed out/); + }); + + test("forwards idleTimeoutMs to request and response socket timers", async () => { + let reqIdleMs: number | undefined; + let resIdleMs: number | undefined; + const requestMock = mock(( + _options: unknown, + onResponse?: (res: EventEmitter & { + statusCode: number; + headers: Record; + setTimeout: (ms: number, cb: () => void) => void; + resume: () => void; + }) => void, + ) => { + const req = new EventEmitter() as EventEmitter & { + setTimeout: (ms: number, cb: () => void) => void; + end: () => void; + destroy: () => void; + }; + req.destroy = () => {}; + req.setTimeout = (ms) => { reqIdleMs = ms; }; + req.end = () => { + const res = new EventEmitter() as EventEmitter & { + statusCode: number; + headers: Record; + setTimeout: (ms: number, cb: () => void) => void; + resume: () => void; + }; + res.statusCode = 200; + res.headers = { "content-type": "image/png" }; + res.resume = () => {}; + res.setTimeout = (ms) => { resIdleMs = ms; }; + queueMicrotask(() => { + onResponse?.(res); + queueMicrotask(() => res.emit("end")); + }); + }; + return req; + }); + mock.module("node:https", () => ({ default: { request: requestMock }, request: requestMock })); + + const { pinnedHttpsGet } = await import("../../src/images/artifacts"); + const resp = await pinnedHttpsGet( + "https://cdn.example/img.png", + { address: "93.184.216.34", family: 4 }, + undefined, + { idleTimeoutMs: 12_345 }, + ); + expect(reqIdleMs).toBe(12_345); + expect(resIdleMs).toBe(12_345); + await resp.arrayBuffer(); + }); +}); diff --git a/tests/images/plan.test.ts b/tests/images/plan.test.ts new file mode 100644 index 000000000..922f09936 --- /dev/null +++ b/tests/images/plan.test.ts @@ -0,0 +1,192 @@ +import { afterAll, beforeAll, beforeEach, 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; +let planImageBridge: typeof import("../../src/images/plan")["planImageBridge"]; +let MAX_IMAGE_TIMEOUT_MS: typeof import("../../src/images/plan")["MAX_IMAGE_TIMEOUT_MS"]; + +/** Mutable token that the mocked getValidAccessToken resolves to. */ +let tokenResult: string | null = null; + +beforeAll(async () => { + process.env.OPENCODEX_HOME = join(tmpdir(), "ocx-test-" + randomUUID()); + const actualOauth = await import("../../src/oauth/index"); + mock.module("../../src/oauth/index", () => ({ + ...actualOauth, + getValidAccessToken: async () => tokenResult, + })); + ({ planImageBridge, MAX_IMAGE_TIMEOUT_MS } = await import("../../src/images/plan")); +}); +afterAll(() => { if (PREV_HOME === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = PREV_HOME; mock.restore(); }); + +beforeEach(() => { + tokenResult = null; +}); + +function makeConfig( + providers: Record>, + images?: { bridgeEnabled?: boolean; bridgeModel?: string; timeoutMs?: number }, +): OcxConfig { + return { + port: 0, + defaultProvider: "test", + providers: Object.fromEntries( + Object.entries(providers).map(([k, v]) => [k, { adapter: "openai-chat", 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-chat", baseUrl: "https://api.anthropic.com" } as OcxProviderConfig; +const openaiRouted = { adapter: "openai-chat", 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("images.timeoutMs is forwarded onto the plan", async () => { + const cfg = makeConfig( + { xai: { baseUrl: "https://api.x.ai", apiKey: "test-token" } }, + { bridgeEnabled: true, timeoutMs: 120_000 }, + ); + const plan = await planImageBridge(cfg, makeParsed(true), routed); + expect(plan!.timeoutMs).toBe(120_000); + }); + + test("images.timeoutMs above ceiling is clamped", async () => { + const cfg = makeConfig( + { xai: { baseUrl: "https://api.x.ai", apiKey: "test-token" } }, + { bridgeEnabled: true, timeoutMs: 999_999_999 }, + ); + const plan = await planImageBridge(cfg, makeParsed(true), routed); + expect(plan!.timeoutMs).toBe(MAX_IMAGE_TIMEOUT_MS); + }); + + 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; + }); + + test("authMode oauth prefers OAuth over a stale apiKey", async () => { + tokenResult = "oauth-token"; + const cfg = makeConfig( + { xai: { baseUrl: "https://api.x.ai/v1", apiKey: "stale-key", authMode: "oauth" } }, + { bridgeEnabled: true }, + ); + const plan = await planImageBridge(cfg, makeParsed(true), routed); + expect(plan!.auth.token).toBe("oauth-token"); + tokenResult = null; + }); + + test("authMode key does not fall back to stored OAuth", async () => { + tokenResult = "oauth-token"; + const cfg = makeConfig( + { xai: { baseUrl: "https://api.x.ai/v1", apiKey: "", authMode: "key" } }, + { bridgeEnabled: true }, + ); + expect(await planImageBridge(cfg, makeParsed(true), routed)).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..16754d06e --- /dev/null +++ b/tests/images/synthetic-tool.test.ts @@ -0,0 +1,78 @@ +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' → false without hosted image tool context", () => { + expect(isImageGenName("imagegen")).toBe(false); + }); + + 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("generate_image alone → undefined (alias requires hosted entry)", () => { + expect( + extractHostedImageGeneration([{ type: "function", name: "generate_image", parameters: { type: "object" } }]), + ).toBeUndefined(); + }); + + test("generate_image with hosted image_generation → matched", () => { + const result = extractHostedImageGeneration([ + { type: "image_generation" }, + { type: "function", name: "generate_image", parameters: { type: "object" } }, + ]); + expect(result).toBeDefined(); + expect(result!.toolNames.has("generate_image")).toBe(true); + expect(result!.toolNames.has("image_generation")).toBe(true); + }); + + 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..b13143d59 --- /dev/null +++ b/tests/images/xai-client.test.ts @@ -0,0 +1,144 @@ +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("custom timeoutMs is composed into the abort signal", async () => { + const calls = stubFetch(200, { data: [{ b64_json: "dGVzdA==" }] }); + await callXaiImages({ prompt: "x" }, AUTH, undefined, 5_000); + const passed = calls[0]!.init?.signal as AbortSignal; + expect(passed).toBeDefined(); + expect(passed.aborted).toBe(false); + }); + + test("timeoutMs composes a deadline that aborts the fetch signal", async () => { + let seenSignal: AbortSignal | undefined; + globalThis.fetch = (async (_input, init) => { + seenSignal = init?.signal; + return new Response(JSON.stringify({ data: [{ b64_json: "dGVzdA==" }] }), { + status: 200, + headers: { "content-type": "application/json" }, + }); + }) as typeof fetch; + await callXaiImages({ prompt: "x" }, AUTH, undefined, 50); + expect(seenSignal).toBeDefined(); + expect(seenSignal!.aborted).toBe(false); + await new Promise(resolve => setTimeout(resolve, 60)); + expect(seenSignal!.aborted).toBe(true); + }); + + test("trailing slash on baseUrl does not produce double-slash URL", async () => { + const calls = stubFetch(200, { data: [{ b64_json: "dGVzdA==" }] }); + await callXaiImages({ prompt: "x" }, { baseUrl: "https://api.x.ai/v1/", token: "test-token" }); + expect(calls[0]!.url).toBe("https://api.x.ai/v1/images/generations"); + }); + + 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/images/z-fulfill.test.ts b/tests/images/z-fulfill.test.ts new file mode 100644 index 000000000..fc5501b91 --- /dev/null +++ b/tests/images/z-fulfill.test.ts @@ -0,0 +1,197 @@ +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; +let fulfillImageCall: typeof import("../../src/images/fulfill")["fulfillImageCall"]; + +beforeAll(async () => { + process.env.OPENCODEX_HOME = join(tmpdir(), "ocx-test-" + randomUUID()); + mock.restore(); + mock.module("../../src/images/xai-client", () => ({ + callXaiImages: async (req: XaiImageRequest, _auth: unknown, _signal?: AbortSignal, timeoutMs?: number) => { + xaiCalls.push(req); + capturedTimeoutMs = timeoutMs; + if (xaiError) throw xaiError; + return xaiResult; + }, + })); + mock.module("../../src/images/artifacts", () => ({ + createImageBudget: () => ({ spent: 0 }), + materializeInlineImage: async () => materializeFn(matIdx++), + downloadImageToArtifact: async () => downloadFn(dlIdx++), + })); + ({ fulfillImageCall } = await import(`../../src/images/fulfill?fulfill=${Date.now()}`)); +}); +afterAll(() => { if (PREV_HOME === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = PREV_HOME; mock.restore(); }); + +// --- 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`; + +let capturedTimeoutMs: number | undefined; + +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; + capturedTimeoutMs = undefined; + 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("plan.timeoutMs is forwarded to callXaiImages", async () => { + reset(); + const timedPlan = { ...plan, timeoutMs: 12_345 } as ImageBridgePlan; + await fulfillImageCall( + { id: "c1", name: "image_gen", arguments: JSON.stringify({ prompt: "a cat" }) }, + timedPlan, { spent: 0 }, + ); + expect(capturedTimeoutMs).toBe(12_345); + }); + + 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"); + }); + + test("plan.defaultSize and defaultQuality fill omitted args", async () => { + reset(); + const sizedPlan = { + ...plan, + defaultSize: "1024x1024", + defaultQuality: "hd", + } as ImageBridgePlan; + await fulfillImageCall( + { id: "c1", name: "image_gen", arguments: JSON.stringify({ prompt: "a cat" }) }, + sizedPlan, { spent: 0 }, + ); + expect(xaiCalls[0]!.size).toBe("1024x1024"); + expect(xaiCalls[0]!.quality).toBe("hd"); + }); + + test("explicit size/quality override plan defaults", async () => { + reset(); + const sizedPlan = { + ...plan, + defaultSize: "1024x1024", + defaultQuality: "hd", + } as ImageBridgePlan; + await fulfillImageCall( + { id: "c1", name: "image_gen", arguments: JSON.stringify({ prompt: "a cat", size: "512x512", quality: "standard" }) }, + sizedPlan, { spent: 0 }, + ); + expect(xaiCalls[0]!.size).toBe("512x512"); + expect(xaiCalls[0]!.quality).toBe("standard"); + }); +}); diff --git a/tests/images/z-handler-activation.test.ts b/tests/images/z-handler-activation.test.ts new file mode 100644 index 000000000..324b87732 --- /dev/null +++ b/tests/images/z-handler-activation.test.ts @@ -0,0 +1,214 @@ +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; + +// --- 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; + +let handleResponses: typeof import("../../src/server/responses")["handleResponses"]; + +beforeAll(async () => { + process.env.OPENCODEX_HOME = join(tmpdir(), "ocx-test-" + randomUUID()); + + 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; + }, + })); + + const actualLoop = await import("../../src/images/loop"); + mock.module("../../src/images/loop", () => ({ + ...actualLoop, + runWithImageBridge: async () => { + imageBridgeRun = true; + return new Response("data: {\"type\":\"done\"}\n\n", { + status: 200, headers: { "content-type": "text/event-stream" }, + }); + }, + })); + + mock.module("../../src/web-search/index", () => ({ + 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, + })); + + ({ handleResponses } = await import("../../src/server/responses")); +}); + +afterAll(() => { + if (PREV_HOME === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = PREV_HOME; + mock.restore(); +}); + +/** 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-chat", baseUrl: "https://api.x.ai/v1", 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); + 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; + 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; + 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 → image bridge wins (web-search loop has no runTurn support)", async () => { + imageBridgeRun = false; webSearchRun = false; runTurnCalled = false; + useRunTurnAdapter = true; + mockWsPlan = { backend: "openai" }; + try { + const res = await post(true, [{ type: "web_search" }, { type: "image_generation" }]); + expect(webSearchRun).toBe(false); + expect(imageBridgeRun).toBe(true); + expect(runTurnCalled).toBe(false); + expect(res.headers.get("content-type")).toBe("text/event-stream"); + } finally { + useRunTurnAdapter = false; + } + }); + + test("image-only on a runTurn adapter → image bridge activates before runTurn early-return", async () => { + imageBridgeRun = false; webSearchRun = false; runTurnCalled = false; + useRunTurnAdapter = true; + mockWsPlan = undefined; + try { + const res = await post(true, [{ type: "image_generation" }]); + expect(imageBridgeRun).toBe(true); + expect(webSearchRun).toBe(false); + expect(runTurnCalled).toBe(false); + expect(res.headers.get("content-type")).toBe("text/event-stream"); + } finally { + useRunTurnAdapter = false; + } + }); +}); 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(); diff --git a/tests/storage-cleanup.test.ts b/tests/storage-cleanup.test.ts index 85bbc0d6e..b82584241 100644 --- a/tests/storage-cleanup.test.ts +++ b/tests/storage-cleanup.test.ts @@ -661,6 +661,8 @@ describe("executeArchivedCleanup", () => { expect(Buffer.compare(beforeState, readFileSync(join(home, "state_5.sqlite")))).toBe(0); }); + // Windows CI: injected satellite rollback paths (especially goals) can measure 6–13s + // there and trip bun's default 5s harness timeout. test.each([ ["failAfterLogsMutation", { failAfterLogsMutation: true }], ["failAfterMemoriesMutation", { failAfterMemoriesMutation: true }], @@ -714,6 +716,7 @@ describe("executeArchivedCleanup", () => { expect(stateAfter.query("SELECT id, rollout_path, archived FROM threads ORDER BY id").all()).toEqual(threads); stateAfter.close(); }, + { timeout: 30_000 }, ); test("satellite restore failure keeps recovery trashDir and manifest", () => {