diff --git a/open-sse/handlers/imageGeneration.ts b/open-sse/handlers/imageGeneration.ts index 94c51576e9..7940c34a83 100644 --- a/open-sse/handlers/imageGeneration.ts +++ b/open-sse/handlers/imageGeneration.ts @@ -6,6 +6,10 @@ import { CHATGPT_WEB_RETIRED_MESSAGE, isCommonChatGptWebRetiredProviderId, } from "@/shared/constants/chatgptWebRetirement"; +import { + isMicrosoftDesignerWebRetiredProviderId, + MICROSOFT_DESIGNER_WEB_RETIRED_MESSAGE, +} from "@/shared/constants/designerWebRetirement"; import { getImageProvider, parseImageModel } from "../config/imageRegistry.ts"; import { HTTP_STATUS } from "../config/constants.ts"; @@ -38,10 +42,6 @@ import { getConfiguredTimeout, } from "@/shared/utils/fetchTimeout"; import { sanitizeErrorMessage, sanitizeUpstreamDetails } from "../utils/error.ts"; -import { - isMicrosoftDesignerWebRetiredProviderId, - MICROSOFT_DESIGNER_WEB_RETIRED_MESSAGE, -} from "@/shared/constants/designerWebRetirement"; import { handleSDWebUIImageGeneration } from "./imageGeneration/providers/sdWebUI.ts"; import { handleHyperbolicImageGeneration } from "./imageGeneration/providers/hyperbolic.ts"; @@ -180,6 +180,17 @@ const OPENAI_IMAGE_TO_IMAGE_MODELS = new Set([ const IMAGE_ASPECT_RATIO_PATTERN = /^\d+:\d+$/; const IMAGE_SIZE_PATTERN = /^(?:1K|2K|4K)$/; +/** + * `fetchRemoteImage` options for any URL that did not originate from an OmniRoute-controlled + * host (caller-supplied `image_url`, upstream-returned result URLs). + * + * GHSA-34rg-3pqj-35g9 / #13883: pin `public-only` (never the operator outbound policy, which + * would let a request body reach loopback/LAN) and `pinDns: true` to close the DNS-rebinding + * TOCTOU where a second, un-pinned resolution at connect time could answer differently than + * the validated lookup and bypass the guard. + */ +const UNTRUSTED_REMOTE_IMAGE_FETCH_OPTIONS = { guard: "public-only", pinDns: true } as const; + /** * Resolve the upstream images endpoint for a custom (OpenAI-compatible) image * provider node (#3205). @@ -239,9 +250,17 @@ function normalizeImageAspectRatio(value: unknown, fallbackSize: unknown): strin return mapImageSize(typeof fallbackSize === "string" ? fallbackSize : null); } -function normalizeImageGenerationSize(snakeCaseValue: unknown, camelCaseValue: unknown): string { - const value = snakeCaseValue ?? camelCaseValue; - if (typeof value !== "string") return "1K"; +/** + * Normalize the caller's `image_size` for Antigravity's `imageConfig.imageSize`. + * + * This is the output-resolution axis (`1K` | `2K` | `4K`, the values Gemini image models + * accept), distinct from the `size`/`aspect_ratio` axis handled by `normalizeImageAspectRatio`. + * Returns `undefined` when the caller sent nothing usable (absent or non-string), so the key + * is left out and the upstream default applies; a string that is not one of the accepted + * values is clamped to `1K` because upstream rejects anything else. + */ +function normalizeImageGenerationSize(value: unknown): string | undefined { + if (typeof value !== "string") return undefined; const normalized = value.trim().toUpperCase(); return IMAGE_SIZE_PATTERN.test(normalized) ? normalized : "1K"; } @@ -394,6 +413,9 @@ export async function handleImageGeneration({ clientHeaders = null, peerLocality = null, }) { + // Retirement guards: the retired-provider sets hold bare provider ids only, so testing + // the `/` prefix (or the whole model when it carries no slash) covers both the + // `provider/model` and bare-id request shapes. const requestedModel = typeof body?.model === "string" ? body.model : ""; const slash = requestedModel.indexOf("/"); const requestedPrefix = slash > 0 ? requestedModel.slice(0, slash) : requestedModel; @@ -408,15 +430,13 @@ export async function handleImageGeneration({ }; } - const requestedProvider = slash > 0 ? requestedModel.slice(0, slash) : null; if ( isCommonChatGptWebRetiredProviderId(resolvedProvider) || - isCommonChatGptWebRetiredProviderId(requestedProvider) || - isCommonChatGptWebRetiredProviderId(requestedModel) + isCommonChatGptWebRetiredProviderId(requestedPrefix) ) { return { success: false, - status: 410, + status: HTTP_STATUS.GONE, error: CHATGPT_WEB_RETIRED_MESSAGE, code: CHATGPT_WEB_RETIRED_ERROR_CODE, }; @@ -1030,7 +1050,7 @@ async function handleGeminiImageGeneration({ model, providerConfig, body, creden typeof body.n === "number" && Number.isFinite(body.n) && body.n > 0 ? Math.floor(body.n) : 1; const promptText = typeof body.prompt === "string" ? body.prompt : String(body.prompt ?? ""); const aspectRatio = normalizeImageAspectRatio(body.aspect_ratio, body.size); - const imageSize = normalizeImageGenerationSize(body.image_size, body.imageSize); + const imageSize = normalizeImageGenerationSize(body.image_size); // Summarized request for call log const logRequestBody = { @@ -1068,7 +1088,7 @@ async function handleGeminiImageGeneration({ model, providerConfig, body, creden candidateCount, imageConfig: { aspectRatio, - imageSize, + ...(imageSize ? { imageSize } : {}), }, }, }, @@ -1088,7 +1108,7 @@ async function handleGeminiImageGeneration({ model, providerConfig, body, creden const promptPreview = promptText.slice(0, 60); log.info( "IMAGE", - `antigravity/${model} (gemini) | prompt: "${promptPreview}..." | ${aspectRatio} ${imageSize}` + `antigravity/${model} (gemini) | prompt: "${promptPreview}..." | ${aspectRatio} ${imageSize ?? "default"}` ); } @@ -2243,11 +2263,8 @@ export async function resolveImageSource(source) { } if (isHttpUrl(trimmed)) { - // GHSA-34rg-3pqj-35g9 / #13883: caller-input URL — pin `public-only` (never the operator - // outbound policy, which would let a request body reach loopback/LAN) and `pinDns: true` - // to close the DNS-rebinding TOCTOU where a second, un-pinned resolution at connect time - // could answer differently than the validated lookup and bypass the guard. - const remoteImage = await fetchRemoteImage(trimmed, { guard: "public-only", pinDns: true }); + // Caller-input URL — see UNTRUSTED_REMOTE_IMAGE_FETCH_OPTIONS. + const remoteImage = await fetchRemoteImage(trimmed, UNTRUSTED_REMOTE_IMAGE_FETCH_OPTIONS); return { buffer: remoteImage.buffer, base64: remoteImage.buffer.toString("base64"), @@ -3276,10 +3293,9 @@ export async function normalizeNanoBananaTaskResult(taskData, body, log) { if (urlCandidates.length > 0) { const firstUrl = urlCandidates[0]; - // GHSA-34rg-3pqj-35g9 / #13883: upstream-supplied result URL, not an OmniRoute- - // controlled host — pin `public-only`, never the operator outbound policy, and - // `pinDns: true` to close the DNS-rebinding TOCTOU (see `resolveImageSource`). - const remoteImage = await fetchRemoteImage(firstUrl, { guard: "public-only", pinDns: true }); + // Upstream-supplied result URL, not an OmniRoute-controlled host — see + // UNTRUSTED_REMOTE_IMAGE_FETCH_OPTIONS. + const remoteImage = await fetchRemoteImage(firstUrl, UNTRUSTED_REMOTE_IMAGE_FETCH_OPTIONS); const base64 = remoteImage.buffer.toString("base64"); return [{ b64_json: base64, revised_prompt: body.prompt }]; } diff --git a/tests/unit/image-generation-handler.test.ts b/tests/unit/image-generation-handler.test.ts index 8749c91e90..61c1ed0307 100644 --- a/tests/unit/image-generation-handler.test.ts +++ b/tests/unit/image-generation-handler.test.ts @@ -24,7 +24,8 @@ process.on("exit", () => { const { IMAGE_PROVIDERS, parseImageModel, getAllImageModels } = await import("../../open-sse/config/imageRegistry.ts"); -const { handleImageGeneration } = await import("../../open-sse/handlers/imageGeneration.ts"); +const { handleImageGeneration, handleOpenAIImageEdit } = + await import("../../open-sse/handlers/imageGeneration.ts"); const { setPinnedFetchTestOverride } = await import("../../src/shared/network/remoteImageFetch.ts"); function mockFetch(impl) { globalThis.fetch = impl; @@ -149,6 +150,98 @@ test("handleImageGeneration uses synthetic OpenAI-compatible routing for resolve } }); +// LEDGER-12 — `hasUsableImage` (fetchImageEndpoint): a 2xx whose items carry no usable +// `b64_json`/`url` must surface as a retryable 502 so image combos fall back, on both the +// generation and the edit path. +function mockOpenAICompatibleUpstream(payload: unknown) { + return async () => + new Response(JSON.stringify(payload), { + status: 200, + headers: { "content-type": "application/json" }, + }); +} + +const usableImageRequest = { + body: { model: "custom-provider/super-image", prompt: "retro poster" }, + credentials: { + apiKey: "custom-key", + baseUrl: "https://custom.example.com/v1/images/generations", + }, + resolvedProvider: "custom-provider", + log: null, +}; + +test("handleImageGeneration treats a 200 whose only item has a blank url as a retryable 502", async () => { + const originalFetch = globalThis.fetch; + globalThis.fetch = mockOpenAICompatibleUpstream({ data: [{ url: "" }] }); + try { + const result = await handleImageGeneration(usableImageRequest); + assert.equal(result.success, false); + assert.equal(result.status, 502); + assert.match(String(result.error), /without an image payload/); + } finally { + globalThis.fetch = originalFetch; + } +}); + +test("handleImageGeneration treats a 200 whose only item is not an object as a retryable 502", async () => { + const originalFetch = globalThis.fetch; + globalThis.fetch = mockOpenAICompatibleUpstream({ data: ["https://cdn.example.com/x.png"] }); + try { + const result = await handleImageGeneration(usableImageRequest); + assert.equal(result.success, false); + assert.equal(result.status, 502); + } finally { + globalThis.fetch = originalFetch; + } +}); + +test("handleImageGeneration keeps a well-formed 200 image payload as success", async () => { + const originalFetch = globalThis.fetch; + globalThis.fetch = mockOpenAICompatibleUpstream({ + created: 123, + data: [{ url: "" }, { b64_json: "ZmFrZQ==" }], + }); + try { + const result = await handleImageGeneration(usableImageRequest); + assert.equal(result.success, true); + assert.equal(result.status, undefined); + assert.deepEqual(result.data, { created: 123, data: [{ url: "" }, { b64_json: "ZmFrZQ==" }] }); + } finally { + globalThis.fetch = originalFetch; + } +}); + +test("handleOpenAIImageEdit applies the same usable-image gate to the edit path", async () => { + const originalFetch = globalThis.fetch; + const editRequest = { + model: "super-image", + provider: "custom-provider", + credentials: { + apiKey: "custom-key", + providerSpecificData: { baseUrl: "https://custom.example.com/v1" }, + }, + prompt: "make it blue", + imageBytes: Buffer.from([0x89, 0x50, 0x4e, 0x47]), + imageMime: "image/png", + log: null, + }; + try { + globalThis.fetch = mockOpenAICompatibleUpstream({ data: [{ b64_json: "" }] }); + const blank = await handleOpenAIImageEdit(editRequest); + assert.equal(blank.success, false); + assert.equal(blank.status, 502); + assert.match(String(blank.error), /without an image payload/); + + globalThis.fetch = mockOpenAICompatibleUpstream({ data: [{ b64_json: "ZmFrZQ==" }] }); + const ok = await handleOpenAIImageEdit(editRequest); + assert.equal(ok.success, true); + assert.deepEqual(ok.data.data, [{ b64_json: "ZmFrZQ==" }]); + } finally { + globalThis.fetch = originalFetch; + } +}); + test("handleImageGeneration polls KIE image tasks and returns URLs on success", async () => { const originalFetch = globalThis.fetch; let createPayload; @@ -787,6 +880,69 @@ test("handleImageGeneration sends Antigravity image requests with native image_g } }); +// LEDGER-6 — `image_size` is forwarded to Antigravity's `imageConfig.imageSize` only when the +// caller supplied it; an unrecognised string is clamped to "1K", everything else leaves the +// key out so the upstream default applies. +async function captureAntigravityImageConfig(extraBody: Record) { + const originalFetch = globalThis.fetch; + let captured; + globalThis.fetch = async (_url, options = {}) => { + captured = JSON.parse(String(options.body || "{}")); + return new Response(JSON.stringify({ response: { candidates: [] } }), { + status: 200, + headers: { "content-type": "application/json" }, + }); + }; + try { + const result = await handleImageGeneration({ + body: { + model: "antigravity/gemini-3.1-flash-image-preview", + prompt: "painted beach", + aspect_ratio: "3:4", + ...extraBody, + }, + credentials: { accessToken: "ag-token", projectId: "project-123" }, + log: null, + }); + assert.equal(result.success, true); + return captured.request.generationConfig.imageConfig; + } finally { + globalThis.fetch = originalFetch; + } +} + +test("handleImageGeneration omits Antigravity imageSize when image_size is not supplied", async () => { + assert.deepEqual(await captureAntigravityImageConfig({}), { aspectRatio: "3:4" }); +}); + +test("handleImageGeneration forwards a valid Antigravity image_size normalized to upper case", async () => { + assert.deepEqual(await captureAntigravityImageConfig({ image_size: "2K" }), { + aspectRatio: "3:4", + imageSize: "2K", + }); + assert.deepEqual(await captureAntigravityImageConfig({ image_size: " 4k " }), { + aspectRatio: "3:4", + imageSize: "4K", + }); +}); + +test("handleImageGeneration clamps an unrecognised Antigravity image_size string to 1K", async () => { + assert.deepEqual(await captureAntigravityImageConfig({ image_size: "1024x1024" }), { + aspectRatio: "3:4", + imageSize: "1K", + }); +}); + +test("handleImageGeneration omits Antigravity imageSize for a non-string image_size", async () => { + assert.deepEqual(await captureAntigravityImageConfig({ image_size: 2 }), { aspectRatio: "3:4" }); +}); + +test("handleImageGeneration ignores the camelCase imageSize alias on Antigravity requests", async () => { + assert.deepEqual(await captureAntigravityImageConfig({ imageSize: "4K" }), { + aspectRatio: "3:4", + }); +}); + test("handleImageGeneration rejects Antigravity image requests without projectId", async () => { const originalFetch = globalThis.fetch; globalThis.fetch = async () => {