pjandro/OmniRoute
1
0
Fork 0
mirror of https://github.com/diegosouzapw/OmniRoute.git synced 2026-09-22 14:52:22 +03:00
Code Issues Packages Projects Releases Wiki Activity

fix(review): resolve LEDGER-4/40/6/24/28/29/12/32/42 (omni-code-review)

Browse Source
This commit is contained in:
diegosouzapw
2026-09-21 14:19:27 -03:00
parent 11d71e959c
commit 8a587a7321
2 changed files with 196 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

62
open-sse/handlers/imageGeneration.ts
View File

@@ -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 `<provider>/` 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 }];
}

158
tests/unit/image-generation-handler.test.ts
View File

@@ -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<string, unknown>) {
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 () => {