diff --git a/src/lib/bounded-body.ts b/src/lib/bounded-body.ts index 87378cec3..1d761dd07 100644 --- a/src/lib/bounded-body.ts +++ b/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; + /** 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, 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 { + 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((_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 = ""; diff --git a/src/server/images.ts b/src/server/images.ts index a35ca84e6..edc710933 100644 --- a/src/server/images.ts +++ b/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 { + 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 = {}; 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); } } diff --git a/src/server/search.ts b/src/server/search.ts index 76ace97cd..94508f3d7 100644 --- a/src/server/search.ts +++ b/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 = {}; 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 */ } + } } } diff --git a/tests/bounded-body.test.ts b/tests/bounded-body.test.ts index 2bf9efff8..fa90e3ec5 100644 --- a/tests/bounded-body.test.ts +++ b/tests/bounded-body.test.ts @@ -2,6 +2,7 @@ import { describe, expect, test } from "bun:test"; import { BOUNDED_BODY_MAX_BYTES, boundedBodyBufferGrowthsForTests, + readBoundedResponseBytes, readBoundedResponseBody, } from "../src/lib/bounded-body"; import { UPSTREAM_JSON_BODY_READ_OPTIONS } from "../src/server/responses/core"; @@ -364,4 +365,92 @@ describe("readBoundedResponseBody", () => { expect(response.bodyUsed).toBe(true); expect(cloneCalls).toBe(0); }); + + test("reads arbitrary response bytes exactly through the raw primitive", async () => { + const expected = new Uint8Array([0x00, 0xff, 0x80, 0xc3, 0x28]); + const response = responseFromChunks(expected.subarray(0, 2), expected.subarray(2)); + + const result = await readBoundedResponseBytes(response, { maxBytes: expected.byteLength }); + + expect(result.oversized).toBe(false); + expect(Array.from(result.bytes)).toEqual(Array.from(expected)); + }); + + test("raw byte reads discard the prefix and cancel without draining the stream", async () => { + let cancelled = false; + let tailPulled = false; + const chunks = [new Uint8Array(3), new Uint8Array(3), new Uint8Array([0x7f]), new Uint8Array([0x7e])]; + const response = new Response(new ReadableStream({ + pull(controller) { + const chunk = chunks.shift(); + if (!chunk) return controller.close(); + if (chunk.byteLength === 1 && chunk[0] === 0x7e) tailPulled = true; + controller.enqueue(chunk); + }, + cancel() { cancelled = true; }, + })); + + const result = await readBoundedResponseBytes(response, { maxBytes: 5 }); + + expect(result.oversized).toBe(true); + expect(result.bytes.byteLength).toBe(0); + expect(cancelled).toBe(true); + // WHATWG streams may prefetch one queued chunk, but cancellation must stop further draining. + expect(tailPulled).toBe(false); + }); + + test("raw byte reads preserve the parent abort reason and cancel the stream", async () => { + const parent = new AbortController(); + const reason = { code: "client-stopped" }; + let cancelled = false; + const response = new Response(new ReadableStream({ + cancel() { cancelled = true; }, + })); + const reading = readBoundedResponseBytes(response, { maxBytes: 5, signal: parent.signal }); + parent.abort(reason); + + let caught: unknown; + try { await reading; } catch (error) { caught = error; } + expect(caught).toBe(reason); + expect(cancelled).toBe(true); + }); + + test("raw byte cancellation rejection is observed", async () => { + const unhandled: unknown[] = []; + const listener = (reason: unknown) => unhandled.push(reason); + process.on("unhandledRejection", listener); + try { + // Bun's test runner fails a test on a real unhandled rejection even when a + // process listener is installed. Prove the pinned runtime's event path in an + // isolated process, then keep this process clean for the negative assertion. + const control = Bun.spawnSync({ + cmd: [ + process.execPath, + "-e", + 'process.on("unhandledRejection", () => console.log("observed"));' + + 'void Promise.reject(new Error("control"));setTimeout(() => {}, 10);', + ], + stdout: "pipe", + stderr: "pipe", + }); + expect(control.exitCode).toBe(0); + expect(new TextDecoder().decode(control.stdout)).toContain("observed"); + + let cancelCalls = 0; + const response = new Response(new ReadableStream({ + start(controller) { controller.enqueue(new Uint8Array(6)); }, + cancel() { + cancelCalls++; + return Promise.reject(new Error("cancel failed")); + }, + })); + const result = await readBoundedResponseBytes(response, { maxBytes: 5 }); + expect(result.oversized).toBe(true); + await new Promise(resolve => setTimeout(resolve, 10)); + expect(cancelCalls).toBe(1); + expect(unhandled).toEqual([]); + } finally { + process.off("unhandledRejection", listener); + } + }); }); diff --git a/tests/server-images.test.ts b/tests/server-images.test.ts index 437c76c25..86e09624f 100644 --- a/tests/server-images.test.ts +++ b/tests/server-images.test.ts @@ -8,10 +8,11 @@ import { existsSync, mkdirSync, rmSync } from "node:fs"; import { join } from "node:path"; import { saveCodexAccountCredential } from "../src/codex/account-store"; import { clearAccountNeedsReauth, clearAccountQuota } from "../src/codex/auth-api"; -import { clearCodexUpstreamHealth, clearThreadAccountMap } from "../src/codex/routing"; +import { clearCodexUpstreamHealth, clearThreadAccountMap, getCodexUpstreamHealth } from "../src/codex/routing"; import { saveConfig } from "../src/config"; import { selectImagesProvider } from "../src/providers/openai-sidecar"; import { startServer } from "../src/server"; +import { handleImages, IMAGES_RESPONSE_MAX_BYTES, readImageResponseBytes } from "../src/server/images"; import { saveCredential } from "../src/oauth/store"; import type { OcxConfig } from "../src/types"; import { ANTIGRAVITY_REQUEST_UA } from "../src/adapters/google-antigravity-wire"; @@ -127,6 +128,33 @@ function keyedProvider(_baseUrl = "") { return { adapter: "openai-responses", baseUrl: "https://api.openai.com/v1", apiKey: "sk-platform-key" }; } +test("image response byte reader enforces the stream cap when Content-Length is understated", async () => { + const chunks = [ + new Uint8Array(5), + new Uint8Array([0x01]), + new Uint8Array([0x7f]), + new Uint8Array([0x7e]), + ]; + let canceled = false; + let tailPulled = false; + const response = new Response(new ReadableStream({ + pull(controller) { + const chunk = chunks.shift(); + if (!chunk) return controller.close(); + if (chunk.byteLength === 1 && chunk[0] === 0x7e) tailPulled = true; + controller.enqueue(chunk); + }, + cancel() { canceled = true; }, + }), { headers: { "content-length": "1" } }); + + const observed = await readImageResponseBytes(response, { maxBytes: 5 }); + expect(observed.oversized).toBe(true); + expect(observed.bytes.byteLength).toBe(0); + expect(canceled).toBe(true); + // One queued chunk may be prefetched; cancellation must stop later draining. + expect(tailPulled).toBe(false); +}); + test("POST /v1/images/generations relays to the ChatGPT forward provider with forwarded auth", async () => { const captured: CapturedRequest[] = []; const upstream = fakeImagesUpstream(captured); @@ -718,6 +746,225 @@ test("relays upstream error status and body verbatim", async () => { } }); +test("relays arbitrary upstream image bytes with status and content-type unchanged", async () => { + const payload = new Uint8Array([0x00, 0xff, 0x80, 0x7f]); + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + const requestUrl = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; + if (new URL(requestUrl).hostname === "chatgpt.com") { + return new Response(payload, { + status: 418, + headers: { "content-type": "application/octet-stream" }, + }); + } + return originalFetch(input, init); + }) as typeof fetch; + saveConfig(forwardConfig()); + + const server = startServer(0); + try { + const response = await fetch(new URL("/v1/images/generations", server.url), { + method: "POST", + headers: { + "content-type": "application/json", + authorization: `Bearer ${DIRECT_CHATGPT_TOKEN}`, + "chatgpt-account-id": "acct-123", + }, + body: JSON.stringify({ prompt: "a cat", model: "gpt-image-2" }), + }); + expect(response.status).toBe(418); + expect(response.headers.get("content-type")).toContain("application/octet-stream"); + expect(new Uint8Array(await response.arrayBuffer())).toEqual(payload); + } finally { + await server.stop(true); + } +}); + +test("relays a bodyless upstream image status without synthesizing a body", async () => { + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + const requestUrl = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; + if (new URL(requestUrl).hostname === "chatgpt.com") { + return new Response(null, { status: 204 }); + } + return originalFetch(input, init); + }) as typeof fetch; + saveConfig(forwardConfig()); + + const server = startServer(0); + try { + const response = await fetch(new URL("/v1/images/generations", server.url), { + method: "POST", + headers: { + "content-type": "application/json", + authorization: `Bearer ${DIRECT_CHATGPT_TOKEN}`, + "chatgpt-account-id": "acct-123", + }, + body: JSON.stringify({ prompt: "a cat", model: "gpt-image-2" }), + }); + expect(response.status).toBe(204); + expect((await response.arrayBuffer()).byteLength).toBe(0); + } finally { + await server.stop(true); + } +}); + +test("records an oversized forward response status before returning the size error", async () => { + let upstreamCanceled = false; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + const requestUrl = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; + if (new URL(requestUrl).hostname === "chatgpt.com") { + return new Response(new ReadableStream({ + cancel() { upstreamCanceled = true; }, + }), { + status: 429, + headers: { + "content-length": String(IMAGES_RESPONSE_MAX_BYTES + 1), + "content-type": "application/json", + "retry-after": "60", + }, + }); + } + return originalFetch(input, init); + }) as typeof fetch; + saveConfig({ + ...forwardConfig(), + providers: { openai: { ...canonicalOpenAiProvider, codexAccountMode: "pool" } }, + codexAccounts: [ + { id: "main", email: "main@example.test", isMain: true }, + { id: "pool-a", email: "pool@example.test", isMain: false, chatgptAccountId: "acct-pool-a" }, + ], + activeCodexAccountId: "pool-a", + } as OcxConfig); + saveCodexAccountCredential("pool-a", { + accessToken: "pool-access-token", + refreshToken: "pool-refresh-token", + expiresAt: Date.now() + 3_600_000, + chatgptAccountId: "acct-pool-a", + }); + + const server = startServer(0); + try { + const response = await fetch(new URL("/v1/images/generations", server.url), { + method: "POST", + headers: { "content-type": "application/json", authorization: "Bearer caller-token" }, + body: JSON.stringify({ prompt: "a cat", model: "gpt-image-2" }), + }); + expect(response.status).toBe(502); + expect(((await response.json()) as { error: { message: string } }).error.message) + .toContain("response too large"); + expect(upstreamCanceled).toBe(true); + expect(getCodexUpstreamHealth("pool-a")).toMatchObject({ lastFailureStatus: 429 }); + } finally { + await server.stop(true); + } +}); + +test("an image body that stalls after headers retains the 504 deadline", async () => { + let upstreamCanceled = false; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + const requestUrl = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; + if (new URL(requestUrl).hostname === "chatgpt.com") { + return new Response(new ReadableStream({ + cancel() { upstreamCanceled = true; }, + }), { headers: { "content-type": "application/json" } }); + } + return originalFetch(input, init); + }) as typeof fetch; + saveConfig({ ...forwardConfig(), images: { timeoutMs: 50 } } as OcxConfig); + + const server = startServer(0); + try { + const response = await fetch(new URL("/v1/images/generations", server.url), { + method: "POST", + headers: { + "content-type": "application/json", + authorization: `Bearer ${DIRECT_CHATGPT_TOKEN}`, + "chatgpt-account-id": "acct-123", + }, + body: JSON.stringify({ prompt: "a cat", model: "gpt-image-2" }), + }); + expect(response.status).toBe(504); + expect(((await response.json()) as { error: { message: string } }).error.message).toContain("timed out"); + expect(upstreamCanceled).toBe(true); + } finally { + await server.stop(true); + } +}, 5_000); + +test("a client abort during image body reading maps to 499 and cancels upstream", async () => { + let markReaderAttached: (() => void) | undefined; + const readerAttached = new Promise(resolve => { markReaderAttached = resolve; }); + let pulls = 0; + let upstreamCanceled = false; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + const requestUrl = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; + if (new URL(requestUrl).hostname === "chatgpt.com") { + return new Response(new ReadableStream({ + pull(controller) { + pulls++; + if (pulls === 1) { + // Fill the initial queue. The second pull only happens after the + // handler's reader consumes this chunk, proving reader attachment. + controller.enqueue(new Uint8Array([0x01])); + return; + } + markReaderAttached?.(); + }, + cancel() { upstreamCanceled = true; }, + }), { headers: { "content-type": "application/json" } }); + } + return originalFetch(input, init); + }) as typeof fetch; + + const parent = new AbortController(); + const request = new Request("http://127.0.0.1/v1/images/generations", { + method: "POST", + headers: { + "content-type": "application/json", + authorization: `Bearer ${DIRECT_CHATGPT_TOKEN}`, + "chatgpt-account-id": "acct-123", + }, + body: JSON.stringify({ prompt: "a cat", model: "gpt-image-2" }), + signal: parent.signal, + }); + const reading = handleImages(request, forwardConfig(), "generations", { model: "", provider: "" }); + await readerAttached; + parent.abort(new Error("client stopped")); + + const response = await reading; + expect(response.status).toBe(499); + expect(upstreamCanceled).toBe(true); +}); + +test("a client abort before the image reader attaches cancels the untouched upstream body", async () => { + const parent = new AbortController(); + let upstreamCanceled = false; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + const requestUrl = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; + if (new URL(requestUrl).hostname === "chatgpt.com") { + const response = new Response(new ReadableStream({ + cancel() { upstreamCanceled = true; }, + }), { headers: { "content-type": "application/json" } }); + parent.abort(new Error("client stopped before body read")); + return response; + } + return originalFetch(input, init); + }) as typeof fetch; + + const request = new Request("http://127.0.0.1/v1/images/generations", { + method: "POST", + headers: { + "content-type": "application/json", + authorization: `Bearer ${DIRECT_CHATGPT_TOKEN}`, + "chatgpt-account-id": "acct-123", + }, + body: JSON.stringify({ prompt: "a cat", model: "gpt-image-2" }), + signal: parent.signal, + }); + const response = await handleImages(request, forwardConfig(), "generations", { model: "", provider: "" }); + expect(response.status).toBe(499); + expect(upstreamCanceled).toBe(true); +}); + test("a hung upstream times out with 504 after config.images.timeoutMs", async () => { const upstream = Bun.serve({ port: 0, diff --git a/tests/server-search.test.ts b/tests/server-search.test.ts index 2a1e2f680..958b85dc4 100644 --- a/tests/server-search.test.ts +++ b/tests/server-search.test.ts @@ -18,6 +18,7 @@ import { import { loadConfig, saveConfig } from "../src/config"; import { startServer } from "../src/server"; import { clearRequestLogsForTests, getRequestLogEntries } from "../src/server/request-log"; +import { handleSearch, SEARCH_RESPONSE_MAX_BYTES } from "../src/server/search"; import type { OcxConfig } from "../src/types"; import { fakeChatGptJwt } from "./helpers/fake-chatgpt-jwt"; import { installIsolatedCodexHome, type IsolatedCodexHome } from "./helpers/isolated-codex-home"; @@ -451,6 +452,181 @@ test("relays search upstream error status and body verbatim", async () => { } }); +test("relays arbitrary search response bytes and content type verbatim", async () => { + const payload = new Uint8Array([0x00, 0xff, 0x80, 0xc3, 0x28]); + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + const requestUrl = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; + if (new URL(requestUrl).hostname === "chatgpt.com") { + return new Response(payload, { status: 418, headers: { "content-type": "application/octet-stream" } }); + } + return originalFetch(input, init); + }) as typeof fetch; + saveConfig(forwardConfig()); + + const server = startServer(0); + try { + const response = await fetch(new URL("/v1/alpha/search", server.url), { + method: "POST", + headers: { + "content-type": "application/json", + authorization: `Bearer ${DIRECT_CHATGPT_TOKEN}`, + "chatgpt-account-id": "acct-123", + }, + body: JSON.stringify({ id: "search-session", model: "gpt-test" }), + }); + expect(response.status).toBe(418); + expect(response.headers.get("content-type")).toContain("application/octet-stream"); + expect(Array.from(new Uint8Array(await response.arrayBuffer()))).toEqual(Array.from(payload)); + } finally { + await server.stop(true); + } +}); + +test("cancels an oversized streaming search response without draining the stream", async () => { + const cap = SEARCH_RESPONSE_MAX_BYTES; + let upstreamCanceled = false; + let tailPulled = false; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + const requestUrl = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; + if (new URL(requestUrl).hostname !== "chatgpt.com") return originalFetch(input, init); + const chunks = [ + new Uint8Array(cap), + new Uint8Array([0x01]), + new Uint8Array([0x7f]), + new Uint8Array([0x7e]), + ]; + return new Response(new ReadableStream({ + pull(controller) { + const chunk = chunks.shift(); + if (!chunk) return controller.close(); + if (chunk.byteLength === 1 && chunk[0] === 0x7e) tailPulled = true; + controller.enqueue(chunk); + }, + cancel() { upstreamCanceled = true; }, + }), { headers: { "content-type": "application/json" } }); + }) as typeof fetch; + saveConfig(forwardConfig()); + + const server = startServer(0); + try { + const response = await fetch(new URL("/v1/alpha/search", server.url), { + method: "POST", + headers: { + "content-type": "application/json", + authorization: `Bearer ${DIRECT_CHATGPT_TOKEN}`, + "chatgpt-account-id": "acct-123", + }, + body: JSON.stringify({ id: "search-session", model: "gpt-test" }), + }); + expect(response.status).toBe(502); + expect(((await response.json()) as { error: { message: string } }).error.message) + .toContain("search response too large"); + expect(upstreamCanceled).toBe(true); + // WHATWG streams may prefetch one queued chunk, but cancellation must stop further draining. + expect(tailPulled).toBe(false); + } finally { + await server.stop(true); + } +}); + +test("a search body that stalls after headers retains the total 504 deadline", async () => { + let upstreamCanceled = false; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + const requestUrl = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; + if (new URL(requestUrl).hostname === "chatgpt.com") { + return new Response(new ReadableStream({ + cancel() { upstreamCanceled = true; }, + }), { headers: { "content-type": "application/json" } }); + } + return originalFetch(input, init); + }) as typeof fetch; + saveConfig({ ...forwardConfig(), search: { timeoutMs: 50 } } as OcxConfig); + + const server = startServer(0); + try { + const response = await fetch(new URL("/v1/alpha/search", server.url), { + method: "POST", + headers: { + "content-type": "application/json", + authorization: `Bearer ${DIRECT_CHATGPT_TOKEN}`, + "chatgpt-account-id": "acct-123", + }, + body: JSON.stringify({ id: "search-session", model: "gpt-test" }), + }); + expect(response.status).toBe(504); + expect(((await response.json()) as { error: { message: string } }).error.message).toContain("timed out"); + expect(upstreamCanceled).toBe(true); + } finally { + await server.stop(true); + } +}, 5_000); + +test("a client abort during search body reading maps to 499 and cancels upstream", async () => { + let markBodyStarted: (() => void) | undefined; + const bodyStarted = new Promise(resolve => { markBodyStarted = resolve; }); + let upstreamCanceled = false; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + const requestUrl = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; + if (new URL(requestUrl).hostname === "chatgpt.com") { + return new Response(new ReadableStream({ + pull() { markBodyStarted?.(); }, + cancel() { upstreamCanceled = true; }, + }), { headers: { "content-type": "application/json" } }); + } + return originalFetch(input, init); + }) as typeof fetch; + + const parent = new AbortController(); + const request = new Request("http://127.0.0.1/v1/alpha/search", { + method: "POST", + headers: { + "content-type": "application/json", + authorization: `Bearer ${DIRECT_CHATGPT_TOKEN}`, + "chatgpt-account-id": "acct-123", + }, + body: JSON.stringify({ id: "search-session", model: "gpt-test" }), + signal: parent.signal, + }); + const reading = handleSearch(request, forwardConfig(), { model: "", provider: "" }); + await bodyStarted; + parent.abort(new Error("client stopped")); + + const response = await reading; + expect(response.status).toBe(499); + expect(upstreamCanceled).toBe(true); +}); + +test("a client abort before the search reader attaches cancels the untouched upstream body", async () => { + const parent = new AbortController(); + let upstreamCanceled = false; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + const requestUrl = typeof input === "string" ? input : input instanceof URL ? input.toString() : input.url; + if (new URL(requestUrl).hostname === "chatgpt.com") { + const response = new Response(new ReadableStream({ + cancel() { upstreamCanceled = true; }, + }), { headers: { "content-type": "application/json" } }); + parent.abort(new Error("client stopped before body read")); + return response; + } + return originalFetch(input, init); + }) as typeof fetch; + + const request = new Request("http://127.0.0.1/v1/alpha/search", { + method: "POST", + headers: { + "content-type": "application/json", + authorization: `Bearer ${DIRECT_CHATGPT_TOKEN}`, + "chatgpt-account-id": "acct-123", + }, + body: JSON.stringify({ id: "search-session", model: "gpt-test" }), + signal: parent.signal, + }); + + const response = await handleSearch(request, forwardConfig(), { model: "", provider: "" }); + expect(response.status).toBe(499); + expect(upstreamCanceled).toBe(true); +}); + test("a hung search upstream times out with 504 after config.search.timeoutMs", async () => { const upstream = Bun.serve({ port: 0,