fix(server): bound search and image response bodies

src/lib/bounded-body.ts

@@ -44,6 +44,26 @@ export interface BoundedBodyResult {
 	displaySafe: boolean;
 }
 
+export interface BoundedBytesOptions {
+	/** Abort the read with this signal. Its reason is rethrown by identity. */
+	signal?: AbortSignal;
+	/** Maximum number of raw bytes retained from the response body. */
+	maxBytes: number;
+}
+
+export interface BoundedBytesResult {
+	/**
+	 * Exact raw bytes retained from the response. Empty when the cap was exceeded.
+	 *
+	 * This is a view over internal storage, so `bytes.buffer.byteLength` can exceed
+	 * `bytes.byteLength`. Consumers must honor the view's byteOffset and byteLength
+	 * instead of reading or transferring the backing buffer directly.
+	 */
+	bytes: Uint8Array<ArrayBuffer>;
+	/** True when the body was observed to exceed the byte cap. */
+	oversized: boolean;
+}
+
 const TOTAL_TIMEOUT = Symbol("bounded body total timeout");
 const INACTIVITY_TIMEOUT = Symbol("bounded body inactivity timeout");
 
@@ -83,6 +103,90 @@ function cancelWithoutWaiting(reader: ReadableStreamDefaultReader<Uint8Array>, r
 	}
 }
 
