[WRONG BRANCH] fix(server): bound image relay response bodies

src/server/images.ts

@@ -52,6 +52,45 @@ const IMAGES_UPSTREAM_TIMEOUT_MS = 300_000;
  */
 const IMAGES_RESPONSE_MAX_BYTES = 100 * 1024 * 1024;
 
+async function readImageResponseBody(response: Response): Promise<Uint8Array<ArrayBuffer> | undefined> {
+  const contentLength = response.headers.get("content-length");
+  if (contentLength !== null) {
+    const declaredBytes = Number(contentLength);
+    if (Number.isFinite(declaredBytes) && declaredBytes > IMAGES_RESPONSE_MAX_BYTES) {
+      try { void response.body?.cancel().catch(() => undefined); } catch { /* best-effort upstream abort */ }
+      return undefined;
+    }
+  }
+
+  const reader = response.body?.getReader();
+  if (!reader) return new Uint8Array();
+  let payload = new Uint8Array(Math.min(IMAGES_RESPONSE_MAX_BYTES, 64 * 1024));
+  let length = 0;
+  try {
+    for (;;) {
+      const { done, value } = await reader.read();
+      if (done) return payload.subarray(0, length);
+      if (!value || value.byteLength === 0) continue;
+      if (value.byteLength > IMAGES_RESPONSE_MAX_BYTES - length) {
+        try { void reader.cancel().catch(() => undefined); } catch { /* best-effort upstream abort */ }
+        return undefined;
+      }
+      if (length + value.byteLength > payload.byteLength) {
+        const grown = new Uint8Array(Math.min(
+          IMAGES_RESPONSE_MAX_BYTES,
+          Math.max(payload.byteLength * 2, length + value.byteLength),
+        ));
+        grown.set(payload.subarray(0, length));
+        payload = grown;
+      }
+      payload.set(value, length);
+      length += value.byteLength;
+    }
+  } finally {
+    reader.releaseLock();
+  }
+}
+
 const CCA_IMAGE_MODEL = "gemini-3.1-flash-image";
 
 /**
@@ -449,12 +488,11 @@ export async function handleImages(
       body: JSON.stringify(body),
       signal: linkedSignal.signal,
     });
-    // Buffer rather than stream: the payload is one JSON document (base64 image, typically a few
-    // MB), and buffering keeps the timeout window covering the whole exchange. Cap the size to
-    // prevent an oversized response from exhausting process memory.
-    const payload = await upstreamResponse.arrayBuffer();
-    if (payload.byteLength > IMAGES_RESPONSE_MAX_BYTES) {
-      return formatErrorResponse(502, "upstream_error", `image ${endpoint} response too large (${payload.byteLength} bytes)`);
+    // The payload is one JSON document, so buffer it under a strict byte ceiling. The reader
+    // cancels as soon as the response exceeds the ceiling rather than allocating the whole body.
+    const payload = await readImageResponseBody(upstreamResponse);
+    if (!payload) {
+      return formatErrorResponse(502, "upstream_error", `image ${endpoint} response too large (exceeded ${IMAGES_RESPONSE_MAX_BYTES} bytes)`);
     }
     forward?.recordOutcome?.(upstreamResponse.status);
     const relayHeaders: Record<string, string> = {};

tests/server-images.test.ts

@@ -290,6 +290,45 @@ test("falls back to a keyed openai-responses provider when no forward provider e
   }
 });
 
+test("rejects an oversized image relay response before buffering its body", async () => {
+  let bodyCanceled = false;
+  globalThis.fetch = (async (input: RequestInfo | URL) => {
+    const requestUrl = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url;
+    if (requestUrl === "https://api.openai.com/v1/images/generations") {
+      return new Response(new ReadableStream<Uint8Array>({
+        cancel() { bodyCanceled = true; },
+      }), {
+        status: 200,
+        headers: {
+          "content-type": "application/json",
+          "content-length": String(100 * 1024 * 1024 + 1),
+        },
+      });
+    }
+    return originalFetch(input);
+  }) as typeof fetch;
+
+  const config = {
+    port: 0,
+    defaultProvider: "openai-apikey",
+    openaiProviderTierVersion: 2,
+    providers: {
+      openai: disabledOpenAiProvider,
+      "openai-apikey": keyedProvider(),
+    },
+  } as OcxConfig;
+  const { handleImages } = await import("../src/server/images");
+  const response = await handleImages(new Request("http://localhost/v1/images/generations", {
+    method: "POST",
+    headers: { "content-type": "application/json" },
+    body: JSON.stringify({ prompt: "a cat", model: "gpt-image-2" }),
+  }), config, "generations", { model: "", provider: "" } as never);
+
+  expect(response.status).toBe(502);
+  expect(await response.text()).toContain("response too large");
+  expect(bodyCanceled).toBe(true);
+});
+
 test("an explicit custom Images provider uses its configured endpoint, key, and headers", async () => {
   const captured: CapturedRequest[] = [];
   const upstream = Bun.serve({