From 7393805c55f90c1ee37e6b793ccd43d6071b4fc7 Mon Sep 17 00:00:00 2001 From: Your Name Date: Sat, 25 Jul 2026 04:38:07 +0800 Subject: [PATCH 01/15] feat(images): add Grok image bridge for non-OpenAI models MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When a non-OpenAI model receives an image_generation hosted tool from Codex, OpenAI's server-side execution is unavailable. This intercepts the tool call, routes it to xAI Grok Imagine, materializes the image to ~/.opencodex/artifacts/, and feeds the result back to the model — all inside a streaming SSE loop. Architecture mirrors the proven web-search sidecar pattern: - parser.ts: extract hosted image_generation tool, stash into parsed._imageGeneration - plan.ts: decide whether to activate (xAI provider + token available, non-OpenAI route) - loop.ts: agentic loop (max 3 rounds) — intercept image tool calls, fulfill via xAI, inject results, re-call model, bridge to SSE - synthetic-tool.ts: buildImageTool() injection + isImageGenName() detection - xai-client.ts: xAI /images/generations and /images/edits API client - fulfill.ts: execute single image call, materialize to disk (never throws) - artifacts.ts: extended with downloadImageToArtifact() + magic byte format detection New files: src/images/{types,xai-client,synthetic-tool,fulfill,plan,loop,index}.ts tests/images/{xai-client,plan,fulfill,synthetic-tool,loop}.test.ts docs-site/src/content/docs/guides/image-bridge.md Modified files: src/types.ts: +_imageGeneration stash, +OcxTool.imageGeneration flag, +OcxImagesConfig fields src/responses/parser.ts: extractHostedImageGeneration extraction src/server/responses/core.ts: image bridge trigger point src/images/artifacts.ts: +downloadImageToArtifact, +guessExtFromMagic, shared helpers --- .../src/content/docs/guides/image-bridge.md | 67 +++ src/images/artifacts.ts | 140 ++++++ src/images/fulfill.ts | 75 +++ src/images/index.ts | 4 + src/images/loop.ts | 430 ++++++++++++++++++ src/images/plan.ts | 64 +++ src/images/synthetic-tool.ts | 70 +++ src/images/types.ts | 19 + src/images/xai-client.ts | 133 ++++++ src/lib/destination-policy.ts | 51 ++- src/responses/parser.ts | 6 + src/server/responses/core.ts | 31 ++ src/types.ts | 10 + tests/images/artifacts-ssrf.test.ts | 51 +++ tests/images/fulfill.test.ts | 146 ++++++ tests/images/loop.test.ts | 183 ++++++++ tests/images/plan.test.ts | 143 ++++++ tests/images/synthetic-tool.test.ts | 62 +++ tests/images/xai-client.test.ts | 114 +++++ tests/responses-parser.test.ts | 14 + 20 files changed, 1812 insertions(+), 1 deletion(-) create mode 100644 docs-site/src/content/docs/guides/image-bridge.md create mode 100644 src/images/artifacts.ts create mode 100644 src/images/fulfill.ts create mode 100644 src/images/index.ts create mode 100644 src/images/loop.ts create mode 100644 src/images/plan.ts create mode 100644 src/images/synthetic-tool.ts create mode 100644 src/images/types.ts create mode 100644 src/images/xai-client.ts create mode 100644 tests/images/artifacts-ssrf.test.ts create mode 100644 tests/images/fulfill.test.ts create mode 100644 tests/images/loop.test.ts create mode 100644 tests/images/plan.test.ts create mode 100644 tests/images/synthetic-tool.test.ts create mode 100644 tests/images/xai-client.test.ts 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 0000000000..0f7102f3e2 --- /dev/null +++ b/docs-site/src/content/docs/guides/image-bridge.md @@ -0,0 +1,67 @@ +--- +title: Image Bridge +description: Route image_generation hosted-tool calls to xAI Grok Imagine when using a non-OpenAI provider. +--- + +## Overview + +When you route Codex through a non-OpenAI model (Claude, Gemini, Grok, etc.), the +`image_generation` **hosted tool** normally doesn't work — it requires OpenAI's server-side +execution environment. The Image Bridge detects these calls and transparently reroutes them to +xAI Grok Imagine, so the model you're actually chatting with can still generate images. + +## Prerequisites + +- **Enable the bridge** by setting `images.bridgeEnabled: true` in your config (it is off by + default to avoid unexpected xAI charges — see [Configuration](#configuration) below). +- An xAI provider configured in settings with `baseUrl: "https://api.x.ai/v1"` and the + `openai-chat` adapter. +- Authentication via `authMode: "oauth"` (`ocx login xai` — uses a stored, auto-refreshed + bearer token) or `authMode: "key"` (a configured API key). +- A non-OpenAI model selected as your active provider. (When the active provider is OpenAI, + the native hosted tool is used directly and the bridge is bypassed.) + +## Configuration + +Image Bridge options live under `images` in `~/.opencodex/config.json`. Bridging is +**opt-in** — you must set `bridgeEnabled: true` to enable paid xAI Grok Imagine generation: + +```json +{ + "images": { + "bridgeEnabled": true, + "bridgeModel": "grok-imagine-image-quality", + "maxRounds": 3 + } +} +``` + +| Option | Default | Description | +| --- | --- | --- | +| `bridgeEnabled` | `false` | Master switch. Set `true` to enable bridging. Off by default to avoid unexpected xAI charges. | +| `bridgeModel` | `grok-imagine-image-quality` | The xAI image model id to send prompts to. | +| `maxRounds` | `3` | Maximum number of image-generation loop iterations per turn. | + +## How It Works + +1. When Codex sends a request with `image_generation` in the tools array, OpenCodex detects it + during request preprocessing. +2. The hosted tool is replaced with a **synthetic function tool** that the routed model can call + normally — the model sees a callable tool rather than an opaque hosted tool it can't execute. +3. When the model invokes that tool, OpenCodex intercepts the call and sends the prompt to xAI's + image generation API. +4. Generated images are saved to `~/.opencodex/artifacts/` and the **local file path** is returned + to the model as the tool result. +5. The model continues the conversation with knowledge of the generated image and its location. + +From the model's perspective nothing changed — it called a tool and got a result. From the user's +perspective, image generation works with any routed provider instead of silently failing. + +## Limitations + +- **Only xAI Grok Imagine is supported.** DALL-E and other image providers may be added later. +- **Web search takes priority.** If both web search and image generation are requested in the same + turn, the web-search bridge runs and image generation is skipped for that turn. +- **xAI costs apply.** Image generation via xAI requires an active xAI subscription or API credits. +- **Streaming only.** The bridge works by intercepting the SSE response stream; requests with + `stream: false` are rejected with a 400 error. diff --git a/src/images/artifacts.ts b/src/images/artifacts.ts new file mode 100644 index 0000000000..f4eb700461 --- /dev/null +++ b/src/images/artifacts.ts @@ -0,0 +1,140 @@ +import { mkdir, writeFile } from "node:fs/promises"; +import { join } from "node:path"; +import { getConfigDir } from "../config"; +import { assessUrlDestination, assertUrlResolvesPublic } from "../lib/destination-policy"; + +const MAX_DECODED_BYTES_PER_IMAGE = 50 * 1024 * 1024; +const MAX_DECODED_BYTES_PER_RESPONSE = 100 * 1024 * 1024; +const MAX_DOWNLOAD_BYTES = 50 * 1024 * 1024; // 50 MiB + +// Strict alphabet check: Buffer.from(..., "base64") silently ignores invalid +// characters, so malformed payloads would otherwise decode to garbage bytes. +const BASE64_RE = /^[A-Za-z0-9+/]*={0,2}$/; + +export interface ImageBudget { + spent: number; +} + +export function createImageBudget(): ImageBudget { + return { spent: 0 }; +} + +function getArtifactsDir(): string { + return join(getConfigDir(), "artifacts"); +} + +function timestampPrefix(): string { + const now = new Date(); + return [ + now.getFullYear(), + String(now.getMonth() + 1).padStart(2, "0"), + String(now.getDate()).padStart(2, "0"), + "-", + String(now.getHours()).padStart(2, "0"), + String(now.getMinutes()).padStart(2, "0"), + String(now.getSeconds()).padStart(2, "0"), + "-", + String(now.getMilliseconds()).padStart(3, "0"), + ].join(""); +} + +export function guessExtFromMagic(bytes: Uint8Array): string { + const sig = Buffer.from(bytes.slice(0, 12)).toString("latin1"); + if (sig.startsWith("\x89PNG")) return "png"; + if (sig.startsWith("\xff\xd8\xff")) return "jpg"; + if (sig.startsWith("RIFF") && sig.slice(8, 12) === "WEBP") return "webp"; + if (sig.startsWith("GIF8")) return "gif"; + return "png"; +} + +export async function materializeInlineImage( + base64Data: string, + budget?: ImageBudget, +): Promise { + const dir = getArtifactsDir(); + await mkdir(dir, { recursive: true, mode: 0o700 }); + + const normalized = base64Data.replace(/\s+/g, ""); + if (!BASE64_RE.test(normalized) || normalized.length % 4 !== 0) { + throw new Error("inline image data is not valid base64"); + } + // Validate decoded size from the base64 length *before* allocating a Buffer, so a + // malicious or broken upstream cannot force a large allocation / OOM. + const padding = normalized.endsWith("==") ? 2 : normalized.endsWith("=") ? 1 : 0; + const decodedBytes = (normalized.length / 4) * 3 - padding; + if (decodedBytes === 0) throw new Error("inline image data is empty after base64 decode"); + if (decodedBytes > MAX_DECODED_BYTES_PER_IMAGE) throw new Error(`inline image exceeds ${MAX_DECODED_BYTES_PER_IMAGE} byte per-image cap`); + if (budget && budget.spent + decodedBytes > MAX_DECODED_BYTES_PER_RESPONSE) { + throw new Error(`inline image response exceeds ${MAX_DECODED_BYTES_PER_RESPONSE} byte per-response cap`); + } + + const buf = Buffer.from(normalized, "base64"); + if (budget) budget.spent += buf.length; + + // Sniff actual format from decoded bytes rather than trusting the declared mimeType. + const ext = guessExtFromMagic(buf); + const filePath = join(dir, `img-${timestampPrefix()}-${crypto.randomUUID()}.${ext}`); + await writeFile(filePath, buf, { mode: 0o600 }); + return filePath; +} + +export async function downloadImageToArtifact( + url: string, + budget?: ImageBudget, + signal?: AbortSignal, +): 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. + // Reject non-HTTP(S) schemes, literal private/loopback/link-local/metadata addresses. + const assessment = assessUrlDestination(url); + if (assessment && assessment.kind !== "public" && assessment.kind !== "hostname") { + throw new Error(`image URL targets ${assessment.detail}`); + } + // DNS check: resolve hostname and reject if it points at private/internal space. + await assertUrlResolvesPublic(url); + const resp = await fetch(url, { signal, redirect: "error" }); + if (!resp.ok) throw new Error("image download failed: " + resp.status); + + // Stream the body with a hard byte cap so a missing/lying Content-Length or a + // compromised CDN URL cannot exhaust memory before the size check runs. + if (!resp.body) throw new Error("image download returned no body"); + const reader = resp.body.getReader(); + const chunks: Uint8Array[] = []; + let total = 0; + try { + for (;;) { + const { value, done } = await reader.read(); + if (done) break; + total += value.byteLength; + if (total > MAX_DOWNLOAD_BYTES) { + throw new Error(`image download exceeds ${MAX_DOWNLOAD_BYTES} byte cap`); + } + chunks.push(value); + } + } finally { + try { await reader.cancel(); } catch { /* ignore cancel errors */ } + reader.releaseLock(); + } + + const bytes = new Uint8Array(total); + let offset = 0; + for (const c of chunks) { bytes.set(c, offset); offset += c.byteLength; } + + if (budget && budget.spent + bytes.length > MAX_DECODED_BYTES_PER_RESPONSE) { + throw new Error(`image download exceeds ${MAX_DECODED_BYTES_PER_RESPONSE} byte per-response budget`); + } + + const ext = guessExtFromMagic(bytes); + const dir = getArtifactsDir(); + await mkdir(dir, { recursive: true, mode: 0o700 }); + if (budget) budget.spent += bytes.length; + + const filePath = join(dir, `dl-${timestampPrefix()}-${crypto.randomUUID()}.${ext}`); + await writeFile(filePath, bytes, { mode: 0o600 }); + return filePath; +} diff --git a/src/images/fulfill.ts b/src/images/fulfill.ts new file mode 100644 index 0000000000..4751f35ac9 --- /dev/null +++ b/src/images/fulfill.ts @@ -0,0 +1,75 @@ +import type { ImageBridgePlan, ImageCallResult } from "./types"; +import { callXaiImages } from "./xai-client"; +import { materializeInlineImage, downloadImageToArtifact, type ImageBudget } from "./artifacts"; + +/** + * Fulfill ONE image-generation tool call end-to-end: parse args, call xAI, materialize the returned + * images to disk, and return a structured result. NEVER throws — all errors become `{ ok: false }` + * so the caller can inject the error as a tool result and let the model respond gracefully. + */ +export async function fulfillImageCall( + call: { id: string; name: string; arguments: string }, + plan: ImageBridgePlan, + budget: ImageBudget, + signal?: AbortSignal, +): Promise { + let args: unknown; + try { + args = JSON.parse(call.arguments || "{}"); + } catch { + return { ok: false, model: plan.model, prompt: "", files: [], count: 0, error: "invalid arguments JSON" }; + } + if (typeof args !== "object" || args === null) { + return { ok: false, model: plan.model, prompt: "", files: [], count: 0, error: "invalid arguments JSON" }; + } + const obj = args as Record; + + const prompt = + typeof obj.prompt === "string" ? obj.prompt : typeof obj.input === "string" ? obj.input : ""; + if (!prompt) { + return { ok: false, model: plan.model, prompt: "", files: [], count: 0, error: "missing prompt" }; + } + + const n = typeof obj.n === "number" ? Math.max(1, Math.min(4, Math.floor(obj.n))) : 1; + const imageUrl = + typeof obj.image_url === "string" ? obj.image_url : typeof obj.image === "string" ? obj.image : undefined; + const size = typeof obj.size === "string" ? obj.size : undefined; + const quality = typeof obj.quality === "string" ? obj.quality : undefined; + + let result; + try { + result = await callXaiImages({ prompt, model: plan.model, n, imageUrl, size, quality }, plan.auth, signal); + } catch (e) { + const error = e instanceof Error ? e.message : String(e); + return { ok: false, model: plan.model, prompt, files: [], count: 0, error }; + } + + const files: string[] = []; + for (const img of result.images ?? []) { + try { + if (img.b64_json) { + files.push(await materializeInlineImage(img.b64_json, budget)); + } else if (img.url) { + files.push(await downloadImageToArtifact(img.url, budget, signal)); + } + } catch (e) { + // Partial success is OK — silently skip this image and continue. + console.warn(`[images] failed to materialize image: ${e instanceof Error ? e.message : String(e)}`); + } + } + + 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 0000000000..6f0adb1fd2 --- /dev/null +++ b/src/images/index.ts @@ -0,0 +1,4 @@ +export { planImageBridge, findXaiProvider, resolveXaiToken } from "./plan"; +export { runWithImageBridge } from "./loop"; +export type { ImageBridgePlan, ImageCallResult } from "./types"; +export { buildImageTool, extractHostedImageGeneration, IMAGE_GEN_TOOL_NAME, isImageGenName } from "./synthetic-tool"; diff --git a/src/images/loop.ts b/src/images/loop.ts new file mode 100644 index 0000000000..e5b2c57865 --- /dev/null +++ b/src/images/loop.ts @@ -0,0 +1,430 @@ +/** + * Image bridge agentic loop — adapted from src/web-search/loop.ts but significantly simpler. + * + * The routed (non-OpenAI) model runs in a bounded loop. Each iteration is streamed and fully + * buffered internally. If the model calls an image-generation tool, the bridge fulfills it via + * the xAI sidecar, injects the result as a tool_result, and loops (bounded by maxRounds). When + * the model produces a real tool call or the budget is exhausted, the passthrough events are + * replayed to the bridge for final SSE output. + * + * Removed vs web-search: no sidecar backend selection, no 429 key-failover, no forced-answer + * nudge, no failed-query dedup, no describeImages/structuredOutput, no recordSidecarOutcome. + */ +import type { ProviderAdapter } from "../adapters/base"; +import { createAdapterEventQueue } from "../adapters/run-turn-queue"; +import type { AdapterEvent, OcxMessage, OcxParsedRequest, OcxThinkingContent } from "../types"; +import { namespacedToolName } from "../types"; +import { bridgeToResponsesSSE } from "../bridge"; +import { clearableDeadline } from "../lib/abort"; +import { readBoundedResponseBody } from "../lib/bounded-body"; +import { fetchWithResetRetry } from "../lib/upstream-retry"; +import { parseStreamWithProgress, RoutedModelInactivityError, WebSearchStreamProtocolError } from "../web-search/progress-stream"; +import { fulfillImageCall } from "./fulfill"; +import { createImageBudget } from "./artifacts"; +import type { ImageBridgePlan } from "./types"; + +const SSE_HEADERS = { + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", +}; + +const CONNECT_TIMEOUT_MS = 200_000; +const STALL_TIMEOUT_MS = 200_000; +const DEFAULT_MAX_ROUNDS = 3; + +interface ImageCall { + id: string; + name: string; + args: string; +} + +/** + * Split an iteration's adapter events into (a) the image-generation tool calls to intercept and + * (b) the events to pass through to Codex. An image tool-call's own start/delta/end events are + * dropped (Codex never sees the synthetic tool); every other event — text, thinking, real tool + * calls, done — is preserved in order. + */ +function scanEventsForImageCall(events: AdapterEvent[], toolNames: Set): { + calls: ImageCall[]; + passthrough: AdapterEvent[]; + hasRealToolCall: boolean; +} { + const calls: ImageCall[] = []; + const passthrough: AdapterEvent[] = []; + let hasRealToolCall = false; + let pending: { name: string; id: string; argsBuf: string; events: AdapterEvent[] } | null = null; + const flushPending = (): void => { + if (pending && !toolNames.has(pending.name)) { + passthrough.push(...pending.events); + hasRealToolCall = true; + } + pending = null; + }; + for (const e of events) { + if (e.type === "tool_call_start") { + flushPending(); + pending = { name: e.name, id: e.id, argsBuf: "", events: [e] }; + } else if (e.type === "tool_call_delta" && pending) { + pending.argsBuf += e.arguments; + pending.events.push(e); + } else if (e.type === "tool_call_end" && pending) { + pending.events.push(e); + if (toolNames.has(pending.name)) { + calls.push({ id: pending.id, name: pending.name, args: pending.argsBuf }); + } else { + passthrough.push(...pending.events); + hasRealToolCall = true; + } + pending = null; + } else { + flushPending(); + passthrough.push(e); + } + } + flushPending(); + return { calls, passthrough, hasRealToolCall }; +} + +async function* replay(events: AdapterEvent[]): AsyncGenerator { + for (const e of events) yield e; +} + +/** + * Collect the thinking block that preceded an image tool call, so the replayed assistant turn can + * carry it. Anthropic extended thinking REQUIRES the assistant message containing tool_use to start + * with its signed thinking blocks. + */ +function extractIterationThinking(events: AdapterEvent[]): OcxThinkingContent | null { + let thinking = ""; + let signature: string | undefined; + const redacted: string[] = []; + for (const e of events) { + if (e.type === "thinking_delta") thinking += e.thinking; + else if (e.type === "thinking_signature") signature = e.signature; + else if (e.type === "redacted_thinking") redacted.push(e.data); + } + if (!thinking && !signature && redacted.length === 0) return null; + return { + type: "thinking", + thinking, + ...(signature ? { signature } : {}), + ...(redacted.length > 0 ? { redacted } : {}), + }; +} + +function jsonError(status: number, message: string): Response { + return new Response(JSON.stringify({ error: { message, type: "upstream_error", code: null } }), { + status, + headers: { "Content-Type": "application/json" }, + }); +} + +/** Hard provider/parse failure inside an iteration. The eager first iteration converts it to a + * non-2xx jsonError; later (already-streaming) iterations surface it as an in-stream error event. */ +class LoopError extends Error { + constructor(readonly status: number, message: string) { + super(message); + this.name = "LoopError"; + } +} + +export interface ImageBridgeDeps { + parsed: OcxParsedRequest; + adapter: ProviderAdapter; + plan: ImageBridgePlan; + /** Headers forwarded from the original request (e.g. Codex auth). Cloned per iteration. */ + forwardHeaders?: Headers; + /** Called before each routed-model dispatch in the bridge loop, for attempt telemetry. */ + onAttemptSend?: () => void; + abortSignal?: AbortSignal; + onFirstOutput?: () => void; + /** Max image-generation rounds before forcing a final answer. Defaults to 3. */ + maxRounds?: number; +} + +/** + * Run the main (non-OpenAI) model in a small agentic loop. Each upstream iteration is streamed and + * fully buffered internally so raw byte progress is observable without leaking the synthetic tool or + * preliminary assistant output. If the model invokes image generation, run it via the xAI sidecar, + * inject the answer as a tool_result, and loop (bounded by `maxRounds`). + */ +export async function runWithImageBridge(deps: ImageBridgeDeps): Promise { + const { parsed, adapter, plan, abortSignal } = deps; + const maxRounds = Math.max(0, deps.maxRounds ?? DEFAULT_MAX_ROUNDS); + const HARD_CAP = maxRounds + 1; + + const messages: OcxMessage[] = [...parsed.context.messages]; + const allTools = parsed.context.tools ?? []; + // For the forced-final pass we drop image tools so the model MUST answer from the results already + // in `messages` (can't generate again) — this guarantees a non-empty final answer. + const toolsNoImage = allTools.filter(t => !t.imageGeneration); + const budget = createImageBudget(); + + // Link an internal AbortController to the turn signal so a client cancel of the SSE body aborts + // in-flight model fetches AND the sidecar. + const internalAbort = new AbortController(); + const linkAbort = (): void => internalAbort.abort(abortSignal?.reason); + if (abortSignal) { + if (abortSignal.aborted) linkAbort(); + else abortSignal.addEventListener("abort", linkAbort, { once: true }); + } + const signal = internalAbort.signal; + + interface IterationResponse { + response: Response; + responseAdapter: ProviderAdapter; + } + type IterationSplit = ReturnType; + + // Acquire one iteration's final response headers. The first call is drained eagerly so an initial + // connect/header/HTTP failure stays a non-2xx JSON response. + const prepareIterationEvents = async function* (forceFinal: boolean): AsyncGenerator { + const iterParsed: OcxParsedRequest = { + ...parsed, stream: true, + context: { ...parsed.context, messages, tools: forceFinal ? toolsNoImage : allTools }, + }; + + // runTurn adapters (Cursor) own all upstream communication via an emit callback. They don't + // expose buildRequest/fetchResponse/parseStream to the bridge, so collect their events through + // an AdapterEventQueue and wrap them in a pseudo-response whose parseStream replays them. + if (adapter.runTurn) { + const queue = createAdapterEventQueue({ + onBacklogExceeded: () => internalAbort.abort("runTurn backlog exceeded"), + }); + void adapter + .runTurn( + iterParsed, + { headers: deps.forwardHeaders ? new Headers(deps.forwardHeaders) : new Headers(), abortSignal: signal }, + queue.push, + ) + .then(() => queue.close()) + .catch(err => { + queue.push({ type: "error", message: err instanceof Error ? err.message : String(err) }); + queue.close(); + }); + + const events = await queue.collect(); + deps.onAttemptSend?.(); + + // runTurn adapters signal errors via {type:"error"} events, not HTTP status codes. + const errorEvent = events.find(e => e.type === "error"); + if (errorEvent && errorEvent.type === "error") { + throw new LoopError(502, errorEvent.message); + } + + // Wrap as an adapter whose parseStream replays the collected events. The pseudo-response + // carries an empty (but non-null) body so parseStreamWithProgress can acquire a reader + // without crashing; the wrapped parseStream never reads from it. + const wrappedAdapter: ProviderAdapter = { + ...adapter, + async *parseStream() { + for (const e of events) yield e; + }, + }; + return { response: new Response(new Uint8Array(0), { status: 200 }), responseAdapter: wrappedAdapter }; + } + + const headerDeadline = clearableDeadline(CONNECT_TIMEOUT_MS, signal); + try { + const request = await adapter.buildRequest(iterParsed, { + headers: deps.forwardHeaders ? new Headers(deps.forwardHeaders) : new Headers(), + abortSignal: headerDeadline.signal, + }); + deps.onAttemptSend?.(); + const response = adapter.fetchResponse + ? await adapter.fetchResponse(request, { + abortSignal: headerDeadline.signal, + timeoutMs: CONNECT_TIMEOUT_MS, + returnRawErrors: true, + stream: true, + }) + : await fetchWithResetRetry( + () => { + const h = new Headers(request.headers); + if (!h.has("accept-encoding")) h.set("accept-encoding", "identity"); + return fetch(request.url, { + method: request.method, + headers: h, + body: request.body, + signal: headerDeadline.signal, + }); + }, + { abortSignal: headerDeadline.signal, label: "image-bridge-loop" }, + ); + + // Final headers have arrived. Clear only the deadline timer before ANY body read. + headerDeadline.clear(); + if (!response.ok) { + let body: Awaited>; + try { + body = await readBoundedResponseBody(response, { signal }); + } catch { + if (signal.aborted) throw new LoopError(499, "client closed request during image-bridge"); + throw new LoopError(response.status, `Provider error ${response.status}`); + } + let formatted = ""; + if (body.displaySafe && !body.truncated && body.text.trim() && adapter.formatErrorBody) { + try { + formatted = adapter.formatErrorBody(response.status, response.headers, body.text).trim(); + } catch { /* formatter hooks are best-effort */ } + } + const suffix = formatted ? `: ${formatted.slice(0, 400)}` : ""; + throw new LoopError(response.status, `Provider error ${response.status}${suffix}`); + } + return { response, responseAdapter: adapter }; + } catch (error) { + if (headerDeadline.didExpire()) { + throw new LoopError(504, `Provider response-header timeout after ${CONNECT_TIMEOUT_MS}ms during image-bridge`); + } + if (signal.aborted) throw new LoopError(499, "client closed request during image-bridge"); + if (error instanceof LoopError) throw error; + throw new LoopError(502, `Provider unreachable: ${error instanceof Error ? error.message : String(error)}`); + } finally { + headerDeadline.clear(); + } + }; + + const prepareIterationDrained = async (forceFinal: boolean): Promise => { + const it = prepareIterationEvents(forceFinal); + let r = await it.next(); + while (!r.done) r = await it.next(); + return r.value; + }; + + // Consume and validate one successful response body. Only invisible heartbeat events escape while + // semantic output remains buffered for safe scanning. + const consumeIterationEvents = async function* (prepared: IterationResponse): AsyncGenerator { + const events: AdapterEvent[] = []; + try { + const parse = prepared.responseAdapter.parseStream.bind(prepared.responseAdapter); + for await (const event of parseStreamWithProgress(prepared.response, parse, { + signal, + inactivityTimeoutMs: STALL_TIMEOUT_MS, + })) { + if (event.type === "heartbeat") yield event; + else events.push(event); + } + } catch (error) { + if (signal.aborted) throw new LoopError(499, "client closed request during image-bridge"); + if (error instanceof RoutedModelInactivityError) throw new LoopError(504, error.message); + if (error instanceof WebSearchStreamProtocolError) throw new LoopError(502, error.message); + throw new LoopError(502, `Provider stream error: ${error instanceof Error ? error.message : String(error)}`); + } + + const terminalIndexes = events.flatMap((event, index) => + event.type === "done" || event.type === "incomplete" || event.type === "error" ? [index] : []); + if (terminalIndexes.length !== 1 || terminalIndexes[0] !== events.length - 1) { + throw new LoopError(502, `Image-bridge adapter stream protocol error: expected one final terminal event, received ${terminalIndexes.length}`); + } + const terminal = events[terminalIndexes[0]!]; + if (terminal.type === "error") throw new LoopError(502, terminal.message); + return scanEventsForImageCall(events, plan.toolNames); + }; + + // Eagerly acquire only the FIRST iteration's final headers so connect/header/HTTP failures remain + // non-2xx JSON. + let firstPrepared: IterationResponse; + try { + firstPrepared = await prepareIterationDrained(maxRounds <= 0); + } catch (e) { + if (abortSignal) abortSignal.removeEventListener("abort", linkAbort); + if (e instanceof LoopError) return jsonError(e.status, e.message); + throw e; + } + + const toolNsMap = new Map(); + const freeform = new Set(); + const toolSearch = new Set(); + for (const t of parsed.context.tools ?? []) { + if (t.namespace) toolNsMap.set(namespacedToolName(t.namespace, t.name), { namespace: t.namespace, name: t.name }); + if (t.freeform) freeform.add(t.name); + if (t.toolSearch) toolSearch.add(t.name); + } + + // Drive the remaining iterations live. Image generation runs interleaved with the real sidecar + // timing; the final answer's passthrough events come last. + async function* produce(): AsyncGenerator { + let prepared = firstPrepared; + try { + for (let i = 0; i < HARD_CAP; i++) { + const forceFinal = i >= maxRounds; + try { + // First loop turn reuses the eager HEADERS. Subsequent header acquisitions run here. + if (i > 0) { + yield { type: "heartbeat" }; + prepared = yield* prepareIterationEvents(forceFinal); + } + // Raw-byte progress heartbeats reach the bridge; semantic events remain buffered. + const split = yield* consumeIterationEvents(prepared); + + // Loop (fulfill + re-ask) ONLY when the model's actionable output is purely image_gen. A + // real tool call means this turn is terminal for Codex — finalize so those calls reach + // Codex. forceFinal also finalizes. + const shouldLoop = split.calls.length > 0 && !split.hasRealToolCall && !forceFinal; + if (!shouldLoop) { + yield* replay(split.passthrough); + return; + } + + // Fulfill each image call, then inject assistant + toolResult into messages. + const iterationThinking = extractIterationThinking(split.passthrough); + for (const [callIndex, call] of split.calls.entries()) { + yield { type: "heartbeat" }; + const result = await fulfillImageCall( + { id: call.id, name: call.name, arguments: call.args }, + plan, budget, signal, + ); + if (signal.aborted) throw new LoopError(499, "client closed request during image-bridge"); + const now = Date.now(); + let parsedArgs: Record = {}; + try { + const raw: unknown = JSON.parse(call.args || "{}"); + if (raw !== null && typeof raw === "object" && !Array.isArray(raw)) { + parsedArgs = raw as Record; + } + } catch { /* malformed args */ } + messages.push({ + role: "assistant", + content: [ + ...(callIndex === 0 && iterationThinking ? [iterationThinking] : []), + { type: "toolCall" as const, id: call.id, name: call.name, arguments: parsedArgs }, + ], + timestamp: now, + }); + messages.push({ + role: "toolResult", + toolCallId: call.id, + toolName: call.name, + content: JSON.stringify(result), + isError: !result.ok, + timestamp: now, + }); + } + } catch (e) { + yield { + type: "error", + message: e instanceof LoopError ? e.message : (e instanceof Error ? e.message : String(e)), + ...(e instanceof LoopError ? { status: e.status } : {}), + }; + return; + } + } + } finally { + if (abortSignal) abortSignal.removeEventListener("abort", linkAbort); + } + } + + const sse = bridgeToResponsesSSE( + produce(), parsed.modelId, toolNsMap, freeform, toolSearch, () => { + internalAbort.abort("client closed responses stream"); + }, undefined, + { + responseId: "", + hideThinkingSummary: parsed.options.hideThinkingSummary, + ...(deps.onFirstOutput ? { onFirstOutput: deps.onFirstOutput } : {}), + }, + ); + return new Response(sse, { headers: SSE_HEADERS }); +} diff --git a/src/images/plan.ts b/src/images/plan.ts new file mode 100644 index 0000000000..5bb27a4867 --- /dev/null +++ b/src/images/plan.ts @@ -0,0 +1,64 @@ +import type { OcxConfig, OcxParsedRequest, OcxProviderConfig } from "../types"; +import type { ImageBridgePlan } from "./types"; +import { getValidAccessToken } from "../oauth/index"; +import { resolveEnvValue } from "../config"; +import { getProviderRegistryEntry } from "../providers/registry"; +import { IMAGE_GEN_TOOL_NAME } from "./synthetic-tool"; + +const DEFAULT_MODEL = "grok-imagine-image-quality"; + +export function findXaiProvider(config: OcxConfig): { name: string; provider: OcxProviderConfig } | undefined { + // Primary: well-known name "xai" + const xai = config.providers["xai"]; + if (xai && xai.disabled !== true) return { name: "xai", provider: xai }; + // Fallback: hostname match for custom-named xAI configs + for (const [name, p] of Object.entries(config.providers)) { + if (p.disabled) continue; + try { + const host = new URL(p.baseUrl).hostname; + if (host === "api.x.ai" || host === "cli-chat-proxy.grok.com") return { name, provider: p }; + } catch { /* invalid baseUrl */ } + } + return undefined; +} + +export async function resolveXaiToken(providerName: string, provider: OcxProviderConfig): Promise { + const apiKey = resolveEnvValue(provider.apiKey)?.trim(); + if (apiKey) return apiKey; + // Built-in OAuth token only for the canonical "xai" provider — never for custom-named configs. + if (providerName !== "xai") return undefined; + try { + return await getValidAccessToken("xai"); + } catch { + return undefined; + } +} + +export async function planImageBridge( + config: OcxConfig, + parsed: OcxParsedRequest, + routedProvider: OcxProviderConfig, +): Promise { + if (config.images?.bridgeEnabled !== true) return undefined; + if (!parsed._imageGeneration) return undefined; + // Don't intercept for OpenAI native passthrough + const host = (() => { try { return new URL(routedProvider.baseUrl).hostname; } catch { return ""; } })(); + if (host === "api.openai.com") return undefined; + const found = findXaiProvider(config); + if (!found) return undefined; + const token = await resolveXaiToken(found.name, found.provider); + if (!token) return undefined; + // Pin the baseUrl to the registry entry, ignoring any config-level baseUrl override. + const registryEntry = getProviderRegistryEntry("xai"); + const pinnedBaseUrl = (registryEntry?.baseUrl ?? "https://api.x.ai/v1").replace(/\/+$/, ""); + // The synthetic tool injected into the conversation is named IMAGE_GEN_TOOL_NAME, + // which is what the model will actually call. Merge it with any original hosted tool names. + const toolNames = new Set(parsed._imageGeneration.toolNames); + toolNames.add(IMAGE_GEN_TOOL_NAME); + return { + provider: found.provider, + auth: { baseUrl: pinnedBaseUrl, token }, + model: config.images?.bridgeModel ?? DEFAULT_MODEL, + toolNames, + }; +} diff --git a/src/images/synthetic-tool.ts b/src/images/synthetic-tool.ts new file mode 100644 index 0000000000..a7707a0ff8 --- /dev/null +++ b/src/images/synthetic-tool.ts @@ -0,0 +1,70 @@ +import type { OcxTool } from "../types"; + +/** The function name the chat model sees + the name the loop intercepts. */ +export const IMAGE_GEN_TOOL_NAME = "image_gen"; + +const IMAGE_GEN_NAMES = new Set([ + "image_gen", "image_generation", "imagegen", + "generate_image", "generateimage", +]); + +export function isImageGenName(name: string): boolean { + return IMAGE_GEN_NAMES.has(name.toLowerCase()); +} + +/** + * Scan a Responses request's `tools[]` for hosted image-generation entries (`{type:"image_generation"}` or + * `{type:"image_gen"}`) and function entries whose name matches `isImageGenName`. Returns the set of all + * matched tool names plus the first matched raw tool object (so its config can be replayed), or undefined + * when image generation isn't enabled. + */ +export function extractHostedImageGeneration( + tools: unknown[] | undefined, +): { toolNames: Set; originalTool?: Record } | undefined { + if (!Array.isArray(tools)) return undefined; + const toolNames = new Set(); + let originalTool: Record | undefined; + for (const t of tools) { + if (!t || typeof t !== "object") continue; + const obj = t as Record; + if (obj.type === "image_generation" || obj.type === "image_gen") { + if (!originalTool) originalTool = obj; + const name = typeof obj.name === "string" ? obj.name : (obj.type as string); + toolNames.add(name); + } else if (obj.type === "function") { + // Responses API function tools have a flat shape: {type:"function", name:"...", parameters:{...}}. + // Also handle the nested Chat Completions shape {type:"function", function:{name:"..."}} for safety. + const fnName = + typeof obj.name === "string" ? obj.name : (obj as { function?: { name?: string } }).function?.name; + if (fnName && isImageGenName(fnName)) { + if (!originalTool) originalTool = obj; + toolNames.add(fnName); + } + } + } + if (toolNames.size === 0) return undefined; + return { toolNames, originalTool }; +} + +/** + * The synthetic function tool exposed to a chat/anthropic model in place of the dropped hosted + * image_generation. The model calls it like any function; the proxy intercepts the call and runs the + * real generation via the sidecar (the call is never relayed to Codex). `imageGeneration:true` flags it. + */ +export function buildImageTool(): OcxTool { + return { + name: IMAGE_GEN_TOOL_NAME, + description: + "Generate or edit an image. Returns absolute local filesystem path(s). " + + "Use when the user asks to create, draw, or edit 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 0000000000..b0d139aeff --- /dev/null +++ b/src/images/types.ts @@ -0,0 +1,19 @@ +import type { OcxProviderConfig } from "../types"; + +export interface ImageBridgePlan { + provider: OcxProviderConfig; + auth: { baseUrl: string; token: string }; + model: string; + toolNames: Set; +} + +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 0000000000..4bb3b67b0d --- /dev/null +++ b/src/images/xai-client.ts @@ -0,0 +1,133 @@ +/** + * xAI image generation/editing client. + * + * Calls xAI's `/images/generations` or `/images/edits` endpoint, composing a + * 60 s timeout with the caller's abort signal so the deadline covers the entire + * response body read. Non-2xx responses throw with the original status code — + * no 502 compression — so callers can distinguish rate-limit / auth failures + * from transient errors. + */ + +export interface XaiImageRequest { + prompt: string; + model?: string; // default "grok-imagine-image-quality" + n?: number; // 1-4 + size?: string; + quality?: string; + imageUrl?: string; // if set → /images/edits +} + +export interface XaiImageResult { + images: Array<{ b64_json?: string; url?: string }>; +} + +const XAI_IMAGES_TIMEOUT_MS = 60_000; +const XAI_DEFAULT_MODEL = "grok-imagine-image-quality"; + +// xAI only accepts aspect_ratio + resolution, not OpenAI's size/quality. Map +// OpenAI's "WxH" size to the closest standard ratio (log-space distance) and +// OpenAI quality buckets onto xAI's 1k/2k resolution. Unknown values are +// dropped rather than forwarded, to avoid sending parameters xAI rejects. +const XAI_ASPECT_RATIOS: ReadonlyArray = [ + ["1:1", 1], + ["3:4", 0.75], + ["4:3", 4 / 3], + ["9:16", 0.5625], + ["16:9", 16 / 9], +]; + +function mapSizeToAspectRatio(size?: string): string | undefined { + if (!size) return undefined; + const m = /^(\d+)x(\d+)$/.exec(size); + if (!m) return undefined; + const ratio = parseInt(m[1], 10) / parseInt(m[2], 10); + let best = XAI_ASPECT_RATIOS[0]!; + let bestDiff = Infinity; + for (const [label, r] of XAI_ASPECT_RATIOS) { + const diff = Math.abs(Math.log(ratio / r)); + if (diff < bestDiff) { + bestDiff = diff; + best = [label, r] as const; + } + } + return best[0]; +} + +function mapQualityToResolution(quality?: string): string | undefined { + if (!quality) return undefined; + const q = quality.toLowerCase(); + if (q === "hd" || q === "high") return "2k"; + if (q === "standard" || q === "low" || q === "auto") return "1k"; + return undefined; +} + +export async function callXaiImages( + req: XaiImageRequest, + auth: { baseUrl: string; token: string }, + signal?: AbortSignal, +): Promise { + const isEdit = typeof req.imageUrl === "string" && req.imageUrl.length > 0; + const endpoint = isEdit ? "/images/edits" : "/images/generations"; + + const body: Record = { + model: req.model ?? XAI_DEFAULT_MODEL, + prompt: req.prompt, + n: req.n ?? 1, + }; + const aspectRatio = mapSizeToAspectRatio(req.size); + const resolution = mapQualityToResolution(req.quality); + if (aspectRatio) body.aspect_ratio = aspectRatio; + if (resolution) body.resolution = resolution; + if (isEdit) { + body.image = { url: req.imageUrl as string, type: "image_url" }; + } + + const timeout = AbortSignal.timeout(XAI_IMAGES_TIMEOUT_MS); + const linkedSignal = signal ? AbortSignal.any([signal, timeout]) : timeout; + + const resp = await fetch(`${auth.baseUrl}${endpoint}`, { + method: "POST", + headers: { + "Authorization": `Bearer ${auth.token}`, + "Content-Type": "application/json", + }, + body: JSON.stringify(body), + signal: linkedSignal, + }); + + if (!resp.ok) { + throw new Error("xAI images API returned " + resp.status); + } + + // Read the body as text under the linked signal, then parse. The 60 s timeout + // and caller abort cover the read. A 200 MiB hard cap on the text prevents a + // runaway response from exhausting memory before materialization caps apply. + const MAX_RESPONSE_BYTES = 200 * 1024 * 1024; + const reader = resp.body?.getReader(); + if (!reader) throw new Error("xAI images API returned no body"); + const decoder = new TextDecoder(); + let text = ""; + let totalBytes = 0; + try { + for (;;) { + const { value, done } = await reader.read(); + if (done) break; + totalBytes += value.byteLength; + if (totalBytes > MAX_RESPONSE_BYTES) throw new Error("xAI images API response exceeds size cap"); + text += decoder.decode(value, { stream: true }); + } + text += decoder.decode(); + } finally { + try { await reader.cancel(); } catch { /* ignore cancel errors */ } + reader.releaseLock(); + } + + const json = JSON.parse(text) as { data?: Array<{ b64_json?: string; url?: string }> }; + + const images = (json.data ?? []).map((entry) => ({ + b64_json: entry.b64_json, + url: entry.url, + })); + + return { images }; +} diff --git a/src/lib/destination-policy.ts b/src/lib/destination-policy.ts index 54cffcbd19..1ec932bd4c 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" @@ -178,3 +178,52 @@ export async function providerDestinationResolvedError( } return null; } + +export interface UrlDestinationAssessment { + kind: DestinationKind; + detail: string; +} + +/** + * Synchronous literal URL destination assessment — classifies the hostname + * without DNS resolution. Returns null for unparseable URLs. + */ +export function assessUrlDestination(url: string): UrlDestinationAssessment | null { + return assessDestination(url); +} + +/** + * Async DNS-resolved URL safety check. Resolves A/AAAA records and rejects + * if any address is loopback, private, link-local, unspecified, or metadata. + * Throws on unsafe destination; returns void on safe/public destination. + * DNS resolution failures are treated as unsafe (fail-closed). + */ +export async function assertUrlResolvesPublic(url: string): Promise { + let hostname: string; + try { + hostname = normalizeHostname(new URL(url.trim()).hostname); + } catch { + throw new Error("image URL is not a valid URL"); + } + if (!hostname) throw new Error("image URL has no hostname"); + const literalAssessment = assessDestination(url); + if (literalAssessment && literalAssessment.kind !== "public" && literalAssessment.kind !== "hostname") { + throw new Error(`image URL targets ${literalAssessment.detail}`); + } + // For literal IPs and localhost, the sync path already classified them. + if (isIP(hostname) !== 0 || hostname === "localhost" || hostname.endsWith(".localhost")) return; + let addresses: { address: string }[]; + try { + addresses = await lookup(hostname, { all: true, verbatim: true }); + } catch { + // If DNS fails, we can't verify — fail-closed (unlike provider config-time validation, + // this is a runtime fetch to an untrusted URL, so be conservative). + throw new Error(`image URL hostname ${hostname} could not be resolved`); + } + for (const { address } of addresses) { + const ipKind = isIP(address); + const assessment = ipKind === 4 ? classifyIpv4(address) : ipKind === 6 ? classifyIpv6(normalizeHostname(address)) : null; + if (!assessment || assessment.kind === "public") continue; + throw new Error(`image URL hostname ${hostname} resolves to ${assessment.detail} (${address})`); + } +} diff --git a/src/responses/parser.ts b/src/responses/parser.ts index 3c13b477e1..c0a311c1ce 100644 --- a/src/responses/parser.ts +++ b/src/responses/parser.ts @@ -16,6 +16,7 @@ import { compactionItemToText } from "./compaction"; import { previousResponseReplayPrefixLength } from "./state"; import { decodeReasoningEnvelope } from "./reasoning-envelope"; import { extractHostedWebSearch, WEB_SEARCH_TOOL_NAME } from "../web-search/synthetic-tool"; +import { extractHostedImageGeneration } from "../images/synthetic-tool"; function isObj(v: unknown): v is Record { return typeof v === "object" && v !== null && !Array.isArray(v); @@ -612,6 +613,10 @@ export function parseRequest(body: unknown): OcxParsedRequest { // gpt-mini sidecar for routed providers. buildTools still drops the hosted tool; the sidecar path // re-injects a synthetic function tool only when it will actually handle the call. const webSearch = extractHostedWebSearch(data.tools as unknown[] | undefined); + const imageGen = extractHostedImageGeneration([ + ...(data.tools as unknown[] ?? []), + ...loadedToolSpecs, + ]); // Detect structured-output mode (Responses `text.format`) so the web-search sidecar can render its // tool_result as JSON rather than prose that could corrupt the model's schema-constrained answer. const structuredOutput = detectStructuredOutput(data.text); @@ -625,6 +630,7 @@ export function parseRequest(body: unknown): OcxParsedRequest { _rawBody: body, ...(replayedInputPrefixLength > 0 ? { _replayPrefixLen: replayedInputPrefixLength } : {}), ...(webSearch ? { _webSearch: webSearch } : {}), + ...(imageGen ? { _imageGeneration: imageGen } : {}), ...(structuredOutput ? { _structuredOutput: true } : {}), ...(compactionRequest ? { _compactionRequest: true } : {}), ...(contextCompactionBoundary ? { _contextCompactionBoundary: true } : {}), diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index fbed025540..104eb833b1 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -39,6 +39,7 @@ import { UnsupportedOAuthProviderError, } from "../../oauth"; import { buildWebSearchTool, planWebSearch, runWithWebSearch, shouldResolveOpenAiWebSearchSidecar } from "../../web-search"; +import { buildImageTool, planImageBridge, runWithImageBridge } from "../../images"; import { describeImagesInPlace, planVisionSidecar, shouldResolveOpenAiVisionSidecar, stripImagesInPlace } from "../../vision"; import { createAdapterEventQueue, preflightAdapterEvents } from "../../adapters/run-turn-queue"; import { @@ -1526,6 +1527,36 @@ export async function handleResponses( }); } + // Image bridge: check BEFORE the runTurn early-return so runTurn adapters (e.g. Cursor) + // also route through the bridge when image_generation is requested. The bridge loop + // internally supports both standard and runTurn adapter paths. + const imgPlan = await planImageBridge(config, parsed, route.provider); + if (imgPlan) { + // The bridge forces stream:true internally and returns SSE. Non-streaming requests can't be + // served — reject explicitly rather than returning SSE to a client expecting JSON. + if (!parsed.stream) { + return formatErrorResponse(400, "invalid_request_error", "image bridge requires stream=true"); + } + parsed.context.tools = [...(parsed.context.tools ?? []), buildImageTool()]; + const imgResponse = await runWithImageBridge({ + parsed, adapter, + plan: imgPlan, + forwardHeaders: selectedForwardHeaders, + onAttemptSend: () => noteAttemptSend(logCtx.activeAttempt, logCtx.usageLogInputTokens), + abortSignal: options.abortSignal, + ...(config.images?.maxRounds != null ? { maxRounds: config.images.maxRounds } : {}), + ...(options.onFirstOutput ? { onFirstOutput: options.onFirstOutput } : {}), + }); + if (imgResponse.body) { + const imgTurnAc = new AbortController(); + return new Response(trackStreamLifetime(imgResponse.body, imgTurnAc), { + status: imgResponse.status, + headers: imgResponse.headers, + }); + } + return imgResponse; + } + if (adapter.runTurn) { const runTurnAbort = new AbortController(); linkAbortSignal(runTurnAbort, options.abortSignal); diff --git a/src/types.ts b/src/types.ts index fb0e244d3f..9b9afff3a1 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; } /** @@ -741,6 +745,12 @@ export interface OcxImagesConfig { provider?: string; /** Upstream timeout (ms) for one /v1/images relay. Default 300000 — generation is slow. */ timeoutMs?: number; + /** Master switch for the image bridge. Default false — set true to enable paid xAI Grok Imagine generation. */ + bridgeEnabled?: boolean; + /** xAI image model id. Default "grok-imagine-image-quality" (see DEFAULT_MODEL in images/plan.ts). */ + bridgeModel?: string; + /** Max image-generation loop iterations before forced-final. Default 3 (see DEFAULT_MAX_ROUNDS in images/loop.ts). */ + maxRounds?: number; } export interface OcxSearchConfig { diff --git a/tests/images/artifacts-ssrf.test.ts b/tests/images/artifacts-ssrf.test.ts new file mode 100644 index 0000000000..37f36942ee --- /dev/null +++ b/tests/images/artifacts-ssrf.test.ts @@ -0,0 +1,51 @@ +import { describe, expect, test } from "bun:test"; +import { assessUrlDestination, assertUrlResolvesPublic } from "../../src/lib/destination-policy"; + +describe("SSRF: assessUrlDestination", () => { + test("loopback IPv4 → loopback", () => { + expect(assessUrlDestination("http://127.0.0.1/test")?.kind).toBe("loopback"); + }); + test("link-local → link-local", () => { + expect(assessUrlDestination("http://169.254.1.1/latest")?.kind).toBe("link-local"); + }); + test("private 10.x → private", () => { + expect(assessUrlDestination("http://10.0.0.1/test")?.kind).toBe("private"); + }); + test("private 192.168 → private", () => { + expect(assessUrlDestination("http://192.168.1.1/test")?.kind).toBe("private"); + }); + test("private 172.16 → private", () => { + expect(assessUrlDestination("http://172.16.0.1/test")?.kind).toBe("private"); + }); + test("metadata endpoint → metadata", () => { + expect(assessUrlDestination("http://169.254.170.2/test")?.kind).toBe("metadata"); + }); + test("localhost → localhost", () => { + expect(assessUrlDestination("http://localhost/test")?.kind).toBe("localhost"); + }); + test("public HTTPS → hostname or public", () => { + const kind = assessUrlDestination("https://example.com/image.png")?.kind; + expect(kind === "hostname" || kind === "public").toBe(true); + }); + test("public IP → public", () => { + expect(assessUrlDestination("https://8.8.8.8/image.png")?.kind).toBe("public"); + }); + test("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(); + }); +}); diff --git a/tests/images/fulfill.test.ts b/tests/images/fulfill.test.ts new file mode 100644 index 0000000000..67e06b62b1 --- /dev/null +++ b/tests/images/fulfill.test.ts @@ -0,0 +1,146 @@ +import { afterAll, beforeAll, describe, expect, mock, test } from "bun:test"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { randomUUID } from "node:crypto"; +import type { ImageBridgePlan } from "../../src/images/types"; +import type { XaiImageRequest } from "../../src/images/xai-client"; + +const PREV_HOME = process.env.OPENCODEX_HOME; +beforeAll(() => { process.env.OPENCODEX_HOME = join(tmpdir(), "ocx-test-" + randomUUID()); }); +afterAll(() => { if (PREV_HOME === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = PREV_HOME; }); + +// --- Mutable mock state (reset() restores defaults before each test) --- +let xaiResult: { images: Array<{ b64_json?: string; url?: string }> } = { images: [{ b64_json: "dGVzdA==" }] }; +let xaiError: Error | null = null; +const xaiCalls: XaiImageRequest[] = []; +let matIdx = 0; +let dlIdx = 0; +let materializeFn: (i: number) => Promise = async (i) => `/test/img-${i}.png`; +let downloadFn: (i: number) => Promise = async (i) => `/test/dl-${i}.png`; + +mock.module("../../src/images/xai-client", () => ({ + callXaiImages: async (req: XaiImageRequest) => { xaiCalls.push(req); if (xaiError) throw xaiError; return xaiResult; }, +})); +mock.module("../../src/images/artifacts", () => ({ + createImageBudget: () => ({ spent: 0 }), + materializeInlineImage: async () => materializeFn(matIdx++), + downloadImageToArtifact: async () => downloadFn(dlIdx++), +})); + +const { fulfillImageCall } = await import("../../src/images/fulfill"); + +const plan = { + provider: {} as never, + auth: { baseUrl: "https://api.x.ai", token: "test-token" }, + model: "grok-imagine-image-quality", + toolNames: new Set(["image_gen"]), +} as ImageBridgePlan; + +function reset(): void { + xaiResult = { images: [{ b64_json: "dGVzdA==" }] }; + xaiError = null; + xaiCalls.length = 0; + matIdx = 0; + dlIdx = 0; + materializeFn = async (i) => `/test/img-${i}.png`; + downloadFn = async (i) => `/test/dl-${i}.png`; +} + +describe("fulfillImageCall", () => { + test("valid args → ok:true with file", async () => { + reset(); + const r = await fulfillImageCall( + { id: "c1", name: "image_gen", arguments: JSON.stringify({ prompt: "a cat", n: 2 }) }, + plan, { spent: 0 }, + ); + expect(r.ok).toBe(true); + expect(r.files.length).toBe(1); + }); + + test("missing prompt → ok:false 'missing prompt'", async () => { + reset(); + const r = await fulfillImageCall({ id: "c1", name: "image_gen", arguments: "{}" }, plan, { spent: 0 }); + expect(r.ok).toBe(false); + expect(r.error).toBe("missing prompt"); + }); + + test("invalid JSON args → ok:false 'invalid arguments JSON'", async () => { + reset(); + const r = await fulfillImageCall({ id: "c1", name: "image_gen", arguments: "{bad" }, plan, { spent: 0 }); + expect(r.ok).toBe(false); + expect(r.error).toBe("invalid arguments JSON"); + }); + + test("xAI throws → ok:false with error message", async () => { + reset(); + xaiError = new Error("xAI images API returned 500"); + const r = await fulfillImageCall( + { id: "c1", name: "image_gen", arguments: JSON.stringify({ prompt: "x" }) }, plan, { spent: 0 }, + ); + expect(r.ok).toBe(false); + expect(r.error).toContain("500"); + }); + + test("b64_json result → materialized via materializeInlineImage", async () => { + reset(); + xaiResult = { images: [{ b64_json: "dGVzdA==" }] }; + await fulfillImageCall({ id: "c1", name: "image_gen", arguments: `{"prompt":"x"}` }, plan, { spent: 0 }); + expect(matIdx).toBe(1); + expect(dlIdx).toBe(0); + }); + + test("URL result → materialized via downloadImageToArtifact", async () => { + reset(); + xaiResult = { images: [{ url: "https://cdn.example.com/i.png" }] }; + await fulfillImageCall({ id: "c1", name: "image_gen", arguments: `{"prompt":"x"}` }, plan, { spent: 0 }); + expect(dlIdx).toBe(1); + expect(matIdx).toBe(0); + }); + + test("all images fail → ok:false", async () => { + reset(); + materializeFn = async () => { throw new Error("disk full"); }; + const r = await fulfillImageCall({ id: "c1", name: "image_gen", arguments: `{"prompt":"x"}` }, plan, { spent: 0 }); + expect(r.ok).toBe(false); + expect(r.error).toContain("no usable images"); + }); + + test("one of two images fails → ok:true with 1 file", async () => { + reset(); + xaiResult = { images: [{ b64_json: "AAA=" }, { b64_json: "QkI=" }] }; + materializeFn = async (i) => { if (i === 1) throw new Error("partial fail"); return `/test/img-${i}.png`; }; + const r = await fulfillImageCall({ id: "c1", name: "image_gen", arguments: `{"prompt":"x"}` }, plan, { spent: 0 }); + expect(r.ok).toBe(true); + expect(r.files.length).toBe(1); + }); + + test("forwards prompt, model, and n to callXaiImages", async () => { + reset(); + await fulfillImageCall( + { id: "c1", name: "image_gen", arguments: JSON.stringify({ prompt: "a cat", n: 2 }) }, + plan, { spent: 0 }, + ); + expect(xaiCalls.length).toBe(1); + expect(xaiCalls[0]!.prompt).toBe("a cat"); + expect(xaiCalls[0]!.model).toBe(plan.model); + expect(xaiCalls[0]!.n).toBe(2); + }); + + test("clamps n > 4 down to 4", async () => { + reset(); + await fulfillImageCall( + { id: "c1", name: "image_gen", arguments: JSON.stringify({ prompt: "x", n: 100 }) }, + plan, { spent: 0 }, + ); + expect(xaiCalls[0]!.n).toBe(4); + }); + + test("forwards imageUrl from image_url arg", async () => { + reset(); + await fulfillImageCall( + { id: "c1", name: "image_gen", arguments: JSON.stringify({ prompt: "x", image_url: "https://example.com/i.png" }) }, + plan, { spent: 0 }, + ); + expect(xaiCalls[0]!.imageUrl).toBe("https://example.com/i.png"); + }); +}); diff --git a/tests/images/loop.test.ts b/tests/images/loop.test.ts new file mode 100644 index 0000000000..82d8e7b564 --- /dev/null +++ b/tests/images/loop.test.ts @@ -0,0 +1,183 @@ +import { afterAll, beforeAll, describe, expect, mock, test } from "bun:test"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { randomUUID } from "node:crypto"; +import type { ProviderAdapter, IncomingMeta } from "../../src/adapters/base"; +import type { AdapterEvent, OcxParsedRequest } from "../../src/types"; +import type { ImageBridgePlan, ImageCallResult } from "../../src/images/types"; + +const PREV_HOME = process.env.OPENCODEX_HOME; +beforeAll(() => { process.env.OPENCODEX_HOME = join(tmpdir(), "ocx-test-" + randomUUID()); }); +afterAll(() => { if (PREV_HOME === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = PREV_HOME; }); + +// --- Mock parseStreamWithProgress: simplify to direct delegation --- +mock.module("../../src/web-search/progress-stream", () => ({ + parseStreamWithProgress: async function* (_resp: Response, parse: (r: Response) => AsyncGenerator, _opts: unknown) { + for await (const e of parse(_resp)) yield e; + }, + RoutedModelInactivityError: class extends Error { readonly timeoutMs = 0; }, + WebSearchStreamProtocolError: class extends Error { /* */ }, +})); + +// --- Mock fulfillImageCall --- +let fulfillResult: ImageCallResult = { + ok: true, model: "grok-imagine-image-quality", prompt: "a cat", + files: ["/test/img.png"], count: 1, markdown: "![image](/test/img.png)", +}; +mock.module("../../src/images/fulfill", () => ({ + fulfillImageCall: async (): Promise => fulfillResult, +})); + +const { runWithImageBridge } = await import("../../src/images/loop"); + +// --- Mock adapter: yields canned events per iteration from a queue --- +let streamQueue: AdapterEvent[][] = []; +let buildRequestCalls = 0; +const mockAdapter: ProviderAdapter = { + name: "test", + buildRequest: async () => { buildRequestCalls++; return { url: "https://test/v1/chat", method: "POST", headers: {}, body: "{}" }; }, + fetchResponse: async () => new Response("{}", { status: 200, headers: { "content-type": "application/json" } }), + parseStream: async function* (): AsyncGenerator { + const events = streamQueue.shift(); + if (events) for (const e of events) yield e; + }, +}; + +const plan = { + provider: {} as never, + auth: { baseUrl: "https://api.x.ai", token: "test-token" }, + model: "grok-imagine-image-quality", + toolNames: new Set(["image_gen"]), +} as ImageBridgePlan; + +function makeParsed(): OcxParsedRequest { + return { modelId: "test-model", context: { messages: [], tools: [] }, stream: true, options: {} } as OcxParsedRequest; +} + +const imageCallEvents: AdapterEvent[] = [ + { type: "tool_call_start", id: "call_1", name: "image_gen" }, + { type: "tool_call_delta", arguments: '{"prompt":"a cat"}' }, + { type: "tool_call_end" }, + { type: "done" }, +]; + +async function runAndGetSSE(streams: AdapterEvent[][], fulfill?: ImageCallResult): Promise { + streamQueue = streams.map(s => [...s]); + if (fulfill) fulfillResult = fulfill; + const response = await runWithImageBridge({ parsed: makeParsed(), adapter: mockAdapter, plan }); + return await response.text(); +} + +describe("runWithImageBridge", () => { + test("no image tool call → passthrough text + done", async () => { + const sse = await runAndGetSSE([ + [{ type: "text_delta", text: "hello world" }, { type: "done" }], + ]); + expect(sse).toContain("hello world"); + }); + + test("single image call → fulfilled, second iteration yields text", async () => { + const sse = await runAndGetSSE( + [imageCallEvents, [{ type: "text_delta", text: "Here is your image" }, { type: "done" }]], + { ok: true, model: "grok-imagine-image-quality", prompt: "a cat", files: ["/test/img.png"], count: 1, markdown: "![image](/test/img.png)" }, + ); + expect(sse).toContain("Here is your image"); + }); + + test("fulfillImageCall error → model responds about failure", async () => { + const sse = await runAndGetSSE( + [imageCallEvents, [{ type: "text_delta", text: "Sorry, image generation failed" }, { type: "done" }]], + { ok: false, model: "grok-imagine-image-quality", prompt: "a cat", files: [], count: 0, error: "xAI unreachable" }, + ); + expect(sse).toContain("Sorry, image generation failed"); + }); + + test("image_gen tool call is intercepted — not visible in client SSE", async () => { + const sse = await runAndGetSSE( + [imageCallEvents, [{ type: "text_delta", text: "done" }, { type: "done" }]], + ); + // The tool_call_start event for image_gen should NOT appear in client-facing SSE + expect(sse).not.toContain("image_gen"); + expect(sse).not.toContain("tool_call_start"); + }); + + test("maxRounds: 1 bounds upstream requests and forces final after limit", async () => { + buildRequestCalls = 0; + // Round 0: model calls image_gen (within limit → fulfill + loop) + // Round 1: forced-final pass (forceFinal, image tools stripped from request) + streamQueue = [ + [...imageCallEvents], + [{ type: "text_delta" as const, text: "final answer" }, { type: "done" as const }], + ]; + const response = await runWithImageBridge({ parsed: makeParsed(), adapter: mockAdapter, plan, maxRounds: 1 }); + const sse = await response.text(); + expect(sse).toContain("final answer"); + // Exactly 2 upstream requests: round 0 (image call) + round 1 (forced final) + expect(buildRequestCalls).toBe(2); + }); + + test("maxRounds: 0 forces final immediately — no image tool offered", async () => { + buildRequestCalls = 0; + streamQueue = [ + [{ type: "text_delta" as const, text: "direct answer" }, { type: "done" as const }], + ]; + const response = await runWithImageBridge({ parsed: makeParsed(), adapter: mockAdapter, plan, maxRounds: 0 }); + const sse = await response.text(); + expect(sse).toContain("direct answer"); + // Single upstream request — first iteration is already forced-final + expect(buildRequestCalls).toBe(1); + }); +}); + +// --------------------------------------------------------------------------- +// runTurn adapter path (Cursor) — events arrive via an emit callback, not +// buildRequest/fetchResponse/parseStream. +// --------------------------------------------------------------------------- + +describe("runWithImageBridge — runTurn adapter", () => { + let runTurnEventQueue: AdapterEvent[][] = []; + const runTurnAdapter: ProviderAdapter = { + ...mockAdapter, + runTurn: async (_parsed: OcxParsedRequest, _incoming: IncomingMeta, emit: (e: AdapterEvent) => void) => { + const events = runTurnEventQueue.shift(); + if (events) for (const e of events) emit(e); + }, + }; + + test("runTurn adapter → image call intercepted and fulfilled", async () => { + runTurnEventQueue = [ + [...imageCallEvents], + [{ type: "text_delta", text: "Here is your image" }, { type: "done" }], + ]; + fulfillResult = { + ok: true, model: "grok-imagine-image-quality", prompt: "a cat", + files: ["/test/img.png"], count: 1, markdown: "![image](/test/img.png)", + }; + const response = await runWithImageBridge({ + parsed: makeParsed(), adapter: runTurnAdapter, plan, maxRounds: 1, + }); + const sse = await response.text(); + expect(sse).toContain("Here is your image"); + // The synthetic image_gen tool call must NOT leak to the client + expect(sse).not.toContain("image_gen"); + expect(sse).not.toContain("tool_call_start"); + }); + + test("runTurn adapter → text passthrough (no image call)", async () => { + runTurnEventQueue = [ + [{ type: "text_delta", text: "hello from runTurn" }, { type: "done" }], + ]; + const response = await runWithImageBridge({ parsed: makeParsed(), adapter: runTurnAdapter, plan }); + const sse = await response.text(); + expect(sse).toContain("hello from runTurn"); + }); + + test("runTurn adapter → error event surfaces as upstream failure", async () => { + runTurnEventQueue = [ + [{ type: "error", message: "cursor blew up" }], + ]; + const response = await runWithImageBridge({ parsed: makeParsed(), adapter: runTurnAdapter, plan }); + const sse = await response.text(); + expect(sse).toContain("cursor blew up"); + }); +}); diff --git a/tests/images/plan.test.ts b/tests/images/plan.test.ts new file mode 100644 index 0000000000..85a536ee6f --- /dev/null +++ b/tests/images/plan.test.ts @@ -0,0 +1,143 @@ +import { afterAll, beforeAll, describe, expect, mock, test } from "bun:test"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { randomUUID } from "node:crypto"; +import type { OcxConfig, OcxProviderConfig, OcxParsedRequest } from "../../src/types"; + +const PREV_HOME = process.env.OPENCODEX_HOME; +beforeAll(() => { process.env.OPENCODEX_HOME = join(tmpdir(), "ocx-test-" + randomUUID()); }); +afterAll(() => { if (PREV_HOME === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = PREV_HOME; }); + +/** Mutable token that the mocked getValidAccessToken resolves to. */ +let tokenResult: string | null = null; +mock.module("../../src/oauth/index", () => ({ + getValidAccessToken: async () => tokenResult, +})); + +const { planImageBridge } = await import("../../src/images/plan"); + +function makeConfig( + providers: Record>, + images?: { bridgeEnabled?: boolean; bridgeModel?: string }, +): OcxConfig { + return { + port: 0, + defaultProvider: "test", + providers: Object.fromEntries( + Object.entries(providers).map(([k, v]) => [k, { adapter: "openai", baseUrl: "https://api.test.com", ...v }]), + ), + ...(images ? { images } : {}), + } as OcxConfig; +} + +function makeParsed(withImageGen: boolean): OcxParsedRequest { + return { + modelId: "test-model", + context: { messages: [], tools: [] }, + stream: true, + options: {}, + ...(withImageGen ? { _imageGeneration: { toolNames: new Set(["image_gen"]) } } : {}), + } as OcxParsedRequest; +} + +const routed = { adapter: "openai", baseUrl: "https://api.anthropic.com" } as OcxProviderConfig; +const openaiRouted = { adapter: "openai", baseUrl: "https://api.openai.com" } as OcxProviderConfig; + +describe("planImageBridge", () => { + test("bridgeEnabled false → undefined", async () => { + expect(await planImageBridge(makeConfig({ test: routed }, { bridgeEnabled: false }), makeParsed(true), routed)).toBeUndefined(); + }); + + test("bridgeEnabled not set → undefined (opt-in required)", async () => { + // xAI provider configured but images.bridgeEnabled is absent — must not bridge. + const cfg = makeConfig({ xai: { baseUrl: "https://api.x.ai", apiKey: "test-token" } }); + expect(await planImageBridge(cfg, makeParsed(true), routed)).toBeUndefined(); + }); + + test("_imageGeneration not set → undefined", async () => { + expect(await planImageBridge(makeConfig({ test: routed }, { bridgeEnabled: true }), makeParsed(false), routed)).toBeUndefined(); + }); + + test("routedProvider is api.openai.com → undefined", async () => { + const cfg = makeConfig({ xai: { baseUrl: "https://api.x.ai", apiKey: "test-token" } }, { bridgeEnabled: true }); + expect(await planImageBridge(cfg, makeParsed(true), openaiRouted)).toBeUndefined(); + }); + + test("no xAI provider → undefined", async () => { + expect(await planImageBridge(makeConfig({ test: routed }, { bridgeEnabled: true }), makeParsed(true), routed)).toBeUndefined(); + }); + + test("xAI provider but apiKey empty and no OAuth → undefined", async () => { + tokenResult = null; + const cfg = makeConfig({ xai: { baseUrl: "https://api.x.ai", apiKey: "" } }, { bridgeEnabled: true }); + expect(await planImageBridge(cfg, makeParsed(true), routed)).toBeUndefined(); + }); + + test("xAI provider with API key → returns plan with correct model", async () => { + const cfg = makeConfig({ xai: { baseUrl: "https://api.x.ai", apiKey: "test-token" } }, { bridgeEnabled: true }); + const plan = await planImageBridge(cfg, makeParsed(true), routed); + expect(plan).toBeDefined(); + expect(plan!.model).toBe("grok-imagine-image-quality"); + expect(plan!.auth.token).toBe("test-token"); + expect(plan!.auth.baseUrl).toBe("https://api.x.ai/v1"); + }); + + test("xAI provider with OAuth (getValidAccessToken) → returns plan", async () => { + tokenResult = "fake-oauth-123"; + const cfg = makeConfig({ xai: { baseUrl: "https://api.x.ai" } }, { bridgeEnabled: true }); + const plan = await planImageBridge(cfg, makeParsed(true), routed); + expect(plan).toBeDefined(); + expect(plan!.auth.token).toBe("fake-oauth-123"); + tokenResult = null; + }); + + test("custom-named provider with api.x.ai baseUrl → found via fallback", async () => { + const cfg = makeConfig({ mygrok: { baseUrl: "https://api.x.ai", apiKey: "test-token" } }, { bridgeEnabled: true }); + const plan = await planImageBridge(cfg, makeParsed(true), routed); + expect(plan).toBeDefined(); + expect(plan!.provider).toBe(cfg.providers.mygrok); + }); + + test("custom bridgeModel is honored", async () => { + const cfg = makeConfig( + { xai: { baseUrl: "https://api.x.ai", apiKey: "test-token" } }, + { bridgeEnabled: true, bridgeModel: "custom-img-model" }, + ); + expect((await planImageBridge(cfg, makeParsed(true), routed))!.model).toBe("custom-img-model"); + }); + + test("toolNames includes IMAGE_GEN_TOOL_NAME so the loop can intercept synthetic calls", async () => { + const { IMAGE_GEN_TOOL_NAME } = await import("../../src/images/synthetic-tool"); + const cfg = makeConfig({ xai: { baseUrl: "https://api.x.ai", apiKey: "test-token" } }, { bridgeEnabled: true }); + const plan = await planImageBridge(cfg, makeParsed(true), routed); + expect(plan).toBeDefined(); + // The plan always merges in IMAGE_GEN_TOOL_NAME, even if _imageGeneration.toolNames + // only contained the original hosted tool name. + expect(plan!.toolNames.has(IMAGE_GEN_TOOL_NAME)).toBe(true); + }); + + test("baseUrl is pinned to registry regardless of config override", async () => { + // config 里 xai provider 的 baseUrl 被改成恶意 host + const cfg = makeConfig( + { xai: { adapter: "openai-chat", baseUrl: "https://evil.example.com/v1", apiKey: "test-key" } }, + { bridgeEnabled: true }, + ); + const plan = await planImageBridge(cfg, makeParsed(true), routed); + expect(plan).toBeDefined(); + // auth.baseUrl 必须是 registry pin 的地址，不是 config 里的恶意地址 + expect(plan!.auth.baseUrl).toBe("https://api.x.ai/v1"); + }); + + test("custom-named provider with api.x.ai baseUrl does NOT get built-in OAuth token", async () => { + tokenResult = "should-not-be-used"; + // provider 名为 "my-xai"，baseUrl 指向 api.x.ai，没有 apiKey + const cfg = makeConfig( + { "my-xai": { adapter: "openai-chat", baseUrl: "https://api.x.ai/v1" } }, + { bridgeEnabled: true }, + ); + const plan = await planImageBridge(cfg, makeParsed(true), routed); + // 没有 apiKey 也没有 "xai" 的 OAuth → 没有 token → 没有 plan + expect(plan).toBeUndefined(); + tokenResult = null; + }); +}); diff --git a/tests/images/synthetic-tool.test.ts b/tests/images/synthetic-tool.test.ts new file mode 100644 index 0000000000..7572c94c14 --- /dev/null +++ b/tests/images/synthetic-tool.test.ts @@ -0,0 +1,62 @@ +import { describe, expect, test } from "bun:test"; +import { isImageGenName, extractHostedImageGeneration, buildImageTool } from "../../src/images/synthetic-tool"; + +describe("isImageGenName", () => { + test("'image_gen' → true", () => { + expect(isImageGenName("image_gen")).toBe(true); + }); + + test("'IMAGE_GENERATION' → true (case insensitive)", () => { + expect(isImageGenName("IMAGE_GENERATION")).toBe(true); + }); + + test("'imagegen' → true", () => { + expect(isImageGenName("imagegen")).toBe(true); + }); + + test("'not_image' → false", () => { + expect(isImageGenName("not_image")).toBe(false); + }); +}); + +describe("extractHostedImageGeneration", () => { + test("type 'image_generation' → returns toolNames with that name", () => { + const result = extractHostedImageGeneration([{ type: "image_generation" }]); + expect(result).toBeDefined(); + expect(result!.toolNames.has("image_generation")).toBe(true); + }); + + test("flat Responses function tool with 'image_gen' name → returns with that name", () => { + const result = extractHostedImageGeneration([ + { type: "function", name: "image_gen", parameters: { type: "object" } }, + ]); + expect(result).toBeDefined(); + expect(result!.toolNames.has("image_gen")).toBe(true); + }); + + test("nested Chat Completions function tool with 'image_gen' name → returns with that name", () => { + const result = extractHostedImageGeneration([ + { type: "function", function: { name: "image_gen" } }, + ]); + expect(result).toBeDefined(); + expect(result!.toolNames.has("image_gen")).toBe(true); + }); + + test("no matching tools → undefined", () => { + expect( + extractHostedImageGeneration([{ type: "function", function: { name: "shell" } }]), + ).toBeUndefined(); + }); + + test("undefined → undefined", () => { + expect(extractHostedImageGeneration(undefined)).toBeUndefined(); + }); +}); + +describe("buildImageTool", () => { + test("has name 'image_gen' and imageGeneration flag", () => { + const tool = buildImageTool(); + expect(tool.name).toBe("image_gen"); + expect(tool.imageGeneration).toBe(true); + }); +}); diff --git a/tests/images/xai-client.test.ts b/tests/images/xai-client.test.ts new file mode 100644 index 0000000000..6c8a6cdc70 --- /dev/null +++ b/tests/images/xai-client.test.ts @@ -0,0 +1,114 @@ +import { afterAll, afterEach, beforeAll, describe, expect, test } from "bun:test"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { randomUUID } from "node:crypto"; +import { callXaiImages } from "../../src/images/xai-client"; + +const PREV_HOME = process.env.OPENCODEX_HOME; +beforeAll(() => { process.env.OPENCODEX_HOME = join(tmpdir(), "ocx-test-" + randomUUID()); }); +afterAll(() => { if (PREV_HOME === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = PREV_HOME; }); + +const AUTH = { baseUrl: "https://api.x.ai", token: "test-token" }; +const originalFetch = globalThis.fetch; +afterEach(() => { globalThis.fetch = originalFetch; }); + +/** Replace globalThis.fetch with a stub that captures the request and returns a canned response. */ +function stubFetch(status: number, body: unknown): { url: string; init?: RequestInit }[] { + const calls: { url: string; init?: RequestInit }[] = []; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + calls.push({ url: input.toString(), init }); + return new Response(JSON.stringify(body), { + status, + headers: { "content-type": "application/json" }, + }); + }) as typeof fetch; + return calls; +} + +describe("callXaiImages", () => { + test("no imageUrl → POST /images/generations", async () => { + const calls = stubFetch(200, { data: [{ b64_json: "dGVzdA==" }] }); + await callXaiImages({ prompt: "a cat" }, AUTH); + expect(calls[0]!.url).toContain("/images/generations"); + expect(calls[0]!.init?.method).toBe("POST"); + }); + + test("with imageUrl → POST /images/edits", async () => { + const calls = stubFetch(200, { data: [{ b64_json: "dGVzdA==" }] }); + await callXaiImages({ prompt: "edit this", imageUrl: "https://example.com/img.png" }, AUTH); + expect(calls[0]!.url).toContain("/images/edits"); + }); + + test("request body has correct model, prompt, n", async () => { + const calls = stubFetch(200, { data: [{ b64_json: "dGVzdA==" }] }); + await callXaiImages({ prompt: "a dog", model: "grok-imagine-fast", n: 3 }, AUTH); + const body = JSON.parse((calls[0]!.init?.body as string) ?? "{}"); + expect(body.model).toBe("grok-imagine-fast"); + expect(body.prompt).toBe("a dog"); + expect(body.n).toBe(3); + }); + + test("non-2xx → throws Error containing status code", async () => { + stubFetch(429, { error: "rate limited" }); + await expect(callXaiImages({ prompt: "x" }, AUTH)).rejects.toThrow("429"); + }); + + test("2xx with b64_json → returns normalized XaiImageResult", async () => { + stubFetch(200, { data: [{ b64_json: "dGVzdA==" }] }); + const result = await callXaiImages({ prompt: "x" }, AUTH); + expect(result.images.length).toBe(1); + expect(result.images[0]!.b64_json).toBe("dGVzdA=="); + }); + + test("2xx with url → returns images[0].url", async () => { + stubFetch(200, { data: [{ url: "https://cdn.example.com/img.png" }] }); + const result = await callXaiImages({ prompt: "x" }, AUTH); + expect(result.images[0]!.url).toBe("https://cdn.example.com/img.png"); + }); + + test("caller abort propagates into the composed signal", async () => { + const calls = stubFetch(200, { data: [{ b64_json: "dGVzdA==" }] }); + const controller = new AbortController(); + await callXaiImages({ prompt: "x" }, AUTH, controller.signal); + const passed = calls[0]!.init?.signal as AbortSignal; + expect(passed.aborted).toBe(false); + controller.abort("client gone"); + expect(passed.aborted).toBe(true); + }); + + test("size/quality mapped to aspect_ratio/resolution, no passthrough", async () => { + const calls = stubFetch(200, { data: [{ b64_json: "dGVzdA==" }] }); + await callXaiImages({ prompt: "x", size: "1024x1792", quality: "hd" }, AUTH); + const body = JSON.parse((calls[0]!.init?.body as string) ?? "{}"); + expect(body.aspect_ratio).toBe("9:16"); + expect(body.resolution).toBe("2k"); + expect(body).not.toHaveProperty("size"); + expect(body).not.toHaveProperty("quality"); + }); + + test("square size → 1:1", async () => { + const calls = stubFetch(200, { data: [{ b64_json: "dGVzdA==" }] }); + await callXaiImages({ prompt: "x", size: "1024x1024" }, AUTH); + const body = JSON.parse((calls[0]!.init?.body as string) ?? "{}"); + expect(body.aspect_ratio).toBe("1:1"); + expect(body).not.toHaveProperty("resolution"); + }); + + test("quality: standard → 1k", async () => { + const calls = stubFetch(200, { data: [{ b64_json: "dGVzdA==" }] }); + await callXaiImages({ prompt: "x", quality: "standard" }, AUTH); + const body = JSON.parse((calls[0]!.init?.body as string) ?? "{}"); + expect(body.resolution).toBe("1k"); + expect(body).not.toHaveProperty("aspect_ratio"); + }); + + test("unknown size/quality dropped", async () => { + const calls = stubFetch(200, { data: [{ b64_json: "dGVzdA==" }] }); + await callXaiImages({ prompt: "x", size: "weird", quality: "ultra" }, AUTH); + const body = JSON.parse((calls[0]!.init?.body as string) ?? "{}"); + expect(body).not.toHaveProperty("aspect_ratio"); + expect(body).not.toHaveProperty("resolution"); + expect(body).not.toHaveProperty("size"); + expect(body).not.toHaveProperty("quality"); + }); +}); diff --git a/tests/responses-parser.test.ts b/tests/responses-parser.test.ts index fce5c005e6..7f49ce4356 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(); From 6428f5d5c5a240570e05dd31a753a85d84204998 Mon Sep 17 00:00:00 2001 From: Your Name Date: Mon, 27 Jul 2026 08:01:37 +0800 Subject: [PATCH 02/15] fix(images): address Wibias review blockers for #424 - Enforce HTTPS-only scheme in downloadImageToArtifact (reject ftp/http/gopher) - Fix dispatch order: image bridge defers to web-search when both eligible - Narrow buildImageTool description to generation-only (no 'edit' promise) - Remove URL-interpolating error from fulfill.ts console.warn - Add downloadImageToArtifact SSRF tests (http/ftp reject, https succeeds) - Add handler-activation regression tests (stream/400/dual-tool defer) --- src/images/artifacts.ts | 8 +- src/images/fulfill.ts | 5 +- src/images/synthetic-tool.ts | 4 +- src/server/responses/core.ts | 46 ++++---- tests/images/artifacts-ssrf.test.ts | 44 +++++++- tests/images/handler-activation.test.ts | 143 ++++++++++++++++++++++++ 6 files changed, 222 insertions(+), 28 deletions(-) create mode 100644 tests/images/handler-activation.test.ts diff --git a/src/images/artifacts.ts b/src/images/artifacts.ts index f4eb700461..0ed1a2aed7 100644 --- a/src/images/artifacts.ts +++ b/src/images/artifacts.ts @@ -90,7 +90,13 @@ export async function downloadImageToArtifact( } // SSRF protection: validate the provider-returned URL before fetching. - // Reject non-HTTP(S) schemes, literal private/loopback/link-local/metadata addresses. + // Require HTTPS strictly — plain HTTP and all other schemes (ftp, file, …) are rejected. + let parsedUrl: URL; + try { parsedUrl = new URL(url); } catch { throw new Error("image URL is not valid"); } + if (parsedUrl.protocol !== "https:") { + throw new Error(`image URL must use HTTPS, got ${parsedUrl.protocol}`); + } + // Reject literal private/loopback/link-local/metadata addresses. const assessment = assessUrlDestination(url); if (assessment && assessment.kind !== "public" && assessment.kind !== "hostname") { throw new Error(`image URL targets ${assessment.detail}`); diff --git a/src/images/fulfill.ts b/src/images/fulfill.ts index 4751f35ac9..a8fe6ad7e1 100644 --- a/src/images/fulfill.ts +++ b/src/images/fulfill.ts @@ -52,9 +52,10 @@ export async function fulfillImageCall( } else if (img.url) { files.push(await downloadImageToArtifact(img.url, budget, signal)); } - } catch (e) { + } 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: ${e instanceof Error ? e.message : String(e)}`); + console.warn("[images] failed to materialize image"); } } diff --git a/src/images/synthetic-tool.ts b/src/images/synthetic-tool.ts index a7707a0ff8..f6b1a7cbe5 100644 --- a/src/images/synthetic-tool.ts +++ b/src/images/synthetic-tool.ts @@ -55,8 +55,8 @@ export function buildImageTool(): OcxTool { return { name: IMAGE_GEN_TOOL_NAME, description: - "Generate or edit an image. Returns absolute local filesystem path(s). " + - "Use when the user asks to create, draw, or edit an image.", + "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: { diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index 104eb833b1..de1c122863 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -1532,29 +1532,33 @@ export async function handleResponses( // internally supports both standard and runTurn adapter paths. const imgPlan = await planImageBridge(config, parsed, route.provider); if (imgPlan) { - // The bridge forces stream:true internally and returns SSE. Non-streaming requests can't be - // served — reject explicitly rather than returning SSE to a client expecting JSON. - if (!parsed.stream) { - return formatErrorResponse(400, "invalid_request_error", "image bridge requires stream=true"); - } - parsed.context.tools = [...(parsed.context.tools ?? []), buildImageTool()]; - const imgResponse = await runWithImageBridge({ - parsed, adapter, - plan: imgPlan, - forwardHeaders: selectedForwardHeaders, - onAttemptSend: () => noteAttemptSend(logCtx.activeAttempt, logCtx.usageLogInputTokens), - abortSignal: options.abortSignal, - ...(config.images?.maxRounds != null ? { maxRounds: config.images.maxRounds } : {}), - ...(options.onFirstOutput ? { onFirstOutput: options.onFirstOutput } : {}), - }); - if (imgResponse.body) { - const imgTurnAc = new AbortController(); - return new Response(trackStreamLifetime(imgResponse.body, imgTurnAc), { - status: imgResponse.status, - headers: imgResponse.headers, + // Web-search takes priority when both are eligible (design: image defers to web-search). + const wsPlan = planWebSearch(config, parsed, false, route.provider, route.modelId, openAiSidecar); + if (!wsPlan) { + // The bridge forces stream:true internally and returns SSE. Non-streaming requests can't be + // served — reject explicitly rather than returning SSE to a client expecting JSON. + if (!parsed.stream) { + return formatErrorResponse(400, "invalid_request_error", "image bridge requires stream=true"); + } + parsed.context.tools = [...(parsed.context.tools ?? []), buildImageTool()]; + const imgResponse = await runWithImageBridge({ + parsed, adapter, + plan: imgPlan, + forwardHeaders: selectedForwardHeaders, + onAttemptSend: () => noteAttemptSend(logCtx.activeAttempt, logCtx.usageLogInputTokens), + abortSignal: options.abortSignal, + ...(config.images?.maxRounds != null ? { maxRounds: config.images.maxRounds } : {}), + ...(options.onFirstOutput ? { onFirstOutput: options.onFirstOutput } : {}), }); + if (imgResponse.body) { + const imgTurnAc = new AbortController(); + return new Response(trackStreamLifetime(imgResponse.body, imgTurnAc), { + status: imgResponse.status, + headers: imgResponse.headers, + }); + } + return imgResponse; } - return imgResponse; } if (adapter.runTurn) { diff --git a/tests/images/artifacts-ssrf.test.ts b/tests/images/artifacts-ssrf.test.ts index 37f36942ee..c13180225e 100644 --- a/tests/images/artifacts-ssrf.test.ts +++ b/tests/images/artifacts-ssrf.test.ts @@ -1,5 +1,12 @@ -import { describe, expect, test } from "bun:test"; -import { assessUrlDestination, assertUrlResolvesPublic } from "../../src/lib/destination-policy"; +import { rm } from "node:fs/promises"; +import { describe, expect, mock, test } from "bun:test"; + +// Mock DNS before importing destination-policy — it binds `lookup` at load time. +const lookupMock = mock(async (_hostname: string, _opts: unknown): Promise<{ address: string; family: number }[]> => []); +mock.module("node:dns/promises", () => ({ lookup: lookupMock })); + +const { assessUrlDestination, assertUrlResolvesPublic } = await import("../../src/lib/destination-policy"); +const { downloadImageToArtifact } = await import("../../src/images/artifacts"); describe("SSRF: assessUrlDestination", () => { test("loopback IPv4 → loopback", () => { @@ -49,3 +56,36 @@ describe("SSRF: assertUrlResolvesPublic", () => { await expect(assertUrlResolvesPublic("not-a-url")).rejects.toThrow(); }); }); + +describe("SSRF: downloadImageToArtifact scheme enforcement", () => { + test("http:// → rejects (non-HTTPS)", async () => { + await expect(downloadImageToArtifact("http://public-host/path")).rejects.toThrow(/HTTPS/); + }); + + test("ftp:// → rejects", async () => { + await expect(downloadImageToArtifact("ftp://host/path")).rejects.toThrow(/HTTPS/); + }); + + test("https:// public host → succeeds with mocked fetch", async () => { + // Stub DNS so public-host resolves to a public address. + lookupMock.mockResolvedValue([{ address: "93.184.216.34", family: 4 }]); + const originalFetch = globalThis.fetch; + let downloadedPath: string | undefined; + try { + globalThis.fetch = (async (input: RequestInfo | URL, _init?: RequestInit) => { + const raw = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; + if (!raw.startsWith("https://")) throw new Error("fetch must only be called over HTTPS"); + // Minimal PNG signature so guessExtFromMagic returns "png". + const pngBytes = new Uint8Array([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]); + return new Response(pngBytes, { status: 200 }); + }) as typeof fetch; + + downloadedPath = await downloadImageToArtifact("https://public-host/valid-image"); + expect(downloadedPath).toMatch(/dl-.*\.png$/); + } finally { + globalThis.fetch = originalFetch; + lookupMock.mockClear(); + if (downloadedPath) await rm(downloadedPath).catch(() => {}); + } + }); +}); diff --git a/tests/images/handler-activation.test.ts b/tests/images/handler-activation.test.ts new file mode 100644 index 0000000000..97470bf765 --- /dev/null +++ b/tests/images/handler-activation.test.ts @@ -0,0 +1,143 @@ +import { afterAll, beforeAll, describe, expect, mock, test } from "bun:test"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { randomUUID } from "node:crypto"; +import type { OcxConfig, OcxProviderConfig } from "../../src/types"; +import type { ProviderAdapter } from "../../src/adapters/base"; + +/** + * Dispatch-priority regression test for the image bridge (PR #424). + * + * The image bridge and the web-search sidecar are both opt-in dispatch paths in + * handleResponses(). The design contract is "image defers to web-search": when a + * request is eligible for BOTH, the web-search sidecar wins and the image bridge + * must NOT activate. This was previously broken because planImageBridge ran and + * returned before planWebSearch was ever consulted. + * + * These tests drive handleResponses() end-to-end (real parser + real routing + + * real planImageBridge) with only the adapter, the runners, and the web-search + * planner stubbed, so they exercise the actual dispatch ordering in + * src/server/responses/core.ts. + * + * NOTE: Full server-level integration testing of every adapter path is out of + * scope here — the focus is the dispatch priority ordering at the planImageBridge + * / planWebSearch fork (core.ts ~L1516). + */ + +const PREV_HOME = process.env.OPENCODEX_HOME; +beforeAll(() => { process.env.OPENCODEX_HOME = join(tmpdir(), "ocx-test-" + randomUUID()); }); +afterAll(() => { if (PREV_HOME === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = PREV_HOME; }); + +// --- Activation spies, flipped by the stubbed runners --- +let imageBridgeRun = false; +let webSearchRun = false; +/** Controlled return value for the stubbed planWebSearch (truthy ⇒ web-search plan active). */ +let mockWsPlan: unknown = undefined; + +// --- Stub adapter-resolve: inject a minimal adapter so no real upstream is hit --- +const actualResolver = await import("../../src/server/adapter-resolve"); +mock.module("../../src/server/adapter-resolve", () => ({ + ...actualResolver, + resolveAdapter(provider: OcxProviderConfig) { + return { + 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 }; }, + } as ProviderAdapter; + }, +})); + +// --- Stub the image bridge runner: detect activation without hitting the real loop --- +mock.module("../../src/images/loop", () => ({ + runWithImageBridge: async () => { + imageBridgeRun = true; + return new Response("data: {\"type\":\"done\"}\n\n", { + status: 200, headers: { "content-type": "text/event-stream" }, + }); + }, +})); + +// --- Stub the web-search planner + runner: control eligibility and detect activation --- +mock.module("../../src/web-search/index", () => ({ + // Re-export the symbols parser.ts imports statically (deep path → no mock there). + buildWebSearchTool: () => ({ name: "web_search", parameters: { type: "object", properties: {} } }), + WEB_SEARCH_TOOL_NAME: "web_search", + extractHostedWebSearch: (tools: unknown[]) => { + if (!Array.isArray(tools)) return undefined; + for (const t of tools) { + if (t && typeof t === "object" && (t as Record).type === "web_search") { + return { search_context_size: "medium" }; + } + } + return undefined; + }, + runWithWebSearch: async () => { + webSearchRun = true; + return new Response("data: {\"type\":\"done\"}\n\n", { + status: 200, headers: { "content-type": "text/event-stream" }, + }); + }, + planWebSearch: () => mockWsPlan, + shouldResolveOpenAiWebSearchSidecar: () => false, +})); + +const { handleResponses } = await import("../../src/server/responses"); + +/** Routed (non-OpenAI) keyed provider + an xAI provider with an API key so the real planImageBridge returns a plan. */ +function makeConfig(): OcxConfig { + return { + port: 0, + defaultProvider: "fixture", + providers: { + fixture: { adapter: "openai-chat", baseUrl: "https://fixture.test/v1", authMode: "key", apiKey: "fixture-key" }, + xai: { adapter: "openai", baseUrl: "https://api.x.ai", apiKey: "xai-test-token" }, + }, + images: { bridgeEnabled: true }, + } as OcxConfig; +} + +function post(stream: boolean, tools: unknown[]): Promise { + return handleResponses( + new Request("http://localhost/v1/responses", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ model: "fixture/model", input: "hello", stream, tools }), + }), + makeConfig(), + { model: "", provider: "" } as never, + {}, + ); +} + +describe("image bridge dispatch priority (handler activation)", () => { + test("stream=true + image_generation tool → image bridge activates and returns SSE", async () => { + imageBridgeRun = false; webSearchRun = false; mockWsPlan = undefined; + const res = await post(true, [{ type: "image_generation" }]); + expect(imageBridgeRun).toBe(true); + expect(res.headers.get("content-type")).toBe("text/event-stream"); + }); + + test("stream=false + image_generation tool → 400 (bridge requires stream=true)", async () => { + imageBridgeRun = false; webSearchRun = false; mockWsPlan = undefined; + const res = await post(false, [{ type: "image_generation" }]); + expect(res.status).toBe(400); + // The runner must not execute when the request is rejected upfront. + expect(imageBridgeRun).toBe(false); + expect((await res.text())).toContain("image bridge requires stream=true"); + }); + + test("dual-tool (image_generation + web_search), both eligible → web-search wins, image bridge deferred", async () => { + imageBridgeRun = false; webSearchRun = false; + // A truthy plan ⇒ web-search is eligible; the image bridge must defer to it. + mockWsPlan = { backend: "openai" }; + const res = await post(true, [{ type: "web_search" }, { type: "image_generation" }]); + expect(webSearchRun).toBe(true); + expect(imageBridgeRun).toBe(false); + expect(res.headers.get("content-type")).toBe("text/event-stream"); + }); +}); From 58d06721e5cfb1d6ca094d9b18b0378f7031ec52 Mon Sep 17 00:00:00 2001 From: Your Name Date: Mon, 27 Jul 2026 09:36:02 +0800 Subject: [PATCH 03/15] fix(images): address Wibias R2 review blockers for #424 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Move web-search dispatch before runTurn so dual-tool turns reach web-search even on Cursor/runTurn adapters - Guard image bridge with !routedCompaction to prevent hijacking compaction requests - Fix IPv4-mapped IPv6 hex SSRF bypass in destination-policy (::ffff:7f00:1 now decoded and classified as loopback) - Add SSRF tests: private IP, gopher, IPv6 mapped (dotted+hex), redirect:error fail-closed - Add regression: runTurn dual-tool → web-search wins - Add regression: compaction request → image bridge skipped --- src/lib/destination-policy.ts | 10 ++ src/server/responses/core.ts | 124 +++++++++++++----------- tests/images/artifacts-ssrf.test.ts | 39 ++++++++ tests/images/handler-activation.test.ts | 61 +++++++++++- 4 files changed, 173 insertions(+), 61 deletions(-) diff --git a/src/lib/destination-policy.ts b/src/lib/destination-policy.ts index 1ec932bd4c..3ed0bf5330 100644 --- a/src/lib/destination-policy.ts +++ b/src/lib/destination-policy.ts @@ -79,6 +79,16 @@ function classifyIpv6(hostname: string): DestinationAssessment { if (BLOCKED_METADATA_IPV6.has(hostname)) return { kind: "metadata", detail: "blocked metadata endpoint" }; const mappedIpv4 = hostname.match(/^::ffff:(\d+\.\d+\.\d+\.\d+)$/i)?.[1]; if (mappedIpv4) return classifyIpv4(mappedIpv4); + // Decode hex IPv4-mapped IPv6: ::ffff:7f00:1 → 127.0.0.1 + // The dotted-decimal regex above only matches ::ffff:127.0.0.1; without this, + // hex form bypasses all private/loopback checks (hextet is 0 → classified "public"). + const hexMapped = hostname.match(/^::ffff:([0-9a-f]{1,4}):([0-9a-f]{1,4})$/i); + if (hexMapped) { + const hi = Number.parseInt(hexMapped[1], 16); + const lo = Number.parseInt(hexMapped[2], 16); + const ipv4 = `${(hi >> 8) & 255}.${hi & 255}.${(lo >> 8) & 255}.${lo & 255}`; + return classifyIpv4(ipv4); + } if (hostname === "::1") return { kind: "loopback", detail: "loopback address" }; if (hostname === "::") return { kind: "unspecified", detail: "unspecified address" }; const hextet = firstIpv6Hextet(hostname); diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index de1c122863..0f3aeff694 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -1530,7 +1530,10 @@ export async function handleResponses( // Image bridge: check BEFORE the runTurn early-return so runTurn adapters (e.g. Cursor) // also route through the bridge when image_generation is requested. The bridge loop // internally supports both standard and runTurn adapter paths. - const imgPlan = await planImageBridge(config, parsed, route.provider); + // Routed-compaction turns must NOT hit the bridge: compaction clears tools/_webSearch but + // leaves _imageGeneration, so planImageBridge would activate and return a normal Responses + // completion instead of the synthetic compaction item Codex expects (#424). + const imgPlan = !routedCompaction ? await planImageBridge(config, parsed, route.provider) : undefined; if (imgPlan) { // Web-search takes priority when both are eligible (design: image defers to web-search). const wsPlan = planWebSearch(config, parsed, false, route.provider, route.modelId, openAiSidecar); @@ -1561,6 +1564,67 @@ export async function handleResponses( } } + // Web-search sidecar: Codex enabled web_search but this is a routed (non-OpenAI) model that can't + // run it server-side. Expose web_search as a function tool and run searches via the gpt-mini sidecar + // through the ChatGPT passthrough, looping until the model answers. Otherwise take the normal path. + // Placed BEFORE the runTurn early-return so dual-tool turns (web_search + image_generation) on + // runTurn adapters (e.g. Cursor) dispatch through the web-search sidecar instead of being swallowed + // by the runTurn branch — which would leave dual-tool turns with neither bridge. + const wsPlan = planWebSearch(config, parsed, false, route.provider, route.modelId, openAiSidecar); + if (wsPlan) { + parsed.context.tools = [...(parsed.context.tools ?? []), buildWebSearchTool()]; + noteAttemptSend(logCtx.activeAttempt, logCtx.usageLogInputTokens); + const wsResponse = await runWithWebSearch({ + parsed, adapter, + backend: wsPlan.backend, + forwardProvider: wsPlan.forwardSidecar?.provider, + anthropicSidecar: wsPlan.anthropicSidecar, + hostedTool: wsPlan.hostedTool, + selectedForwardHeaders: wsPlan.forwardSidecar?.headers ?? selectedForwardHeaders, + settings: wsPlan.settings, + maxSearches: wsPlan.maxSearches, + forceEmptyResponseId: true, + abortSignal: options.abortSignal, + ...(options.onFirstOutput ? { onFirstOutput: options.onFirstOutput } : {}), + onRequestBuilt: request => recordAdapterReasoning(logCtx, request), + onUsage: usage => { + logCtx.usageFromBridge = true; + if (usage) { + logCtx.usage = usage; + if (logCtx.activeAttempt) logCtx.activeAttempt.usage = usage; + } + }, + recordSidecarOutcome: wsPlan.forwardSidecar?.recordOutcome, + connectTimeoutMs: config.connectTimeoutMs ?? 200_000, + routedModelStallTimeoutMs: wsPlan.routedModelStallTimeoutMs, + stallTimeoutSec: wsPlan.stallTimeoutSec, + on429: retryAfter => { + const rotated = rotateProviderTransportOn429(config, route.providerName, { + retryAfter, + now: Date.now(), + attemptedKey: route.provider.apiKey, + promptCacheKey: parsed.options.promptCacheKey, + }); + if (!rotated) return null; + route.provider = rotated; + return resolveAdapter( + resolveWireProtocolOverride(route.providerName, route.modelId, route.provider), + config.cacheRetention, + ); + }, + }); + // Register the sidecar stream as an active turn so drainAndShutdown waits for (or aborts) + // in-flight web-search turns instead of skipping them during graceful shutdown. + if (wsResponse.body) { + const wsTurnAc = new AbortController(); + return new Response(trackStreamLifetime(wsResponse.body, wsTurnAc), { + status: wsResponse.status, + headers: wsResponse.headers, + }); + } + return wsResponse; + } + if (adapter.runTurn) { const runTurnAbort = new AbortController(); linkAbortSignal(runTurnAbort, options.abortSignal); @@ -1681,64 +1745,6 @@ export async function handleResponses( return new Response(JSON.stringify(json), { headers: { "Content-Type": "application/json" } }); } - // Web-search sidecar: Codex enabled web_search but this is a routed (non-OpenAI) model that can't - // run it server-side. Expose web_search as a function tool and run searches via the gpt-mini sidecar - // through the ChatGPT passthrough, looping until the model answers. Otherwise take the normal path. - const wsPlan = planWebSearch(config, parsed, false, route.provider, route.modelId, openAiSidecar); - if (wsPlan) { - parsed.context.tools = [...(parsed.context.tools ?? []), buildWebSearchTool()]; - noteAttemptSend(logCtx.activeAttempt, logCtx.usageLogInputTokens); - const wsResponse = await runWithWebSearch({ - parsed, adapter, - backend: wsPlan.backend, - forwardProvider: wsPlan.forwardSidecar?.provider, - anthropicSidecar: wsPlan.anthropicSidecar, - hostedTool: wsPlan.hostedTool, - selectedForwardHeaders: wsPlan.forwardSidecar?.headers ?? selectedForwardHeaders, - settings: wsPlan.settings, - maxSearches: wsPlan.maxSearches, - forceEmptyResponseId: true, - abortSignal: options.abortSignal, - ...(options.onFirstOutput ? { onFirstOutput: options.onFirstOutput } : {}), - onRequestBuilt: request => recordAdapterReasoning(logCtx, request), - onUsage: usage => { - logCtx.usageFromBridge = true; - if (usage) { - logCtx.usage = usage; - if (logCtx.activeAttempt) logCtx.activeAttempt.usage = usage; - } - }, - recordSidecarOutcome: wsPlan.forwardSidecar?.recordOutcome, - connectTimeoutMs: config.connectTimeoutMs ?? 200_000, - routedModelStallTimeoutMs: wsPlan.routedModelStallTimeoutMs, - stallTimeoutSec: wsPlan.stallTimeoutSec, - on429: retryAfter => { - const rotated = rotateProviderTransportOn429(config, route.providerName, { - retryAfter, - now: Date.now(), - attemptedKey: route.provider.apiKey, - promptCacheKey: parsed.options.promptCacheKey, - }); - if (!rotated) return null; - route.provider = rotated; - return resolveAdapter( - resolveWireProtocolOverride(route.providerName, route.modelId, route.provider), - config.cacheRetention, - ); - }, - }); - // Register the sidecar stream as an active turn so drainAndShutdown waits for (or aborts) - // in-flight web-search turns instead of skipping them during graceful shutdown. - if (wsResponse.body) { - const wsTurnAc = new AbortController(); - return new Response(trackStreamLifetime(wsResponse.body, wsTurnAc), { - status: wsResponse.status, - headers: wsResponse.headers, - }); - } - return wsResponse; - } - const upstream = new AbortController(); const cleanupUpstreamAbort = linkAbortSignal(upstream, options.abortSignal); const connectMs = config.connectTimeoutMs ?? 200_000; diff --git a/tests/images/artifacts-ssrf.test.ts b/tests/images/artifacts-ssrf.test.ts index c13180225e..18039f584d 100644 --- a/tests/images/artifacts-ssrf.test.ts +++ b/tests/images/artifacts-ssrf.test.ts @@ -37,6 +37,15 @@ describe("SSRF: assessUrlDestination", () => { 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(); }); @@ -66,6 +75,36 @@ describe("SSRF: downloadImageToArtifact scheme enforcement", () => { await expect(downloadImageToArtifact("ftp://host/path")).rejects.toThrow(/HTTPS/); }); + test("gopher:// → rejects (non-HTTPS)", async () => { + await expect(downloadImageToArtifact("gopher://host/path")).rejects.toThrow(/HTTPS/); + }); + + test("private 10.x via download helper → rejects", async () => { + await expect(downloadImageToArtifact("https://10.0.0.1/img.png")).rejects.toThrow(); + }); + + test("IPv4-mapped IPv6 [::ffff:127.0.0.1] via download helper → rejects", async () => { + await expect(downloadImageToArtifact("https://[::ffff:127.0.0.1]/image.png")).rejects.toThrow(); + }); + + test("IPv4-mapped IPv6 hex [::ffff:7f00:1] via download helper → rejects", async () => { + await expect(downloadImageToArtifact("https://[::ffff:7f00:1]/image.png")).rejects.toThrow(); + }); + + test(`3xx redirect response → rejects (redirect: 'error')`, async () => { + lookupMock.mockResolvedValue([{ address: "93.184.216.34", family: 4 }]); + const originalFetch = globalThis.fetch; + try { + globalThis.fetch = (async (_input: RequestInfo | URL, _init?: RequestInit) => { + return new Response("", { status: 301, headers: { Location: "https://evil.example/redirect" } }); + }) as typeof fetch; + await expect(downloadImageToArtifact("https://public-host/redirect-img")).rejects.toThrow(); + } finally { + globalThis.fetch = originalFetch; + lookupMock.mockClear(); + } + }); + test("https:// public host → succeeds with mocked fetch", async () => { // Stub DNS so public-host resolves to a public address. lookupMock.mockResolvedValue([{ address: "93.184.216.34", family: 4 }]); diff --git a/tests/images/handler-activation.test.ts b/tests/images/handler-activation.test.ts index 97470bf765..64152f4049 100644 --- a/tests/images/handler-activation.test.ts +++ b/tests/images/handler-activation.test.ts @@ -31,6 +31,10 @@ afterAll(() => { if (PREV_HOME === undefined) delete 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; @@ -39,7 +43,7 @@ const actualResolver = await import("../../src/server/adapter-resolve"); mock.module("../../src/server/adapter-resolve", () => ({ ...actualResolver, resolveAdapter(provider: OcxProviderConfig) { - return { + const base = { name: "test", buildRequest: async () => ({ url: provider.baseUrl, method: "POST", headers: {}, body: "" }), async fetchResponse() { @@ -48,7 +52,17 @@ mock.module("../../src/server/adapter-resolve", () => ({ }); }, async *parseStream() { yield { type: "done" as const }; }, - } as ProviderAdapter; + }; + 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; }, })); @@ -140,4 +154,47 @@ describe("image bridge dispatch priority (handler activation)", () => { expect(imageBridgeRun).toBe(false); expect(res.headers.get("content-type")).toBe("text/event-stream"); }); + + test("routed compaction with image_generation tool → image bridge does NOT hijack compaction (#424)", async () => { + imageBridgeRun = false; webSearchRun = false; mockWsPlan = undefined; + // A routed-compaction request carries both _compactionRequest and _imageGeneration: + // compaction clears tools/_webSearch but leaves _imageGeneration, so without the + // routedCompaction guard planImageBridge would activate and return a normal Responses + // completion instead of the synthetic compaction item Codex expects. + const res = await handleResponses( + new Request("http://localhost/v1/responses", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + model: "fixture/model", + input: [{ type: "compaction_trigger" }], + stream: true, + tools: [{ type: "image_generation" }], + }), + }), + makeConfig(), + { model: "", provider: "" } as never, + {}, + ); + expect(imageBridgeRun).toBe(false); + expect(res.headers.get("content-type")).toBe("text/event-stream"); + }); + + test("dual-tool on a runTurn adapter → web-search path wins, runTurn does not eat the request (#424)", async () => { + imageBridgeRun = false; webSearchRun = false; runTurnCalled = false; + useRunTurnAdapter = true; + mockWsPlan = { backend: "openai" }; + try { + const res = await post(true, [{ type: "web_search" }, { type: "image_generation" }]); + // Web-search sidecar must handle the turn. + expect(webSearchRun).toBe(true); + // Image bridge must not activate (design: image defers to web-search). + expect(imageBridgeRun).toBe(false); + // runTurn must NOT be reached — the web-search dispatch now runs before the runTurn early-return. + expect(runTurnCalled).toBe(false); + expect(res.headers.get("content-type")).toBe("text/event-stream"); + } finally { + useRunTurnAdapter = false; + } + }); }); From b96bfa9732be7cb37771b170731df98ca8679f6c Mon Sep 17 00:00:00 2001 From: Wibias <37517432+Wibias@users.noreply.github.com> Date: Mon, 27 Jul 2026 04:07:21 +0200 Subject: [PATCH 04/15] fix(images): address Codex P2 follow-ups for image bridge Clamp maxRounds, honor images.timeoutMs, strip all image tool aliases on forced-final, return SSE headers before runTurn collect, preserve Cursor conversation ids across loop iterations, wire onUsage/logCtx and 429 key rotation, and keep parallel image calls in one assistant turn. --- src/images/fulfill.ts | 7 +- src/images/index.ts | 2 +- src/images/loop.ts | 216 ++++++++++++++++-------- src/images/plan.ts | 3 + src/images/types.ts | 2 + src/images/xai-client.ts | 6 +- src/server/responses/core.ts | 25 ++- src/types.ts | 4 +- tests/images/handler-activation.test.ts | 6 + tests/images/loop.test.ts | 202 +++++++++++++++++++++- tests/images/plan.test.ts | 11 +- tests/images/xai-client.test.ts | 8 + 12 files changed, 415 insertions(+), 77 deletions(-) diff --git a/src/images/fulfill.ts b/src/images/fulfill.ts index a8fe6ad7e1..907923e0ab 100644 --- a/src/images/fulfill.ts +++ b/src/images/fulfill.ts @@ -38,7 +38,12 @@ export async function fulfillImageCall( let result; try { - result = await callXaiImages({ prompt, model: plan.model, n, imageUrl, size, quality }, plan.auth, signal); + 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 }; diff --git a/src/images/index.ts b/src/images/index.ts index 6f0adb1fd2..9a402706ba 100644 --- a/src/images/index.ts +++ b/src/images/index.ts @@ -1,4 +1,4 @@ export { planImageBridge, findXaiProvider, resolveXaiToken } from "./plan"; -export { runWithImageBridge } from "./loop"; +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 index e5b2c57865..e48369ecc1 100644 --- a/src/images/loop.ts +++ b/src/images/loop.ts @@ -7,12 +7,12 @@ * the model produces a real tool call or the budget is exhausted, the passthrough events are * replayed to the bridge for final SSE output. * - * Removed vs web-search: no sidecar backend selection, no 429 key-failover, no forced-answer - * nudge, no failed-query dedup, no describeImages/structuredOutput, no recordSidecarOutcome. + * Removed vs web-search: no sidecar backend selection, no forced-answer nudge, no failed-query + * dedup, no describeImages/structuredOutput, no recordSidecarOutcome. */ import type { ProviderAdapter } from "../adapters/base"; import { createAdapterEventQueue } from "../adapters/run-turn-queue"; -import type { AdapterEvent, OcxMessage, OcxParsedRequest, OcxThinkingContent } from "../types"; +import type { AdapterEvent, OcxMessage, OcxParsedRequest, OcxThinkingContent, OcxUsage } from "../types"; import { namespacedToolName } from "../types"; import { bridgeToResponsesSSE } from "../bridge"; import { clearableDeadline } from "../lib/abort"; @@ -32,7 +32,18 @@ const SSE_HEADERS = { const CONNECT_TIMEOUT_MS = 200_000; const STALL_TIMEOUT_MS = 200_000; -const DEFAULT_MAX_ROUNDS = 3; +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; + +/** + * 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; @@ -140,8 +151,15 @@ export interface ImageBridgeDeps { onAttemptSend?: () => void; abortSignal?: AbortSignal; onFirstOutput?: () => void; - /** Max image-generation rounds before forcing a final answer. Defaults to 3. */ + /** Max image-generation rounds before forcing a final answer. Defaults to 3; clamped to [0, 10]. */ maxRounds?: number; + /** 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; } /** @@ -151,15 +169,23 @@ export interface ImageBridgeDeps { * inject the answer as a tool_result, and loop (bounded by `maxRounds`). */ export async function runWithImageBridge(deps: ImageBridgeDeps): Promise { - const { parsed, adapter, plan, abortSignal } = deps; - const maxRounds = Math.max(0, deps.maxRounds ?? DEFAULT_MAX_ROUNDS); + const { parsed, plan, abortSignal } = deps; + let adapter = deps.adapter; + const maxRounds = clampImageMaxRounds(deps.maxRounds ?? DEFAULT_MAX_ROUNDS); const HARD_CAP = maxRounds + 1; const messages: OcxMessage[] = [...parsed.context.messages]; const allTools = parsed.context.tools ?? []; - // For the forced-final pass we drop image tools so the model MUST answer from the results already - // in `messages` (can't generate again) — this guarantees a non-empty final answer. - const toolsNoImage = allTools.filter(t => !t.imageGeneration); + // 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 @@ -179,7 +205,8 @@ export async function runWithImageBridge(deps: ImageBridgeDeps): Promise; // 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. + // 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, @@ -205,7 +232,28 @@ export async function runWithImageBridge(deps: ImageBridgeDeps): Promise ({ done: true as const, value })), + new Promise<{ done: false }>(resolve => { + const timer = setTimeout(() => resolve({ done: false }), 15_000); + void collectPromise.finally(() => clearTimeout(timer)); + }), + ]); + if (raced.done) events = raced.value; + else yield { type: "heartbeat" }; + } + + // 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; + } + deps.onAttemptSend?.(); // runTurn adapters signal errors via {type:"error"} events, not HTTP status codes. @@ -214,9 +262,6 @@ export async function runWithImageBridge(deps: ImageBridgeDeps): Promise { - const h = new Headers(request.headers); - if (!h.has("accept-encoding")) h.set("accept-encoding", "identity"); - return fetch(request.url, { - method: request.method, - headers: h, - body: request.body, - signal: headerDeadline.signal, - }); - }, - { abortSignal: headerDeadline.signal, label: "image-bridge-loop" }, - ); + const fetchOnce = async (requestAdapter: ProviderAdapter): Promise => { + const request = await requestAdapter.buildRequest(iterParsed, { + headers: deps.forwardHeaders ? new Headers(deps.forwardHeaders) : new Headers(), + abortSignal: headerDeadline.signal, + }); + deps.onAttemptSend?.(); + const response = requestAdapter.fetchResponse + ? await requestAdapter.fetchResponse(request, { + abortSignal: headerDeadline.signal, + timeoutMs: CONNECT_TIMEOUT_MS, + returnRawErrors: true, + stream: true, + }) + : await fetchWithResetRetry( + () => { + const h = new Headers(request.headers); + if (!h.has("accept-encoding")) h.set("accept-encoding", "identity"); + return fetch(request.url, { + method: request.method, + headers: h, + body: request.body, + signal: headerDeadline.signal, + }); + }, + { abortSignal: headerDeadline.signal, label: "image-bridge-loop" }, + ); + 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 (!response.ok) { + if (!prepared.response.ok) { let body: Awaited>; try { - body = await readBoundedResponseBody(response, { signal }); + body = await readBoundedResponseBody(prepared.response, { signal }); } catch { if (signal.aborted) throw new LoopError(499, "client closed request during image-bridge"); - throw new LoopError(response.status, `Provider error ${response.status}`); + throw new LoopError(prepared.response.status, `Provider error ${prepared.response.status}`); } let formatted = ""; - if (body.displaySafe && !body.truncated && body.text.trim() && adapter.formatErrorBody) { + if (body.displaySafe && !body.truncated && body.text.trim() && prepared.responseAdapter.formatErrorBody) { try { - formatted = adapter.formatErrorBody(response.status, response.headers, body.text).trim(); + 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(response.status, `Provider error ${response.status}${suffix}`); + throw new LoopError(prepared.response.status, `Provider error ${prepared.response.status}${suffix}`); } - return { response, responseAdapter: adapter }; + return prepared; } catch (error) { if (headerDeadline.didExpire()) { throw new LoopError(504, `Provider response-header timeout after ${CONNECT_TIMEOUT_MS}ms during image-bridge`); @@ -324,14 +387,18 @@ export async function runWithImageBridge(deps: ImageBridgeDeps): Promise(); @@ -351,13 +418,15 @@ export async function runWithImageBridge(deps: ImageBridgeDeps): Promise= maxRounds; try { - // First loop turn reuses the eager HEADERS. Subsequent header acquisitions run here. - if (i > 0) { + // 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 @@ -368,16 +437,17 @@ export async function runWithImageBridge(deps: ImageBridgeDeps): Promise>; args: Record }> = []; + for (const call of split.calls) { yield { type: "heartbeat" }; const result = await fulfillImageCall( { id: call.id, name: call.name, arguments: call.args }, plan, budget, signal, ); if (signal.aborted) throw new LoopError(499, "client closed request during image-bridge"); - const now = Date.now(); let parsedArgs: Record = {}; try { const raw: unknown = JSON.parse(call.args || "{}"); @@ -385,14 +455,23 @@ export async function runWithImageBridge(deps: ImageBridgeDeps): Promise; } } catch { /* malformed args */ } - messages.push({ - role: "assistant", - content: [ - ...(callIndex === 0 && iterationThinking ? [iterationThinking] : []), - { type: "toolCall" as const, id: call.id, name: call.name, arguments: parsedArgs }, - ], - timestamp: now, - }); + 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, @@ -424,6 +503,7 @@ export async function runWithImageBridge(deps: ImageBridgeDeps): Promise 0 + ? { timeoutMs: Math.floor(config.images.timeoutMs) } + : {}), }; } diff --git a/src/images/types.ts b/src/images/types.ts index b0d139aeff..ed15701dbb 100644 --- a/src/images/types.ts +++ b/src/images/types.ts @@ -5,6 +5,8 @@ export interface ImageBridgePlan { auth: { baseUrl: string; token: string }; model: string; toolNames: Set; + /** Per-call xAI deadline (ms). Defaults inside callXaiImages when omitted. */ + timeoutMs?: number; } export interface ImageCallResult { diff --git a/src/images/xai-client.ts b/src/images/xai-client.ts index 4bb3b67b0d..f66c7dfc42 100644 --- a/src/images/xai-client.ts +++ b/src/images/xai-client.ts @@ -65,6 +65,7 @@ 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"; @@ -82,7 +83,10 @@ export async function callXaiImages( body.image = { url: req.imageUrl as string, type: "image_url" }; } - const timeout = AbortSignal.timeout(XAI_IMAGES_TIMEOUT_MS); + 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 resp = await fetch(`${auth.baseUrl}${endpoint}`, { diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index 0f3aeff694..fb59b60fd2 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -39,7 +39,7 @@ import { UnsupportedOAuthProviderError, } from "../../oauth"; import { buildWebSearchTool, planWebSearch, runWithWebSearch, shouldResolveOpenAiWebSearchSidecar } from "../../web-search"; -import { buildImageTool, planImageBridge, runWithImageBridge } from "../../images"; +import { buildImageTool, planImageBridge, runWithImageBridge, clampImageMaxRounds } from "../../images"; import { describeImagesInPlace, planVisionSidecar, shouldResolveOpenAiVisionSidecar, stripImagesInPlace } from "../../vision"; import { createAdapterEventQueue, preflightAdapterEvents } from "../../adapters/run-turn-queue"; import { @@ -1550,7 +1550,28 @@ export async function handleResponses( forwardHeaders: selectedForwardHeaders, onAttemptSend: () => noteAttemptSend(logCtx.activeAttempt, logCtx.usageLogInputTokens), abortSignal: options.abortSignal, - ...(config.images?.maxRounds != null ? { maxRounds: config.images.maxRounds } : {}), + maxRounds: clampImageMaxRounds(config.images?.maxRounds), + onUsage: usage => { + 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 } : {}), }); if (imgResponse.body) { diff --git a/src/types.ts b/src/types.ts index 9b9afff3a1..0f906998e9 100644 --- a/src/types.ts +++ b/src/types.ts @@ -743,13 +743,13 @@ 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 (see DEFAULT_MAX_ROUNDS in images/loop.ts). */ + /** Max image-generation loop iterations before forced-final. Default 3; clamped to [0, 10]. */ maxRounds?: number; } diff --git a/tests/images/handler-activation.test.ts b/tests/images/handler-activation.test.ts index 64152f4049..8b09aaf3e5 100644 --- a/tests/images/handler-activation.test.ts +++ b/tests/images/handler-activation.test.ts @@ -74,6 +74,12 @@ mock.module("../../src/images/loop", () => ({ status: 200, headers: { "content-type": "text/event-stream" }, }); }, + clampImageMaxRounds: (value: unknown) => { + if (typeof value !== "number" || !Number.isFinite(value)) return 3; + return Math.max(0, Math.min(10, Math.floor(value))); + }, + DEFAULT_MAX_ROUNDS: 3, + MAX_ROUNDS_HARD_LIMIT: 10, })); // --- Stub the web-search planner + runner: control eligibility and detect activation --- diff --git a/tests/images/loop.test.ts b/tests/images/loop.test.ts index 82d8e7b564..6462f9656a 100644 --- a/tests/images/loop.test.ts +++ b/tests/images/loop.test.ts @@ -28,7 +28,7 @@ mock.module("../../src/images/fulfill", () => ({ fulfillImageCall: async (): Promise => fulfillResult, })); -const { runWithImageBridge } = await import("../../src/images/loop"); +const { runWithImageBridge, clampImageMaxRounds, DEFAULT_MAX_ROUNDS, MAX_ROUNDS_HARD_LIMIT } = await import("../../src/images/loop"); // --- Mock adapter: yields canned events per iteration from a queue --- let streamQueue: AdapterEvent[][] = []; @@ -127,6 +127,159 @@ describe("runWithImageBridge", () => { // 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 — does not run unbounded iterations", async () => { + buildRequestCalls = 0; + // Provide only two streams; an unclamped 10000 would hang waiting for more. + streamQueue = [ + [...imageCallEvents], + [{ type: "text_delta" as const, text: "clamped final" }, { type: "done" as const }], + ]; + // Fill remaining slots so force-final after hard limit still has events if clamp failed. + for (let i = 0; i < 20; i++) { + streamQueue.push([{ type: "text_delta" as const, text: `extra ${i}` }, { 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 → at most 11 upstream requests (maxRounds+1), plus the first image round + // that triggers a loop: image call + up to 10 more. With only one image call then text, + // we stop early at 2. + expect(buildRequestCalls).toBe(2); + }); + + 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("429 key rotation rebuilds the adapter and retries the iteration", async () => { + let fetchCalls = 0; + let rotations = 0; + const rotatingAdapter: ProviderAdapter = { + name: "test", + 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 response = await runWithImageBridge({ + parsed: makeParsed(), + adapter: rotatingAdapter, + plan, + on429: () => { + rotations++; + return rotatingAdapter; + }, + }); + const sse = await response.text(); + expect(rotations).toBe(1); + expect(fetchCalls).toBe(2); + expect(sse).toContain("after rotate"); + }); }); // --------------------------------------------------------------------------- @@ -180,4 +333,51 @@ describe("runWithImageBridge — runTurn adapter", () => { 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 → 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/plan.test.ts b/tests/images/plan.test.ts index 85a536ee6f..777e0a9314 100644 --- a/tests/images/plan.test.ts +++ b/tests/images/plan.test.ts @@ -18,7 +18,7 @@ const { planImageBridge } = await import("../../src/images/plan"); function makeConfig( providers: Record>, - images?: { bridgeEnabled?: boolean; bridgeModel?: string }, + images?: { bridgeEnabled?: boolean; bridgeModel?: string; timeoutMs?: number }, ): OcxConfig { return { port: 0, @@ -106,6 +106,15 @@ describe("planImageBridge", () => { 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("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 }); diff --git a/tests/images/xai-client.test.ts b/tests/images/xai-client.test.ts index 6c8a6cdc70..76ed1a7a7f 100644 --- a/tests/images/xai-client.test.ts +++ b/tests/images/xai-client.test.ts @@ -76,6 +76,14 @@ describe("callXaiImages", () => { 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("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); From dd5c49b8679099fc542ecafe04fea4680c5f75eb Mon Sep 17 00:00:00 2001 From: Wibias <37517432+Wibias@users.noreply.github.com> Date: Mon, 27 Jul 2026 04:44:11 +0200 Subject: [PATCH 05/15] fix(images): address Codex and CodeRabbit feedback on #528 Keep runTurn adapters out of the unsupported web-search loop, honor xAI authMode, cap paid image calls, and close the remaining review gaps around usage, transport, tool-choice, and tests. --- .../src/content/docs/guides/image-bridge.md | 29 +++- src/images/artifacts.ts | 2 + src/images/fulfill.ts | 4 +- src/images/loop.ts | 136 ++++++++++++++--- src/images/plan.ts | 18 ++- src/images/types.ts | 3 + src/images/xai-client.ts | 4 +- src/responses/parser.ts | 1 + src/server/responses/core.ts | 143 +++++++++++------- tests/images/artifacts-ssrf.test.ts | 5 +- tests/images/fulfill.test.ts | 19 ++- tests/images/handler-activation.test.ts | 37 +++-- tests/images/loop.test.ts | 33 ++-- tests/images/plan.test.ts | 21 +++ 14 files changed, 334 insertions(+), 121 deletions(-) diff --git a/docs-site/src/content/docs/guides/image-bridge.md b/docs-site/src/content/docs/guides/image-bridge.md index 0f7102f3e2..5734d8246b 100644 --- a/docs-site/src/content/docs/guides/image-bridge.md +++ b/docs-site/src/content/docs/guides/image-bridge.md @@ -14,10 +14,21 @@ xAI Grok Imagine, so the model you're actually chatting with can still generate - **Enable the bridge** by setting `images.bridgeEnabled: true` in your config (it is off by default to avoid unexpected xAI charges — see [Configuration](#configuration) below). -- An xAI provider configured in settings with `baseUrl: "https://api.x.ai/v1"` and the - `openai-chat` adapter. +- 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). + bearer token) or `authMode: "key"` (a configured API key). See + [Providers](/reference/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.) @@ -31,7 +42,8 @@ Image Bridge options live under `images` in `~/.opencodex/config.json`. Bridging "images": { "bridgeEnabled": true, "bridgeModel": "grok-imagine-image-quality", - "maxRounds": 3 + "maxRounds": 3, + "timeoutMs": 60000 } } ``` @@ -40,7 +52,8 @@ Image Bridge options live under `images` in `~/.opencodex/config.json`. Bridging | --- | --- | --- | | `bridgeEnabled` | `false` | Master switch. Set `true` to enable bridging. Off by default to avoid unexpected xAI charges. | | `bridgeModel` | `grok-imagine-image-quality` | The xAI image model id to send prompts to. | -| `maxRounds` | `3` | Maximum number of image-generation loop iterations per turn. | +| `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 @@ -60,8 +73,10 @@ perspective, image generation works with any routed provider instead of silently ## Limitations - **Only xAI Grok Imagine is supported.** DALL-E and other image providers may be added later. -- **Web search takes priority.** If both web search and image generation are requested in the same - turn, the web-search bridge runs and image generation is skipped for that turn. +- **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 index 0ed1a2aed7..1523056e75 100644 --- a/src/images/artifacts.ts +++ b/src/images/artifacts.ts @@ -91,6 +91,8 @@ export async function downloadImageToArtifact( // SSRF protection: validate the provider-returned URL before fetching. // Require HTTPS strictly — plain HTTP and all other schemes (ftp, file, …) are rejected. + // DNS is checked before fetch; pinning the connected peer across a second resolution + // (rebinding) remains a recorded residual for this loopback proxy (destination-policy.ts). let parsedUrl: URL; try { parsedUrl = new URL(url); } catch { throw new Error("image URL is not valid"); } if (parsedUrl.protocol !== "https:") { diff --git a/src/images/fulfill.ts b/src/images/fulfill.ts index 907923e0ab..62d214d2f6 100644 --- a/src/images/fulfill.ts +++ b/src/images/fulfill.ts @@ -33,8 +33,8 @@ export async function fulfillImageCall( const n = typeof obj.n === "number" ? Math.max(1, Math.min(4, Math.floor(obj.n))) : 1; const imageUrl = typeof obj.image_url === "string" ? obj.image_url : typeof obj.image === "string" ? obj.image : undefined; - const size = typeof obj.size === "string" ? obj.size : undefined; - const quality = typeof obj.quality === "string" ? obj.quality : undefined; + const size = typeof obj.size === "string" ? obj.size : plan.defaultSize; + const quality = typeof obj.quality === "string" ? obj.quality : plan.defaultQuality; let result; try { diff --git a/src/images/loop.ts b/src/images/loop.ts index e48369ecc1..01011d98ea 100644 --- a/src/images/loop.ts +++ b/src/images/loop.ts @@ -10,7 +10,7 @@ * Removed vs web-search: no sidecar backend selection, no forced-answer nudge, no failed-query * dedup, no describeImages/structuredOutput, no recordSidecarOutcome. */ -import type { ProviderAdapter } from "../adapters/base"; +import type { AdapterRequest, ProviderAdapter } from "../adapters/base"; import { createAdapterEventQueue } from "../adapters/run-turn-queue"; import type { AdapterEvent, OcxMessage, OcxParsedRequest, OcxThinkingContent, OcxUsage } from "../types"; import { namespacedToolName } from "../types"; @@ -35,6 +35,8 @@ 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]. @@ -67,7 +69,12 @@ function scanEventsForImageCall(events: AdapterEvent[], toolNames: Set): let hasRealToolCall = false; let pending: { name: string; id: string; argsBuf: string; events: AdapterEvent[] } | null = null; const flushPending = (): void => { - if (pending && !toolNames.has(pending.name)) { + 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; } @@ -149,10 +156,18 @@ export interface ImageBridgeDeps { 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; /** @@ -173,6 +188,45 @@ export async function runWithImageBridge(deps: ImageBridgeDeps): Promise 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.usage) hiddenUsage = addUsage(hiddenUsage, e.usage); + } + }; const messages: OcxMessage[] = [...parsed.context.messages]; const allTools = parsed.context.tools ?? []; @@ -232,20 +286,25 @@ export async function runWithImageBridge(deps: ImageBridgeDeps): Promise ({ done: true as const, value })), - new Promise<{ done: false }>(resolve => { - const timer = setTimeout(() => resolve({ done: false }), 15_000); - void collectPromise.finally(() => clearTimeout(timer)); + const collectDeadline = clearableDeadline(stallTimeoutMs, signal); + let events: AdapterEvent[]; + try { + events = await Promise.race([ + collectPromise, + new Promise((_, reject) => { + const onAbort = (): void => { + reject(new LoopError(504, `runTurn inactivity timeout after ${stallTimeoutMs}ms during image-bridge`)); + }; + if (collectDeadline.signal.aborted) onAbort(); + else collectDeadline.signal.addEventListener("abort", onAbort, { once: true }); }), ]); - if (raced.done) events = raced.value; - else yield { type: "heartbeat" }; + } finally { + collectDeadline.clear(); } // Preserve Cursor conversation continuity across image-loop iterations. runTurn mutates @@ -271,18 +330,19 @@ export async function runWithImageBridge(deps: ImageBridgeDeps): 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: CONNECT_TIMEOUT_MS, + timeoutMs: connectTimeoutMs, returnRawErrors: true, stream: true, }) @@ -290,7 +350,7 @@ export async function runWithImageBridge(deps: ImageBridgeDeps): Promise { const h = new Headers(request.headers); if (!h.has("accept-encoding")) h.set("accept-encoding", "identity"); - return fetch(request.url, { + return fetchImpl(request.url, { method: request.method, headers: h, body: request.body, @@ -339,7 +399,7 @@ export async function runWithImageBridge(deps: ImageBridgeDeps): Promise 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") { + 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" }; - const result = await fulfillImageCall( - { id: call.id, name: call.name, arguments: call.args }, - plan, budget, signal, - ); + 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 { @@ -498,12 +583,15 @@ export async function runWithImageBridge(deps: ImageBridgeDeps): Promise { internalAbort.abort("client closed responses stream"); - }, undefined, + }, 2_000, { responseId: "", hideThinkingSummary: parsed.options.hideThinkingSummary, + stallTimeoutSec: deps.stallTimeoutSec, ...(deps.onFirstOutput ? { onFirstOutput: deps.onFirstOutput } : {}), - ...(deps.onUsage ? { onUsage: deps.onUsage } : {}), + ...(deps.onUsage ? { + onUsage: (usage: OcxUsage | undefined) => deps.onUsage?.(addUsage(hiddenUsage, usage)), + } : {}), }, ); return new Response(sse, { headers: SSE_HEADERS }); diff --git a/src/images/plan.ts b/src/images/plan.ts index eb8121b710..0f8e079b0c 100644 --- a/src/images/plan.ts +++ b/src/images/plan.ts @@ -23,10 +23,21 @@ export function findXaiProvider(config: OcxConfig): { name: string; provider: Oc } 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; - // Built-in OAuth token only for the canonical "xai" provider — never for custom-named configs. + // 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 { @@ -55,11 +66,16 @@ export async function planImageBridge( // 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; return { provider: found.provider, auth: { baseUrl: pinnedBaseUrl, token }, model: config.images?.bridgeModel ?? DEFAULT_MODEL, toolNames, + ...(hostedSize ? { defaultSize: hostedSize } : {}), + ...(hostedQuality ? { defaultQuality: hostedQuality } : {}), ...(typeof config.images?.timeoutMs === "number" && Number.isFinite(config.images.timeoutMs) && config.images.timeoutMs > 0 ? { timeoutMs: Math.floor(config.images.timeoutMs) } : {}), diff --git a/src/images/types.ts b/src/images/types.ts index ed15701dbb..d503828ab7 100644 --- a/src/images/types.ts +++ b/src/images/types.ts @@ -7,6 +7,9 @@ export interface ImageBridgePlan { 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 { diff --git a/src/images/xai-client.ts b/src/images/xai-client.ts index f66c7dfc42..fd26fabb06 100644 --- a/src/images/xai-client.ts +++ b/src/images/xai-client.ts @@ -100,7 +100,9 @@ export async function callXaiImages( }); if (!resp.ok) { - throw new Error("xAI images API returned " + resp.status); + 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 diff --git a/src/responses/parser.ts b/src/responses/parser.ts index c0a311c1ce..bc13d7fba2 100644 --- a/src/responses/parser.ts +++ b/src/responses/parser.ts @@ -125,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"; if (tool.type === "tool_search") return "tool_search"; return undefined; } diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index fb59b60fd2..0bc47e30a2 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -39,7 +39,7 @@ import { UnsupportedOAuthProviderError, } from "../../oauth"; import { buildWebSearchTool, planWebSearch, runWithWebSearch, shouldResolveOpenAiWebSearchSidecar } from "../../web-search"; -import { buildImageTool, planImageBridge, runWithImageBridge, clampImageMaxRounds } from "../../images"; +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 { @@ -1527,72 +1527,101 @@ export async function handleResponses( }); } - // Image bridge: check BEFORE the runTurn early-return so runTurn adapters (e.g. Cursor) - // also route through the bridge when image_generation is requested. The bridge loop - // internally supports both standard and runTurn adapter paths. - // Routed-compaction turns must NOT hit the bridge: compaction clears tools/_webSearch but + // 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; - if (imgPlan) { - // Web-search takes priority when both are eligible (design: image defers to web-search). - const wsPlan = planWebSearch(config, parsed, false, route.provider, route.modelId, openAiSidecar); - if (!wsPlan) { - // The bridge forces stream:true internally and returns SSE. Non-streaming requests can't be - // served — reject explicitly rather than returning SSE to a client expecting JSON. - if (!parsed.stream) { - return formatErrorResponse(400, "invalid_request_error", "image bridge requires stream=true"); - } - parsed.context.tools = [...(parsed.context.tools ?? []), buildImageTool()]; - const imgResponse = await runWithImageBridge({ - parsed, adapter, - plan: imgPlan, - forwardHeaders: selectedForwardHeaders, - onAttemptSend: () => noteAttemptSend(logCtx.activeAttempt, logCtx.usageLogInputTokens), - abortSignal: options.abortSignal, - maxRounds: clampImageMaxRounds(config.images?.maxRounds), - onUsage: usage => { - 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 } : {}), - }); - if (imgResponse.body) { - const imgTurnAc = new AbortController(); - return new Response(trackStreamLifetime(imgResponse.body, imgTurnAc), { - status: imgResponse.status, - headers: imgResponse.headers, + 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 => { + 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, }); - } - return imgResponse; + if (!rotated) return null; + route.provider = rotated; + return resolveAdapter( + resolveWireProtocolOverride(route.providerName, route.modelId, route.provider), + config.cacheRetention, + ); + }, + ...(options.onFirstOutput ? { onFirstOutput: options.onFirstOutput } : {}), + }); + if (imgResponse.body) { + const imgTurnAc = new AbortController(); + return new Response(trackStreamLifetime(imgResponse.body, imgTurnAc), { + status: imgResponse.status, + headers: imgResponse.headers, + }); } + return imgResponse; } // Web-search sidecar: Codex enabled web_search but this is a routed (non-OpenAI) model that can't // run it server-side. Expose web_search as a function tool and run searches via the gpt-mini sidecar // through the ChatGPT passthrough, looping until the model answers. Otherwise take the normal path. - // Placed BEFORE the runTurn early-return so dual-tool turns (web_search + image_generation) on - // runTurn adapters (e.g. Cursor) dispatch through the web-search sidecar instead of being swallowed - // by the runTurn branch — which would leave dual-tool turns with neither bridge. - const wsPlan = planWebSearch(config, parsed, false, route.provider, route.modelId, openAiSidecar); - if (wsPlan) { + // 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({ diff --git a/tests/images/artifacts-ssrf.test.ts b/tests/images/artifacts-ssrf.test.ts index 18039f584d..3bc67bfa81 100644 --- a/tests/images/artifacts-ssrf.test.ts +++ b/tests/images/artifacts-ssrf.test.ts @@ -94,11 +94,14 @@ describe("SSRF: downloadImageToArtifact scheme enforcement", () => { test(`3xx redirect response → rejects (redirect: 'error')`, async () => { lookupMock.mockResolvedValue([{ address: "93.184.216.34", family: 4 }]); const originalFetch = globalThis.fetch; + let seenRedirect: RequestRedirect | undefined; try { - globalThis.fetch = (async (_input: RequestInfo | URL, _init?: RequestInit) => { + globalThis.fetch = (async (_input: RequestInfo | URL, init?: RequestInit) => { + seenRedirect = init?.redirect; return new Response("", { status: 301, headers: { Location: "https://evil.example/redirect" } }); }) as typeof fetch; await expect(downloadImageToArtifact("https://public-host/redirect-img")).rejects.toThrow(); + expect(seenRedirect).toBe("error"); } finally { globalThis.fetch = originalFetch; lookupMock.mockClear(); diff --git a/tests/images/fulfill.test.ts b/tests/images/fulfill.test.ts index 67e06b62b1..f37fc4a245 100644 --- a/tests/images/fulfill.test.ts +++ b/tests/images/fulfill.test.ts @@ -18,8 +18,14 @@ 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; mock.module("../../src/images/xai-client", () => ({ - callXaiImages: async (req: XaiImageRequest) => { xaiCalls.push(req); if (xaiError) throw xaiError; return xaiResult; }, + 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 }), @@ -40,6 +46,7 @@ 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`; @@ -57,6 +64,16 @@ describe("fulfillImageCall", () => { 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 }); diff --git a/tests/images/handler-activation.test.ts b/tests/images/handler-activation.test.ts index 8b09aaf3e5..1af527de9d 100644 --- a/tests/images/handler-activation.test.ts +++ b/tests/images/handler-activation.test.ts @@ -67,19 +67,16 @@ mock.module("../../src/server/adapter-resolve", () => ({ })); // --- Stub the image bridge runner: detect activation without hitting the real loop --- +// --- Stub the image bridge runner: detect activation without hitting the real loop --- +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" }, }); }, - clampImageMaxRounds: (value: unknown) => { - if (typeof value !== "number" || !Number.isFinite(value)) return 3; - return Math.max(0, Math.min(10, Math.floor(value))); - }, - DEFAULT_MAX_ROUNDS: 3, - MAX_ROUNDS_HARD_LIMIT: 10, })); // --- Stub the web-search planner + runner: control eligibility and detect activation --- @@ -115,7 +112,7 @@ function makeConfig(): OcxConfig { defaultProvider: "fixture", providers: { fixture: { adapter: "openai-chat", baseUrl: "https://fixture.test/v1", authMode: "key", apiKey: "fixture-key" }, - xai: { adapter: "openai", baseUrl: "https://api.x.ai", apiKey: "xai-test-token" }, + xai: { adapter: "openai-chat", baseUrl: "https://api.x.ai/v1", apiKey: "xai-test-token" }, }, images: { bridgeEnabled: true }, } as OcxConfig; @@ -186,17 +183,31 @@ describe("image bridge dispatch priority (handler activation)", () => { expect(res.headers.get("content-type")).toBe("text/event-stream"); }); - test("dual-tool on a runTurn adapter → web-search path wins, runTurn does not eat the request (#424)", async () => { + 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" }]); - // Web-search sidecar must handle the turn. - expect(webSearchRun).toBe(true); - // Image bridge must not activate (design: image defers to web-search). - expect(imageBridgeRun).toBe(false); - // runTurn must NOT be reached — the web-search dispatch now runs before the runTurn early-return. + // Web-search sidecar cannot drive runTurn adapters — skip it. + expect(webSearchRun).toBe(false); + // Image bridge supports runTurn, so it handles the dual-tool turn instead. + 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 { diff --git a/tests/images/loop.test.ts b/tests/images/loop.test.ts index 6462f9656a..d21f266aa4 100644 --- a/tests/images/loop.test.ts +++ b/tests/images/loop.test.ts @@ -1,4 +1,4 @@ -import { afterAll, beforeAll, describe, expect, mock, test } from "bun:test"; +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"; @@ -33,6 +33,16 @@ const { runWithImageBridge, clampImageMaxRounds, DEFAULT_MAX_ROUNDS, MAX_ROUNDS_ // --- 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; +}); + const mockAdapter: ProviderAdapter = { name: "test", buildRequest: async () => { buildRequestCalls++; return { url: "https://test/v1/chat", method: "POST", headers: {}, body: "{}" }; }, @@ -136,26 +146,21 @@ describe("runWithImageBridge", () => { expect(clampImageMaxRounds(undefined)).toBe(DEFAULT_MAX_ROUNDS); }); - test("maxRounds: 10000 is clamped — does not run unbounded iterations", async () => { + test("maxRounds: 10000 is clamped — hits hard limit when every round calls image_gen", async () => { buildRequestCalls = 0; - // Provide only two streams; an unclamped 10000 would hang waiting for more. - streamQueue = [ - [...imageCallEvents], - [{ type: "text_delta" as const, text: "clamped final" }, { type: "done" as const }], - ]; - // Fill remaining slots so force-final after hard limit still has events if clamp failed. - for (let i = 0; i < 20; i++) { - streamQueue.push([{ type: "text_delta" as const, text: `extra ${i}` }, { type: "done" as const }]); + 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 → at most 11 upstream requests (maxRounds+1), plus the first image round - // that triggers a loop: image call + up to 10 more. With only one image call then text, - // we stop early at 2. - expect(buildRequestCalls).toBe(2); + // 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 () => { diff --git a/tests/images/plan.test.ts b/tests/images/plan.test.ts index 777e0a9314..39df98106b 100644 --- a/tests/images/plan.test.ts +++ b/tests/images/plan.test.ts @@ -149,4 +149,25 @@ describe("planImageBridge", () => { 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; + }); }); From 83987e834a4cce8593b6a1318c60d806e165dbde Mon Sep 17 00:00:00 2001 From: Wibias <37517432+Wibias@users.noreply.github.com> Date: Mon, 27 Jul 2026 05:03:35 +0200 Subject: [PATCH 06/15] fix(images): abort stalled runTurn and merge incomplete usage Cancel the in-flight runTurn when the image-bridge collect deadline fires, fire onAttemptSend at dispatch time, and fold accumulated usage into incomplete terminals as well as done. --- src/images/loop.ts | 13 +++++++++---- tests/images/loop.test.ts | 37 +++++++++++++++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 4 deletions(-) diff --git a/src/images/loop.ts b/src/images/loop.ts index 01011d98ea..1a1e465d5e 100644 --- a/src/images/loop.ts +++ b/src/images/loop.ts @@ -224,7 +224,9 @@ export async function runWithImageBridge(deps: ImageBridgeDeps): Promise { for (const e of events) { - if (e.type === "done" && e.usage) hiddenUsage = addUsage(hiddenUsage, e.usage); + if ((e.type === "done" || e.type === "incomplete") && e.usage) { + hiddenUsage = addUsage(hiddenUsage, e.usage); + } } }; @@ -274,6 +276,8 @@ export async function runWithImageBridge(deps: ImageBridgeDeps): Promise internalAbort.abort("runTurn backlog exceeded"), }); + // Attempt telemetry must fire at dispatch time (parity with fetchOnce), not after collect. + deps.onAttemptSend?.(); void adapter .runTurn( iterParsed, @@ -297,6 +301,9 @@ export async function runWithImageBridge(deps: ImageBridgeDeps): Promise((_, reject) => { const onAbort = (): void => { + // Cancel the fire-and-forget runTurn so a stalled Cursor session does not keep + // running after the bridge has already failed the iteration with 504. + internalAbort.abort(`runTurn inactivity timeout after ${stallTimeoutMs}ms`); reject(new LoopError(504, `runTurn inactivity timeout after ${stallTimeoutMs}ms during image-bridge`)); }; if (collectDeadline.signal.aborted) onAbort(); @@ -313,8 +320,6 @@ export async function runWithImageBridge(deps: ImageBridgeDeps): Promise e.type === "error"); if (errorEvent && errorEvent.type === "error") { @@ -496,7 +501,7 @@ export async function runWithImageBridge(deps: ImageBridgeDeps): Promise= 0; i--) { const e = split.passthrough[i]; - if (e?.type === "done") { + if (e?.type === "done" || e?.type === "incomplete") { split.passthrough[i] = { ...e, usage: addUsage(hiddenUsage, e.usage) }; break; } diff --git a/tests/images/loop.test.ts b/tests/images/loop.test.ts index d21f266aa4..973892b09f 100644 --- a/tests/images/loop.test.ts +++ b/tests/images/loop.test.ts @@ -359,6 +359,43 @@ describe("runWithImageBridge — runTurn adapter", () => { 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 → preserves _cursorConversationId across iterations", async () => { const seenIds: Array = []; const cursorAdapter: ProviderAdapter = { From 875567164cb1e6da9a934e65128bea917fc770f8 Mon Sep 17 00:00:00 2001 From: Wibias <37517432+Wibias@users.noreply.github.com> Date: Mon, 27 Jul 2026 16:55:48 +0200 Subject: [PATCH 07/15] fix(images): pin validated DNS for artifact downloads against rebinding Resolve once via resolvePublicAddresses, then connect with a custom https lookup that keeps SNI/Host on the original hostname so a later private answer cannot retarget the peer. --- src/images/artifacts.ts | 109 ++++++++++++++++++++++++++-- src/lib/destination-policy.ts | 43 ++++++++--- tests/images/artifacts-ssrf.test.ts | 98 +++++++++++++++++++------ 3 files changed, 212 insertions(+), 38 deletions(-) diff --git a/src/images/artifacts.ts b/src/images/artifacts.ts index 1523056e75..b465b0bb8b 100644 --- a/src/images/artifacts.ts +++ b/src/images/artifacts.ts @@ -1,7 +1,9 @@ import { mkdir, writeFile } from "node:fs/promises"; +import https from "node:https"; +import type { IncomingMessage, RequestOptions } from "node:http"; import { join } from "node:path"; import { getConfigDir } from "../config"; -import { assessUrlDestination, assertUrlResolvesPublic } from "../lib/destination-policy"; +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; @@ -15,6 +17,15 @@ 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 }; } @@ -78,10 +89,95 @@ export async function materializeInlineImage( 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. + */ +export function pinnedHttpsGet( + url: string, + pinned: PinnedAddress, + signal?: AbortSignal, +): Promise { + const parsed = new URL(url); + if (parsed.protocol !== "https:") { + throw new Error(`image URL must use HTTPS, got ${parsed.protocol}`); + } + + return new Promise((resolve, reject) => { + if (signal?.aborted) { + reject(signal.reason instanceof Error ? signal.reason : new Error("aborted")); + return; + } + + const options: 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 }, + lookup(_hostname, lookupOptions, callback) { + const cb = typeof lookupOptions === "function" + ? lookupOptions + : callback; + if (!cb) return; + // Pin the validated peer — do not call dns.lookup again. + cb(null, pinned.address, pinned.family as 4 | 6); + }, + }; + + const req = https.request(options, (res: IncomingMessage) => { + const status = res.statusCode ?? 0; + // Match fetch({ redirect: "error" }): never follow 3xx. + if (status >= 300 && status < 400) { + res.resume(); + reject(new Error("image download failed: " + status)); + return; + } + const chunks: Buffer[] = []; + res.on("data", (chunk: Buffer | string) => { + chunks.push(typeof chunk === "string" ? Buffer.from(chunk) : chunk); + }); + res.on("end", () => { + const body = Buffer.concat(chunks); + const headers = new Headers(); + for (const [key, value] of Object.entries(res.headers)) { + if (value === undefined) continue; + if (Array.isArray(value)) { + for (const item of value) headers.append(key, item); + } else { + headers.set(key, value); + } + } + resolve(new Response(body, { status, headers })); + }); + res.on("error", reject); + }); + + const onAbort = () => { + req.destroy(signal?.reason instanceof Error ? signal.reason : new Error("aborted")); + }; + signal?.addEventListener("abort", onAbort, { once: true }); + req.on("error", (err) => { + signal?.removeEventListener("abort", onAbort); + reject(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); @@ -91,8 +187,8 @@ export async function downloadImageToArtifact( // SSRF protection: validate the provider-returned URL before fetching. // Require HTTPS strictly — plain HTTP and all other schemes (ftp, file, …) are rejected. - // DNS is checked before fetch; pinning the connected peer across a second resolution - // (rebinding) remains a recorded residual for this loopback proxy (destination-policy.ts). + // 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:") { @@ -103,9 +199,10 @@ export async function downloadImageToArtifact( if (assessment && assessment.kind !== "public" && assessment.kind !== "hostname") { throw new Error(`image URL targets ${assessment.detail}`); } - // DNS check: resolve hostname and reject if it points at private/internal space. - await assertUrlResolvesPublic(url); - const resp = await fetch(url, { signal, redirect: "error" }); + 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) throw new Error("image download failed: " + resp.status); // Stream the body with a hard byte cap so a missing/lying Content-Length or a diff --git a/src/lib/destination-policy.ts b/src/lib/destination-policy.ts index 3ed0bf5330..97ae24c7c7 100644 --- a/src/lib/destination-policy.ts +++ b/src/lib/destination-policy.ts @@ -205,10 +205,14 @@ export function assessUrlDestination(url: string): UrlDestinationAssessment | nu /** * Async DNS-resolved URL safety check. Resolves A/AAAA records and rejects * if any address is loopback, private, link-local, unspecified, or metadata. - * Throws on unsafe destination; returns void on safe/public destination. + * 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 assertUrlResolvesPublic(url: string): Promise { +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); @@ -220,9 +224,15 @@ export async function assertUrlResolvesPublic(url: string): Promise { if (literalAssessment && literalAssessment.kind !== "public" && literalAssessment.kind !== "hostname") { throw new Error(`image URL targets ${literalAssessment.detail}`); } - // For literal IPs and localhost, the sync path already classified them. - if (isIP(hostname) !== 0 || hostname === "localhost" || hostname.endsWith(".localhost")) return; - let addresses: { address: string }[]; + // 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 }] }; + } + if (hostname === "localhost" || hostname.endsWith(".localhost")) { + throw new Error(`image URL targets localhost destination`); + } + let addresses: { address: string; family: number }[]; try { addresses = await lookup(hostname, { all: true, verbatim: true }); } catch { @@ -230,10 +240,25 @@ export async function assertUrlResolvesPublic(url: string): Promise { // this is a runtime fetch to an untrusted URL, so be conservative). throw new Error(`image URL hostname ${hostname} could not be resolved`); } - for (const { address } of addresses) { - const ipKind = isIP(address); + 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") continue; - throw new Error(`image URL hostname ${hostname} resolves to ${assessment.detail} (${address})`); + 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/tests/images/artifacts-ssrf.test.ts b/tests/images/artifacts-ssrf.test.ts index 3bc67bfa81..b69b9e1f69 100644 --- a/tests/images/artifacts-ssrf.test.ts +++ b/tests/images/artifacts-ssrf.test.ts @@ -5,9 +5,11 @@ import { describe, expect, mock, test } from "bun:test"; const lookupMock = mock(async (_hostname: string, _opts: unknown): Promise<{ address: string; family: number }[]> => []); mock.module("node:dns/promises", () => ({ lookup: lookupMock })); -const { assessUrlDestination, assertUrlResolvesPublic } = await import("../../src/lib/destination-policy"); +const { 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"); @@ -66,6 +68,31 @@ describe("SSRF: assertUrlResolvesPublic", () => { }); }); +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(); + } + }); +}); + describe("SSRF: downloadImageToArtifact scheme enforcement", () => { test("http:// → rejects (non-HTTPS)", async () => { await expect(downloadImageToArtifact("http://public-host/path")).rejects.toThrow(/HTTPS/); @@ -93,39 +120,64 @@ describe("SSRF: downloadImageToArtifact scheme enforcement", () => { test(`3xx redirect response → rejects (redirect: 'error')`, async () => { lookupMock.mockResolvedValue([{ address: "93.184.216.34", family: 4 }]); - const originalFetch = globalThis.fetch; - let seenRedirect: RequestRedirect | undefined; try { - globalThis.fetch = (async (_input: RequestInfo | URL, init?: RequestInit) => { - seenRedirect = init?.redirect; - return new Response("", { status: 301, headers: { Location: "https://evil.example/redirect" } }); - }) as typeof fetch; - await expect(downloadImageToArtifact("https://public-host/redirect-img")).rejects.toThrow(); - expect(seenRedirect).toBe("error"); + await expect(downloadImageToArtifact("https://public-host/redirect-img", undefined, undefined, { + pinnedDownload: async () => new Response("", { + status: 301, + headers: { Location: "https://evil.example/redirect" }, + }), + })).rejects.toThrow(); } finally { - globalThis.fetch = originalFetch; lookupMock.mockClear(); } }); - test("https:// public host → succeeds with mocked fetch", async () => { - // Stub DNS so public-host resolves to a public address. + test("https:// public host → succeeds with pinned download", async () => { lookupMock.mockResolvedValue([{ address: "93.184.216.34", family: 4 }]); - const originalFetch = globalThis.fetch; 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("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. + 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 { - globalThis.fetch = (async (input: RequestInfo | URL, _init?: RequestInit) => { - const raw = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; - if (!raw.startsWith("https://")) throw new Error("fetch must only be called over HTTPS"); - // Minimal PNG signature so guessExtFromMagic returns "png". - const pngBytes = new Uint8Array([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]); - return new Response(pngBytes, { status: 200 }); - }) as typeof fetch; - - downloadedPath = await downloadImageToArtifact("https://public-host/valid-image"); + 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 { - globalThis.fetch = originalFetch; lookupMock.mockClear(); if (downloadedPath) await rm(downloadedPath).catch(() => {}); } From 4b772ca846a5de50301a2051638d6ad32d694723 Mon Sep 17 00:00:00 2001 From: Wibias <37517432+Wibias@users.noreply.github.com> Date: Mon, 27 Jul 2026 17:21:11 +0200 Subject: [PATCH 08/15] fix(images): stream pinned downloads with byte cap and idle timeout Address Ingwannu and CodeRabbit feedback on #528: pinnedHttpsGet returns a streaming body, enforces MAX_DOWNLOAD_BYTES mid-stream, honors lookup all:true, and times out idle peers. Drop unreachable localhost check after assessDestination. Add transport regressions. --- src/images/artifacts.ts | 117 +++++++++++++----- src/lib/destination-policy.ts | 3 - tests/images/artifacts-ssrf.test.ts | 5 +- tests/images/pinned-https-get.test.ts | 167 ++++++++++++++++++++++++++ 4 files changed, 260 insertions(+), 32 deletions(-) create mode 100644 tests/images/pinned-https-get.test.ts diff --git a/src/images/artifacts.ts b/src/images/artifacts.ts index b465b0bb8b..102e5565be 100644 --- a/src/images/artifacts.ts +++ b/src/images/artifacts.ts @@ -1,13 +1,17 @@ import { mkdir, writeFile } from "node:fs/promises"; +import type { IncomingMessage } from "node:http"; import https from "node:https"; -import type { IncomingMessage, RequestOptions } from "node:http"; +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; -const MAX_DOWNLOAD_BYTES = 50 * 1024 * 1024; // 50 MiB +/** 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. @@ -93,16 +97,26 @@ export async function materializeInlineImage( * 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) { @@ -110,7 +124,15 @@ export function pinnedHttpsGet( return; } - const options: RequestOptions & { servername?: string } = { + 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, @@ -118,51 +140,92 @@ export function pinnedHttpsGet( path: `${parsed.pathname}${parsed.search}`, method: "GET", headers: { Host: parsed.host }, + rejectUnauthorized: options?.rejectUnauthorized, lookup(_hostname, lookupOptions, callback) { - const cb = typeof lookupOptions === "function" - ? 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. - cb(null, pinned.address, pinned.family as 4 | 6); + // 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(options, (res: IncomingMessage) => { + const req = https.request(optionsHttps, (res: IncomingMessage) => { const status = res.statusCode ?? 0; // Match fetch({ redirect: "error" }): never follow 3xx. if (status >= 300 && status < 400) { res.resume(); - reject(new Error("image download failed: " + status)); + fail(new Error("image download failed: " + status)); return; } - const chunks: Buffer[] = []; - res.on("data", (chunk: Buffer | string) => { - chunks.push(typeof chunk === "string" ? Buffer.from(chunk) : chunk); - }); - res.on("end", () => { - const body = Buffer.concat(chunks); - const headers = new Headers(); - for (const [key, value] of Object.entries(res.headers)) { - if (value === undefined) continue; - if (Array.isArray(value)) { - for (const item of value) headers.append(key, item); - } else { - headers.set(key, value); - } + 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)); } - resolve(new Response(body, { status, headers })); + } + + 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(); + }, }); - res.on("error", reject); + + if (settled) return; + settled = true; + resolve(new Response(stream, { status, headers })); }); const onAbort = () => { - req.destroy(signal?.reason instanceof Error ? signal.reason : new Error("aborted")); + 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); - reject(err); + fail(err); }); req.on("close", () => signal?.removeEventListener("abort", onAbort)); req.end(); diff --git a/src/lib/destination-policy.ts b/src/lib/destination-policy.ts index 97ae24c7c7..501475dd8f 100644 --- a/src/lib/destination-policy.ts +++ b/src/lib/destination-policy.ts @@ -229,9 +229,6 @@ export async function resolvePublicAddresses(url: string): Promise<{ if (literalKind !== 0) { return { hostname, addresses: [{ address: hostname, family: literalKind }] }; } - if (hostname === "localhost" || hostname.endsWith(".localhost")) { - throw new Error(`image URL targets localhost destination`); - } let addresses: { address: string; family: number }[]; try { addresses = await lookup(hostname, { all: true, verbatim: true }); diff --git a/tests/images/artifacts-ssrf.test.ts b/tests/images/artifacts-ssrf.test.ts index b69b9e1f69..9cb4298cf3 100644 --- a/tests/images/artifacts-ssrf.test.ts +++ b/tests/images/artifacts-ssrf.test.ts @@ -118,7 +118,7 @@ describe("SSRF: downloadImageToArtifact scheme enforcement", () => { await expect(downloadImageToArtifact("https://[::ffff:7f00:1]/image.png")).rejects.toThrow(); }); - test(`3xx redirect response → rejects (redirect: 'error')`, async () => { + 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, { @@ -126,7 +126,7 @@ describe("SSRF: downloadImageToArtifact scheme enforcement", () => { status: 301, headers: { Location: "https://evil.example/redirect" }, }), - })).rejects.toThrow(); + })).rejects.toThrow(/301|failed/); } finally { lookupMock.mockClear(); } @@ -154,6 +154,7 @@ describe("SSRF: downloadImageToArtifact scheme enforcement", () => { 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; diff --git a/tests/images/pinned-https-get.test.ts b/tests/images/pinned-https-get.test.ts new file mode 100644 index 0000000000..d50123be04 --- /dev/null +++ b/tests/images/pinned-https-get.test.ts @@ -0,0 +1,167 @@ +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("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/); + }); +}); From d4478ec85365c4759482ef0a2e84d2b6863cbacc Mon Sep 17 00:00:00 2001 From: Wibias <37517432+Wibias@users.noreply.github.com> Date: Mon, 27 Jul 2026 17:35:30 +0200 Subject: [PATCH 09/15] fix(images): destroy pinned HTTPS response on non-2xx status Reject failed downloads before attaching a streaming body so unread 4xx/5xx payloads cannot keep the socket alive, and cancel custom-seam bodies on !ok. --- src/images/artifacts.ts | 15 ++++-- tests/images/pinned-https-get.test.ts | 67 +++++++++++++++++++++++++++ 2 files changed, 78 insertions(+), 4 deletions(-) diff --git a/src/images/artifacts.ts b/src/images/artifacts.ts index 102e5565be..30775bf94b 100644 --- a/src/images/artifacts.ts +++ b/src/images/artifacts.ts @@ -164,9 +164,11 @@ export function pinnedHttpsGet( const req = https.request(optionsHttps, (res: IncomingMessage) => { const status = res.statusCode ?? 0; - // Match fetch({ redirect: "error" }): never follow 3xx. - if (status >= 300 && status < 400) { - res.resume(); + // 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; } @@ -266,7 +268,12 @@ export async function downloadImageToArtifact( const pinned = pickPinnedAddress(resolved.addresses); const download = options?.pinnedDownload ?? pinnedHttpsGet; const resp = await download(url, pinned, signal); - if (!resp.ok) throw new Error("image download failed: " + resp.status); + 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. diff --git a/tests/images/pinned-https-get.test.ts b/tests/images/pinned-https-get.test.ts index d50123be04..a3038dcb79 100644 --- a/tests/images/pinned-https-get.test.ts +++ b/tests/images/pinned-https-get.test.ts @@ -139,6 +139,73 @@ describe("pinnedHttpsGet transport", () => { expect(received).toBeLessThanOrEqual(maxBytes + small.byteLength); }); + 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, From ad5d04ce65cfe5f1a2ae2ca0bda836e2cf044e75 Mon Sep 17 00:00:00 2001 From: Wibias <37517432+Wibias@users.noreply.github.com> Date: Mon, 27 Jul 2026 17:38:17 +0200 Subject: [PATCH 10/15] fix(docs): point image-bridge Providers link at guides/providers CodeRabbit: /reference/providers/ is not a registered Starlight page. --- docs-site/src/content/docs/guides/image-bridge.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs-site/src/content/docs/guides/image-bridge.md b/docs-site/src/content/docs/guides/image-bridge.md index 5734d8246b..6aefbee6d6 100644 --- a/docs-site/src/content/docs/guides/image-bridge.md +++ b/docs-site/src/content/docs/guides/image-bridge.md @@ -28,7 +28,7 @@ xAI Grok Imagine, so the model you're actually chatting with can still generate - Authentication via `authMode: "oauth"` (`ocx login xai` — uses a stored, auto-refreshed bearer token) or `authMode: "key"` (a configured API key). See - [Providers](/reference/providers/) for the shared auth modes. + [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.) From de19c7e3b941090abc50b3da24729736d104a516 Mon Sep 17 00:00:00 2001 From: Wibias <37517432+Wibias@users.noreply.github.com> Date: Mon, 27 Jul 2026 17:50:41 +0200 Subject: [PATCH 11/15] fix(images): address Codex review on #528 Reject non-global IPv6 peers, stop double-counting hidden usage, reset runTurn idle stalls on progress, reject empty/non-image downloads, backfill Cursor conversation ids from the image loop, and clarify Responses vs /images/generations activation in docs. --- .../content/docs/guides/codex-integration.md | 6 +- .../src/content/docs/guides/image-bridge.md | 8 ++- src/images/artifacts.ts | 18 ++++-- src/images/loop.ts | 49 +++++++++-------- src/lib/destination-policy.ts | 16 +++++- src/server/responses/core.ts | 6 ++ tests/destination-policy-resolved.test.ts | 17 ++++++ tests/images/artifacts-ssrf.test.ts | 55 +++++++++++++++++++ tests/images/loop.test.ts | 47 ++++++++++++++++ 9 files changed, 191 insertions(+), 31 deletions(-) diff --git a/docs-site/src/content/docs/guides/codex-integration.md b/docs-site/src/content/docs/guides/codex-integration.md index cd4fe810af..44eaec6105 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 index 6aefbee6d6..748b3e22fa 100644 --- a/docs-site/src/content/docs/guides/image-bridge.md +++ b/docs-site/src/content/docs/guides/image-bridge.md @@ -57,7 +57,13 @@ Image Bridge options live under `images` in `~/.opencodex/config.json`. Bridging ## How It Works -1. When Codex sends a request with `image_generation` in the tools array, OpenCodex detects it +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. diff --git a/src/images/artifacts.ts b/src/images/artifacts.ts index 30775bf94b..f6ba898d87 100644 --- a/src/images/artifacts.ts +++ b/src/images/artifacts.ts @@ -53,13 +53,19 @@ function timestampPrefix(): string { ].join(""); } -export function guessExtFromMagic(bytes: Uint8Array): string { +/** 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 "png"; + return null; +} + +export function guessExtFromMagic(bytes: Uint8Array): string { + return sniffImageExtension(bytes) ?? "png"; } export async function materializeInlineImage( @@ -87,7 +93,8 @@ export async function materializeInlineImage( if (budget) budget.spent += buf.length; // Sniff actual format from decoded bytes rather than trusting the declared mimeType. - const ext = guessExtFromMagic(buf); + 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; @@ -300,11 +307,14 @@ export async function downloadImageToArtifact( 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 ext = guessExtFromMagic(bytes); const dir = getArtifactsDir(); await mkdir(dir, { recursive: true, mode: 0o700 }); if (budget) budget.spent += bytes.length; diff --git a/src/images/loop.ts b/src/images/loop.ts index 1a1e465d5e..d96b203b48 100644 --- a/src/images/loop.ts +++ b/src/images/loop.ts @@ -15,7 +15,7 @@ import { createAdapterEventQueue } from "../adapters/run-turn-queue"; import type { AdapterEvent, OcxMessage, OcxParsedRequest, OcxThinkingContent, OcxUsage } from "../types"; import { namespacedToolName } from "../types"; import { bridgeToResponsesSSE } from "../bridge"; -import { clearableDeadline } from "../lib/abort"; +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"; @@ -290,28 +290,31 @@ export async function runWithImageBridge(deps: ImageBridgeDeps): Promise { + timedOut = true; + // Cancel the fire-and-forget runTurn so a stalled Cursor session does not keep + // running after the bridge has already failed the iteration with 504. + internalAbort.abort(`runTurn inactivity timeout after ${stallTimeoutMs}ms`); + }); + const events: AdapterEvent[] = []; try { - events = await Promise.race([ - collectPromise, - new Promise((_, reject) => { - const onAbort = (): void => { - // Cancel the fire-and-forget runTurn so a stalled Cursor session does not keep - // running after the bridge has already failed the iteration with 504. - internalAbort.abort(`runTurn inactivity timeout after ${stallTimeoutMs}ms`); - reject(new LoopError(504, `runTurn inactivity timeout after ${stallTimeoutMs}ms during image-bridge`)); - }; - if (collectDeadline.signal.aborted) onAbort(); - else collectDeadline.signal.addEventListener("abort", onAbort, { once: true }); - }), - ]); + idle.reset(); + for await (const event of queue.stream()) { + if (timedOut) break; + idle.reset(); + events.push(event); + } } finally { - collectDeadline.clear(); + 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 @@ -595,7 +598,9 @@ export async function runWithImageBridge(deps: ImageBridgeDeps): Promise deps.onUsage?.(addUsage(hiddenUsage, usage)), + // 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), } : {}), }, ); diff --git a/src/lib/destination-policy.ts b/src/lib/destination-policy.ts index 501475dd8f..357dd7ff90 100644 --- a/src/lib/destination-policy.ts +++ b/src/lib/destination-policy.ts @@ -92,10 +92,20 @@ function classifyIpv6(hostname: string): DestinationAssessment { 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 { diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index 0bc47e30a2..b9799c74e3 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -1584,6 +1584,12 @@ export async function handleResponses( 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; diff --git a/tests/destination-policy-resolved.test.ts b/tests/destination-policy-resolved.test.ts index e700c3a572..d47df9416c 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 index 9cb4298cf3..7036333984 100644 --- a/tests/images/artifacts-ssrf.test.ts +++ b/tests/images/artifacts-ssrf.test.ts @@ -32,6 +32,21 @@ describe("SSRF: assessUrlDestination", () => { 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); @@ -91,6 +106,24 @@ describe("SSRF: resolvePublicAddresses", () => { 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", () => { @@ -151,6 +184,28 @@ describe("SSRF: downloadImageToArtifact scheme enforcement", () => { } }); + 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. diff --git a/tests/images/loop.test.ts b/tests/images/loop.test.ts index 973892b09f..18df1468cb 100644 --- a/tests/images/loop.test.ts +++ b/tests/images/loop.test.ts @@ -254,6 +254,29 @@ describe("runWithImageBridge", () => { 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; @@ -396,6 +419,30 @@ describe("runWithImageBridge — runTurn adapter", () => { expect(sse).toContain("runTurn inactivity timeout"); }); + 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 = { From 0a4c78187605a35c65d3646de8f8127d7e700dbf Mon Sep 17 00:00:00 2001 From: Wibias <37517432+Wibias@users.noreply.github.com> Date: Mon, 27 Jul 2026 18:05:02 +0200 Subject: [PATCH 12/15] fix(images): unblock runTurn idle timeout without waiting on the adapter On idle expiry abort the runTurn signal and close the event queue so the consumer finishes even when runTurn ignores cancellation and never settles. Add a never-settling regression. --- src/images/loop.ts | 8 ++++++-- tests/images/loop.test.ts | 26 ++++++++++++++++++++++++++ 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/src/images/loop.ts b/src/images/loop.ts index d96b203b48..b8fbb3be5b 100644 --- a/src/images/loop.ts +++ b/src/images/loop.ts @@ -295,12 +295,16 @@ export async function runWithImageBridge(deps: ImageBridgeDeps): Promise { timedOut = true; - // Cancel the fire-and-forget runTurn so a stalled Cursor session does not keep - // running after the bridge has already failed the iteration with 504. + // 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 { diff --git a/tests/images/loop.test.ts b/tests/images/loop.test.ts index 18df1468cb..acaebfd82c 100644 --- a/tests/images/loop.test.ts +++ b/tests/images/loop.test.ts @@ -419,6 +419,32 @@ describe("runWithImageBridge — runTurn adapter", () => { 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 = { From 34c6d852d6632ca19666faf25ecd3e78e7e80404 Mon Sep 17 00:00:00 2001 From: Wibias <37517432+Wibias@users.noreply.github.com> Date: Mon, 27 Jul 2026 21:01:20 +0200 Subject: [PATCH 13/15] fix(test): raise Windows CI timeout for injected cleanup rollback cases MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rebase onto dev and extend the storage cleanup harness timeout for injected satellite rollback tests that measure 6–13s on Windows runners. --- tests/storage-cleanup.test.ts | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/storage-cleanup.test.ts b/tests/storage-cleanup.test.ts index 85bbc0d6ea..b825842415 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", () => { From 11b740f93862b72827a6c542d7bd0f9eebabcaa7 Mon Sep 17 00:00:00 2001 From: Wibias <37517432+Wibias@users.noreply.github.com> Date: Mon, 27 Jul 2026 22:23:59 +0200 Subject: [PATCH 14/15] fix(tests): stabilize blocked policy worker test on Windows CI Wait for worker spawn and holdAfterLoadMs before concurrent PUT so slow Windows CI worker startup does not load a disabled policy snapshot. --- tests/api-storage-policy.test.ts | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/tests/api-storage-policy.test.ts b/tests/api-storage-policy.test.ts index 499b078ff5..8324efdbf8 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: 800 }); + // 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(120); + 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); From 2db8b5721a4d8a0a5c469f9a64931b0c1d216f18 Mon Sep 17 00:00:00 2001 From: Wibias <37517432+Wibias@users.noreply.github.com> Date: Mon, 27 Jul 2026 23:03:47 +0200 Subject: [PATCH 15/15] fix(images): address CodeRabbit review on image bridge P2 --- src/images/loop.ts | 5 +- src/images/plan.ts | 12 +- src/images/synthetic-tool.ts | 26 +++- src/images/xai-client.ts | 3 +- src/responses/parser.ts | 4 +- src/server/responses/core.ts | 7 + tests/images/loop.test.ts | 58 ++++--- tests/images/pinned-https-get.test.ts | 87 +++++++++++ tests/images/plan.test.ts | 39 +++-- tests/images/synthetic-tool.test.ts | 20 ++- tests/images/xai-client.test.ts | 22 +++ .../{fulfill.test.ts => z-fulfill.test.ts} | 68 ++++++-- ...n.test.ts => z-handler-activation.test.ts} | 145 +++++++++--------- 13 files changed, 359 insertions(+), 137 deletions(-) rename tests/images/{fulfill.test.ts => z-fulfill.test.ts} (74%) rename tests/images/{handler-activation.test.ts => z-handler-activation.test.ts} (63%) diff --git a/src/images/loop.ts b/src/images/loop.ts index b8fbb3be5b..9a35618316 100644 --- a/src/images/loop.ts +++ b/src/images/loop.ts @@ -12,7 +12,7 @@ */ import type { AdapterRequest, ProviderAdapter } from "../adapters/base"; import { createAdapterEventQueue } from "../adapters/run-turn-queue"; -import type { AdapterEvent, OcxMessage, OcxParsedRequest, OcxThinkingContent, OcxUsage } from "../types"; +import type { AdapterEvent, OcxMessage, OcxParsedRequest, OcxProviderContinuationState, OcxThinkingContent, OcxUsage } from "../types"; import { namespacedToolName } from "../types"; import { bridgeToResponsesSSE } from "../bridge"; import { clearableDeadline, idleDeadline } from "../lib/abort"; @@ -175,6 +175,8 @@ export interface ImageBridgeDeps { * 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; } /** @@ -606,6 +608,7 @@ export async function runWithImageBridge(deps: ImageBridgeDeps): Promise 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 index 0f8e079b0c..9d9c960c00 100644 --- a/src/images/plan.ts +++ b/src/images/plan.ts @@ -6,6 +6,13 @@ 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" @@ -69,6 +76,7 @@ export async function planImageBridge( 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 }, @@ -76,8 +84,6 @@ export async function planImageBridge( toolNames, ...(hostedSize ? { defaultSize: hostedSize } : {}), ...(hostedQuality ? { defaultQuality: hostedQuality } : {}), - ...(typeof config.images?.timeoutMs === "number" && Number.isFinite(config.images.timeoutMs) && config.images.timeoutMs > 0 - ? { timeoutMs: Math.floor(config.images.timeoutMs) } - : {}), + ...(timeoutMs !== undefined ? { timeoutMs } : {}), }; } diff --git a/src/images/synthetic-tool.ts b/src/images/synthetic-tool.ts index f6b1a7cbe5..6a1dc1f0ca 100644 --- a/src/images/synthetic-tool.ts +++ b/src/images/synthetic-tool.ts @@ -3,13 +3,18 @@ import type { OcxTool } from "../types"; /** The function name the chat model sees + the name the loop intercepts. */ export const IMAGE_GEN_TOOL_NAME = "image_gen"; -const IMAGE_GEN_NAMES = new Set([ - "image_gen", "image_generation", "imagegen", - "generate_image", "generateimage", -]); +/** 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 { - return IMAGE_GEN_NAMES.has(name.toLowerCase()); + 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); } /** @@ -22,6 +27,15 @@ 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) { @@ -36,7 +50,7 @@ export function extractHostedImageGeneration( // Also handle the nested Chat Completions shape {type:"function", function:{name:"..."}} for safety. const fnName = typeof obj.name === "string" ? obj.name : (obj as { function?: { name?: string } }).function?.name; - if (fnName && isImageGenName(fnName)) { + if (fnName && isHostedImageGenFunctionName(fnName, hasHostedEntry)) { if (!originalTool) originalTool = obj; toolNames.add(fnName); } diff --git a/src/images/xai-client.ts b/src/images/xai-client.ts index fd26fabb06..2b7f01de56 100644 --- a/src/images/xai-client.ts +++ b/src/images/xai-client.ts @@ -89,7 +89,8 @@ export async function callXaiImages( const timeout = AbortSignal.timeout(deadlineMs); const linkedSignal = signal ? AbortSignal.any([signal, timeout]) : timeout; - const resp = await fetch(`${auth.baseUrl}${endpoint}`, { + const baseUrl = auth.baseUrl.replace(/\/+$/, ""); + const resp = await fetch(`${baseUrl}${endpoint}`, { method: "POST", headers: { "Authorization": `Bearer ${auth.token}`, diff --git a/src/responses/parser.ts b/src/responses/parser.ts index bc13d7fba2..b59ad238b1 100644 --- a/src/responses/parser.ts +++ b/src/responses/parser.ts @@ -16,7 +16,7 @@ import { compactionItemToText } from "./compaction"; import { previousResponseReplayPrefixLength } from "./state"; import { decodeReasoningEnvelope } from "./reasoning-envelope"; import { extractHostedWebSearch, WEB_SEARCH_TOOL_NAME } from "../web-search/synthetic-tool"; -import { extractHostedImageGeneration } from "../images/synthetic-tool"; +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); @@ -125,7 +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"; + if (tool.type === "image_generation" || tool.type === "image_gen") return IMAGE_GEN_TOOL_NAME; if (tool.type === "tool_search") return "tool_search"; return undefined; } diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index b9799c74e3..5f31af77c4 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -1611,6 +1611,13 @@ export async function handleResponses( ); }, ...(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(); diff --git a/tests/images/loop.test.ts b/tests/images/loop.test.ts index acaebfd82c..8315effbc7 100644 --- a/tests/images/loop.test.ts +++ b/tests/images/loop.test.ts @@ -7,28 +7,37 @@ import type { AdapterEvent, OcxParsedRequest } from "../../src/types"; import type { ImageBridgePlan, ImageCallResult } from "../../src/images/types"; const PREV_HOME = process.env.OPENCODEX_HOME; -beforeAll(() => { process.env.OPENCODEX_HOME = join(tmpdir(), "ocx-test-" + randomUUID()); }); -afterAll(() => { if (PREV_HOME === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = PREV_HOME; }); +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"]; -// --- Mock parseStreamWithProgress: simplify to direct delegation --- -mock.module("../../src/web-search/progress-stream", () => ({ - parseStreamWithProgress: async function* (_resp: Response, parse: (r: Response) => AsyncGenerator, _opts: unknown) { - for await (const e of parse(_resp)) yield e; - }, - RoutedModelInactivityError: class extends Error { readonly timeoutMs = 0; }, - WebSearchStreamProtocolError: class extends Error { /* */ }, -})); - -// --- Mock fulfillImageCall --- let fulfillResult: ImageCallResult = { ok: true, model: "grok-imagine-image-quality", prompt: "a cat", files: ["/test/img.png"], count: 1, markdown: "![image](/test/img.png)", }; -mock.module("../../src/images/fulfill", () => ({ - fulfillImageCall: async (): Promise => fulfillResult, -})); -const { runWithImageBridge, clampImageMaxRounds, DEFAULT_MAX_ROUNDS, MAX_ROUNDS_HARD_LIMIT } = await import("../../src/images/loop"); +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[][] = []; @@ -41,6 +50,7 @@ const defaultFulfillResult: ImageCallResult = { beforeEach(() => { fulfillResult = { ...defaultFulfillResult, files: [...defaultFulfillResult.files] }; buildRequestCalls = 0; + streamQueue = []; }); const mockAdapter: ProviderAdapter = { @@ -280,8 +290,9 @@ describe("runWithImageBridge", () => { test("429 key rotation rebuilds the adapter and retries the iteration", async () => { let fetchCalls = 0; let rotations = 0; - const rotatingAdapter: ProviderAdapter = { - name: "test", + 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++; @@ -293,19 +304,24 @@ describe("runWithImageBridge", () => { 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: rotatingAdapter, + adapter: firstAdapter, plan, on429: () => { rotations++; - return rotatingAdapter; + 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"); }); }); diff --git a/tests/images/pinned-https-get.test.ts b/tests/images/pinned-https-get.test.ts index a3038dcb79..e3d1f8017f 100644 --- a/tests/images/pinned-https-get.test.ts +++ b/tests/images/pinned-https-get.test.ts @@ -139,6 +139,42 @@ describe("pinnedHttpsGet transport", () => { 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 @@ -231,4 +267,55 @@ describe("pinnedHttpsGet transport", () => { { 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 index 39df98106b..922f099365 100644 --- a/tests/images/plan.test.ts +++ b/tests/images/plan.test.ts @@ -1,20 +1,30 @@ -import { afterAll, beforeAll, describe, expect, mock, test } from "bun:test"; +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; -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; }); +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; -mock.module("../../src/oauth/index", () => ({ - getValidAccessToken: async () => tokenResult, -})); -const { planImageBridge } = await import("../../src/images/plan"); +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>, @@ -24,7 +34,7 @@ function makeConfig( port: 0, defaultProvider: "test", providers: Object.fromEntries( - Object.entries(providers).map(([k, v]) => [k, { adapter: "openai", baseUrl: "https://api.test.com", ...v }]), + Object.entries(providers).map(([k, v]) => [k, { adapter: "openai-chat", baseUrl: "https://api.test.com", ...v }]), ), ...(images ? { images } : {}), } as OcxConfig; @@ -40,8 +50,8 @@ function makeParsed(withImageGen: boolean): OcxParsedRequest { } as OcxParsedRequest; } -const routed = { adapter: "openai", baseUrl: "https://api.anthropic.com" } as OcxProviderConfig; -const openaiRouted = { adapter: "openai", baseUrl: "https://api.openai.com" } as OcxProviderConfig; +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 () => { @@ -115,6 +125,15 @@ describe("planImageBridge", () => { 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 }); diff --git a/tests/images/synthetic-tool.test.ts b/tests/images/synthetic-tool.test.ts index 7572c94c14..16754d06e3 100644 --- a/tests/images/synthetic-tool.test.ts +++ b/tests/images/synthetic-tool.test.ts @@ -10,8 +10,8 @@ describe("isImageGenName", () => { expect(isImageGenName("IMAGE_GENERATION")).toBe(true); }); - test("'imagegen' → true", () => { - expect(isImageGenName("imagegen")).toBe(true); + test("'imagegen' → false without hosted image tool context", () => { + expect(isImageGenName("imagegen")).toBe(false); }); test("'not_image' → false", () => { @@ -48,6 +48,22 @@ describe("extractHostedImageGeneration", () => { ).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(); }); diff --git a/tests/images/xai-client.test.ts b/tests/images/xai-client.test.ts index 76ed1a7a7f..b13143d59f 100644 --- a/tests/images/xai-client.test.ts +++ b/tests/images/xai-client.test.ts @@ -84,6 +84,28 @@ describe("callXaiImages", () => { 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); diff --git a/tests/images/fulfill.test.ts b/tests/images/z-fulfill.test.ts similarity index 74% rename from tests/images/fulfill.test.ts rename to tests/images/z-fulfill.test.ts index f37fc4a245..fc5501b919 100644 --- a/tests/images/fulfill.test.ts +++ b/tests/images/z-fulfill.test.ts @@ -6,8 +6,27 @@ import type { ImageBridgePlan } from "../../src/images/types"; import type { XaiImageRequest } from "../../src/images/xai-client"; const PREV_HOME = process.env.OPENCODEX_HOME; -beforeAll(() => { process.env.OPENCODEX_HOME = join(tmpdir(), "ocx-test-" + randomUUID()); }); -afterAll(() => { if (PREV_HOME === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = PREV_HOME; }); +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==" }] }; @@ -19,21 +38,6 @@ let materializeFn: (i: number) => Promise = async (i) => `/test/img-${i} let downloadFn: (i: number) => Promise = async (i) => `/test/dl-${i}.png`; let capturedTimeoutMs: number | undefined; -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++), -})); - -const { fulfillImageCall } = await import("../../src/images/fulfill"); const plan = { provider: {} as never, @@ -160,4 +164,34 @@ describe("fulfillImageCall", () => { ); 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/handler-activation.test.ts b/tests/images/z-handler-activation.test.ts similarity index 63% rename from tests/images/handler-activation.test.ts rename to tests/images/z-handler-activation.test.ts index 1af527de9d..324b877328 100644 --- a/tests/images/handler-activation.test.ts +++ b/tests/images/z-handler-activation.test.ts @@ -25,8 +25,6 @@ import type { ProviderAdapter } from "../../src/adapters/base"; */ const PREV_HOME = process.env.OPENCODEX_HOME; -beforeAll(() => { process.env.OPENCODEX_HOME = join(tmpdir(), "ocx-test-" + randomUUID()); }); -afterAll(() => { if (PREV_HOME === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = PREV_HOME; }); // --- Activation spies, flipped by the stubbed runners --- let imageBridgeRun = false; @@ -38,72 +36,79 @@ let runTurnCalled = false; /** Controlled return value for the stubbed planWebSearch (truthy ⇒ web-search plan active). */ let mockWsPlan: unknown = undefined; -// --- Stub adapter-resolve: inject a minimal adapter so no real upstream is hit --- -const actualResolver = await import("../../src/server/adapter-resolve"); -mock.module("../../src/server/adapter-resolve", () => ({ - ...actualResolver, - resolveAdapter(provider: OcxProviderConfig) { - const base = { - name: "test", - buildRequest: async () => ({ url: provider.baseUrl, method: "POST", headers: {}, body: "" }), - async fetchResponse() { - return new Response("data: {\"type\":\"done\"}\n\n", { - status: 200, headers: { "content-type": "text/event-stream" }, - }); - }, - async *parseStream() { yield { type: "done" as const }; }, - }; - if (useRunTurnAdapter) { - return { - ...base, - async runTurn(_parsed: unknown, _incoming: unknown, emit: (event: { type: string }) => void) { - runTurnCalled = true; - emit({ type: "done" }); +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" }, + }); }, - } as ProviderAdapter; - } - return base as ProviderAdapter; - }, -})); - -// --- Stub the image bridge runner: detect activation without hitting the real loop --- -// --- Stub the image bridge runner: detect activation without hitting the real loop --- -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" }, - }); - }, -})); - -// --- Stub the web-search planner + runner: control eligibility and detect activation --- -mock.module("../../src/web-search/index", () => ({ - // Re-export the symbols parser.ts imports statically (deep path → no mock there). - buildWebSearchTool: () => ({ name: "web_search", parameters: { type: "object", properties: {} } }), - WEB_SEARCH_TOOL_NAME: "web_search", - extractHostedWebSearch: (tools: unknown[]) => { - if (!Array.isArray(tools)) return undefined; - for (const t of tools) { - if (t && typeof t === "object" && (t as Record).type === "web_search") { - return { search_context_size: "medium" }; + 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 undefined; - }, - runWithWebSearch: async () => { - webSearchRun = true; - return new Response("data: {\"type\":\"done\"}\n\n", { - status: 200, headers: { "content-type": "text/event-stream" }, - }); - }, - planWebSearch: () => mockWsPlan, - shouldResolveOpenAiWebSearchSidecar: () => false, -})); - -const { handleResponses } = await import("../../src/server/responses"); + 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 { @@ -143,14 +148,12 @@ describe("image bridge dispatch priority (handler activation)", () => { imageBridgeRun = false; webSearchRun = false; mockWsPlan = undefined; const res = await post(false, [{ type: "image_generation" }]); expect(res.status).toBe(400); - // The runner must not execute when the request is rejected upfront. expect(imageBridgeRun).toBe(false); expect((await res.text())).toContain("image bridge requires stream=true"); }); test("dual-tool (image_generation + web_search), both eligible → web-search wins, image bridge deferred", async () => { imageBridgeRun = false; webSearchRun = false; - // A truthy plan ⇒ web-search is eligible; the image bridge must defer to it. mockWsPlan = { backend: "openai" }; const res = await post(true, [{ type: "web_search" }, { type: "image_generation" }]); expect(webSearchRun).toBe(true); @@ -160,10 +163,6 @@ describe("image bridge dispatch priority (handler activation)", () => { test("routed compaction with image_generation tool → image bridge does NOT hijack compaction (#424)", async () => { imageBridgeRun = false; webSearchRun = false; mockWsPlan = undefined; - // A routed-compaction request carries both _compactionRequest and _imageGeneration: - // compaction clears tools/_webSearch but leaves _imageGeneration, so without the - // routedCompaction guard planImageBridge would activate and return a normal Responses - // completion instead of the synthetic compaction item Codex expects. const res = await handleResponses( new Request("http://localhost/v1/responses", { method: "POST", @@ -189,9 +188,7 @@ describe("image bridge dispatch priority (handler activation)", () => { mockWsPlan = { backend: "openai" }; try { const res = await post(true, [{ type: "web_search" }, { type: "image_generation" }]); - // Web-search sidecar cannot drive runTurn adapters — skip it. expect(webSearchRun).toBe(false); - // Image bridge supports runTurn, so it handles the dual-tool turn instead. expect(imageBridgeRun).toBe(true); expect(runTurnCalled).toBe(false); expect(res.headers.get("content-type")).toBe("text/event-stream");