+/**
+ * Consume the original response body as raw bytes under a strict memory ceiling.
+ *
+ * The caller owns any wall-clock deadline through `signal`, which lets one budget
+ * cover both response headers and body consumption. No decoding, cloning, or teeing
+ * occurs, so arbitrary upstream bytes remain unchanged.
+ */
+export async function readBoundedResponseBytes(
+	response: Response,
+	options: BoundedBytesOptions,
+): Promise<BoundedBytesResult> {
+	const signal = options.signal;
+	if (signal?.aborted) throw signal.reason;
+
+	const body = response.body;
+	if (!body) return { bytes: new Uint8Array(0), oversized: false };
+
+	const reader = body.getReader();
+	const maxBytes = options.maxBytes;
+	let retained = new Uint8Array(Math.min(maxBytes, 64 * 1024));
+	let retainedBytes = 0;
+	let mustCancel = false;
+	let cancelReason: unknown;
+
+	let rejectForAbort: ((reason: unknown) => void) | undefined;
+	const aborted = new Promise<never>((_resolve, reject) => {
+		rejectForAbort = reject;
+	});
+	const onAbort = () => rejectForAbort?.(signal?.reason);
+	signal?.addEventListener("abort", onAbort, { once: true });
+	// Close the narrow race between the preflight check and listener install.
+	if (signal?.aborted) onAbort();
+
+	try {
+		while (true) {
+			const read = reader.read();
+			// Observe a late read rejection when abort/cancellation wins the race.
+			void read.catch(() => undefined);
+			const outcome = await Promise.race([read, aborted]);
+			if (signal?.aborted) {
+				mustCancel = true;
+				cancelReason = signal.reason;
+				throw signal.reason;
+			}
+
+			const { value, done } = outcome;
+			if (done) {
+				return { bytes: retained.subarray(0, retainedBytes), oversized: false };
+			}
+			if (!value || value.byteLength === 0) continue;
+
+			if (value.byteLength > maxBytes - retainedBytes) {
+				mustCancel = true;
+				cancelReason = new DOMException("Response body size limit reached", "QuotaExceededError");
+				retained = new Uint8Array(0);
+				retainedBytes = 0;
+				return { bytes: retained, oversized: true };
+			}
+
+			if (retainedBytes + value.byteLength > retained.length) {
+				const grown = new Uint8Array(
+					Math.min(maxBytes, Math.max(retained.length * 2, retainedBytes + value.byteLength)),
+				);
+				grown.set(retained.subarray(0, retainedBytes));
+				retained = grown;
+			}
+			retained.set(value, retainedBytes);
+			retainedBytes += value.byteLength;
+		}
+	} catch (error) {
+		mustCancel = true;
+		cancelReason = error;
+		throw error;
+	} finally {
+		signal?.removeEventListener("abort", onAbort);
+		if (mustCancel) cancelWithoutWaiting(reader, cancelReason);
+		try {
+			reader.releaseLock();
+		} catch {
+			// A pending read can keep the lock briefly while cancel settles.
+		}
+	}
+}
+
 function decodeUtf8(chunks: readonly Uint8Array[], fatal: boolean): string {
 	const decoder = new TextDecoder("utf-8", { fatal });
 	let text = "";

src/server/images.ts

@@ -24,6 +24,7 @@ import {
 } from "../codex/auth-context";
 import { formatCodexProviderForLog } from "../codex/routing";
 import { signalWithTimeout } from "../lib/abort";
+import { readBoundedResponseBytes, type BoundedBytesResult } from "../lib/bounded-body";
 import { sidecarEnter } from "../lib/sidecar-tracker";
 import type { OcxConfig } from "../types";
 import { resolveFirstUsableOpenAiSidecar, selectImagesProvider } from "../providers/openai-sidecar";
@@ -50,7 +51,53 @@ const IMAGES_UPSTREAM_TIMEOUT_MS = 300_000;
  * containing base64-encoded images — typically a few MB. This prevents an oversized or malicious
  * response from exhausting process memory.
  */
-const IMAGES_RESPONSE_MAX_BYTES = 100 * 1024 * 1024;
+export const IMAGES_RESPONSE_MAX_BYTES = 100 * 1024 * 1024;
+
+interface ImageResponseBytesOptions {
+  signal?: AbortSignal;
+  /** Smaller values are used only by focused reader tests. */
+  maxBytes?: number;
+}
+
+function cancelUnlockedResponseBody(response: Response | undefined, reason?: unknown): void {
+  const body = response?.body;
+  if (!body || body.locked) return;
+  try {
+    void body.cancel(reason).catch(() => undefined);
+  } catch {
+    // A broken stream can throw synchronously from cancel().
+  }
+}
+
+/**
+ * Read one image relay response as exact raw bytes under the same cap used by
+ * the public handler. A trustworthy Content-Length can reject obvious excess
+ * before attaching a reader; the streaming reader remains authoritative when
+ * the header is absent or understated.
+ */
+export async function readImageResponseBytes(
+  response: Response,
+  options: ImageResponseBytesOptions = {},
+): Promise<BoundedBytesResult> {
+  const signal = options.signal;
+  if (signal?.aborted) {
+    cancelUnlockedResponseBody(response, signal.reason);
+    throw signal.reason;
+  }
+
+  const maxBytes = options.maxBytes ?? IMAGES_RESPONSE_MAX_BYTES;
+  const contentLength = response.headers.get("content-length");
+  const declaredBytes = contentLength === null ? Number.NaN : Number(contentLength);
+  if (Number.isSafeInteger(declaredBytes) && declaredBytes > maxBytes) {
+    cancelUnlockedResponseBody(
+      response,
+      new DOMException("Response body size limit reached", "QuotaExceededError"),
+    );
+    return { bytes: new Uint8Array(0), oversized: true };
+  }
+
+  return readBoundedResponseBytes(response, { maxBytes, signal });
+}
 
 const CCA_IMAGE_MODEL = "gemini-3.1-flash-image";
 
@@ -440,34 +487,47 @@ export async function handleImages(
   const timeoutMs = config.images?.timeoutMs ?? IMAGES_UPSTREAM_TIMEOUT_MS;
   const linkedSignal = signalWithTimeout(timeoutMs, req.signal);
   const sidecarExit = sidecarEnter("images");
+  let upstreamResponse: Response | undefined;
   try {
     // Images POSTs create paid, non-idempotent work. One fetch only: no reset retry without a
     // source-proven idempotency contract.
-    const upstreamResponse = await fetch(url, {
+    upstreamResponse = await fetch(url, {
       method: "POST",
       headers,
       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)`);
+    const observed = await readImageResponseBytes(upstreamResponse, {
+      maxBytes: IMAGES_RESPONSE_MAX_BYTES,
+      signal: linkedSignal.signal,
+    });
+    if (observed.oversized) {
+      forward?.recordOutcome?.(upstreamResponse.status);
+      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> = {};
     const contentType = upstreamResponse.headers.get("content-type");
     if (contentType) relayHeaders["content-type"] = contentType;
-    return new Response(payload, { status: upstreamResponse.status, headers: relayHeaders });
+    // Fetch represents 204/205/304 responses with a null body. Preserve that
+    // invariant: constructing those statuses with even an empty Uint8Array throws.
+    const relayBody = upstreamResponse.body === null ? null : observed.bytes;
+    const relayResponse = new Response(relayBody, {
+      status: upstreamResponse.status,
+      headers: relayHeaders,
+    });
+    forward?.recordOutcome?.(upstreamResponse.status);
+    return relayResponse;
   } catch (err) {
     // Client cancel first: it aborts the linked signal too, and must not be logged as an
     // upstream failure (499 maps to client_closed_request in the request log).
     if (req.signal.aborted) {
       return formatErrorResponse(499, "client_closed_request", `image ${endpoint} request canceled by client`);
     }
-    if (err instanceof Error && err.name === "TimeoutError") {
+    if (linkedSignal.signal.aborted || (err instanceof Error && err.name === "TimeoutError")) {
       forward?.recordOutcome?.("timeout");
       // codex retries 5xx up to 4 more times; a retried 504 is acceptable for a transient hang.
       return formatErrorResponse(504, "upstream_error", `image ${endpoint} upstream timed out`);
@@ -481,5 +541,9 @@ export async function handleImages(
   } finally {
     sidecarExit();
     linkedSignal.cleanup();
+    // If cancellation won before readImageResponseBytes attached its reader, the
+    // response still owns an upstream body/socket. Consumed or locked bodies are
+    // already owned by the reader and need no second cancellation.
+    cancelUnlockedResponseBody(upstreamResponse);
   }
 }

src/server/search.ts

@@ -21,6 +21,7 @@ import {
 import { codexAccountNamespaceForModel } from "../codex/account-namespace-match";
 import { formatCodexProviderForLog } from "../codex/routing";
 import { signalWithTimeout } from "../lib/abort";
+import { readBoundedResponseBytes } from "../lib/bounded-body";
 import { sidecarEnter } from "../lib/sidecar-tracker";
 import type { OcxConfig } from "../types";
 import {
@@ -44,7 +45,7 @@ import { codexAccountSelectionForTurn } from "./lifecycle";
  * long-running search).
  */
 const SEARCH_UPSTREAM_TIMEOUT_MS = 200_000;
-const SEARCH_RESPONSE_MAX_BYTES = 16 * 1024 * 1024;
+export const SEARCH_RESPONSE_MAX_BYTES = 16 * 1024 * 1024;
 
 export async function handleSearch(
   req: Request,
@@ -144,27 +145,36 @@ export async function handleSearch(
   const timeoutMs = config.search?.timeoutMs ?? SEARCH_UPSTREAM_TIMEOUT_MS;
   const linkedSignal = signalWithTimeout(timeoutMs, req.signal);
   const sidecarExit = sidecarEnter("search");
+  let upstreamResponse: Response | undefined;
   try {
-    const upstreamResponse = await fetch(url, {
+    upstreamResponse = await fetch(url, {
       method: "POST",
       headers,
       body: JSON.stringify(relayBody),
       signal: linkedSignal.signal,
     });
-    const payload = await upstreamResponse.arrayBuffer();
-    if (payload.byteLength > SEARCH_RESPONSE_MAX_BYTES) {
-      return formatErrorResponse(502, "upstream_error", `search response too large (${payload.byteLength} bytes)`);
+    const observed = await readBoundedResponseBytes(upstreamResponse, {
+      maxBytes: SEARCH_RESPONSE_MAX_BYTES,
+      signal: linkedSignal.signal,
+    });
+    if (observed.oversized) {
+      upstream.recordOutcome?.(upstreamResponse.status);
+      return formatErrorResponse(
+        502,
+        "upstream_error",
+        `search response too large (exceeded ${SEARCH_RESPONSE_MAX_BYTES} bytes)`,
+      );
     }
     upstream.recordOutcome?.(upstreamResponse.status);
     const relayHeaders: Record<string, string> = {};
     const contentType = upstreamResponse.headers.get("content-type");
     if (contentType) relayHeaders["content-type"] = contentType;
-    return new Response(payload, { status: upstreamResponse.status, headers: relayHeaders });
+    return new Response(observed.bytes, { status: upstreamResponse.status, headers: relayHeaders });
   } catch (err) {
     if (req.signal.aborted) {
       return formatErrorResponse(499, "client_closed_request", "search request canceled by client");
     }
-    if (err instanceof Error && err.name === "TimeoutError") {
+    if (linkedSignal.signal.aborted || (err instanceof Error && err.name === "TimeoutError")) {
       upstream.recordOutcome?.("timeout");
       return formatErrorResponse(504, "upstream_error", "search upstream timed out");
     }
@@ -177,5 +187,11 @@ export async function handleSearch(
   } finally {
     sidecarExit();
     linkedSignal.cleanup();
+    // A response that aborted before the reader attached still owns its upstream socket.
+    // Consumed/cancelled bodies are already closed; locked bodies remain owned by the reader.
+    const pendingBody = upstreamResponse?.body;
+    if (pendingBody && !pendingBody.locked) {
+      try { void pendingBody.cancel().catch(() => undefined); } catch { /* already closed */ }
+    }
   }
 }