diff --git a/src/codex/quota-rejection.ts b/src/codex/quota-rejection.ts new file mode 100644 index 000000000..e0a2e6fad --- /dev/null +++ b/src/codex/quota-rejection.ts @@ -0,0 +1,224 @@ +import { readBoundedResponseBody } from "../lib/bounded-body"; + +const RESET_ELIGIBLE_CODE_VALUES = [ + "usage_limit_exceeded", + "insufficient_quota", +] as const; + +export type CodexResetEligibleExhaustionCode = + (typeof RESET_ELIGIBLE_CODE_VALUES)[number]; + +export type CodexPreStreamRejectionKind = + | "reset-eligible-exhaustion" + | "generic-rate-limit" + | "unverified-billing-or-quota" + | "transient-server-error" + | "authentication-error" + | "permission-error" + | "other"; + +export interface CodexPreStreamRejection { + kind: CodexPreStreamRejectionKind; + status: number; + alternateRetryEligible: boolean; + resetCreditEligible: boolean; + semanticCode?: CodexResetEligibleExhaustionCode; +} + +const RESET_ELIGIBLE_CODES: ReadonlySet = new Set(RESET_ELIGIBLE_CODE_VALUES); + +const TRANSIENT_SERVER_STATUSES = new Set([500, 502, 503, 504, 520, 521, 522]); +const JSON_NUMBER_PATTERN = /-?(?:0|[1-9]\d*)(?:\.\d+)?(?:[eE][+-]?\d+)?/y; + +function rejection( + status: number, + kind: CodexPreStreamRejectionKind, + options: { + alternateRetryEligible?: boolean; + semanticCode?: CodexResetEligibleExhaustionCode; + } = {}, +): CodexPreStreamRejection { + return { + kind, + status, + alternateRetryEligible: options.alternateRetryEligible === true, + resetCreditEligible: options.semanticCode !== undefined, + ...(options.semanticCode ? { semanticCode: options.semanticCode } : {}), + }; +} + +function hasOwnField(container: Record, field: string): boolean { + return Object.prototype.hasOwnProperty.call(container, field); +} + +type JsonScanResult = { + next: number; + duplicate: boolean; +}; + +function skipJsonWhitespace(text: string, index: number): number { + while (index < text.length && /[\t\n\r ]/.test(text[index] ?? "")) index += 1; + return index; +} + +function scanJsonStringEnd(text: string, index: number): number { + if (text[index] !== '"') throw new SyntaxError("expected JSON string"); + for (let cursor = index + 1; cursor < text.length; cursor += 1) { + const char = text[cursor]; + if (char === '"') return cursor + 1; + if (char === "\\") cursor += 1; + } + throw new SyntaxError("unterminated JSON string"); +} + +function scanJsonValue(text: string, index: number): JsonScanResult { + const start = skipJsonWhitespace(text, index); + if (text[start] === "{") return scanJsonObject(text, start); + if (text[start] === "[") return scanJsonArray(text, start); + if (text[start] === '"') return { next: scanJsonStringEnd(text, start), duplicate: false }; + + for (const literal of ["true", "false", "null"]) { + if (text.startsWith(literal, start)) { + return { next: start + literal.length, duplicate: false }; + } + } + JSON_NUMBER_PATTERN.lastIndex = start; + const number = JSON_NUMBER_PATTERN.exec(text); + if (!number) throw new SyntaxError("expected JSON value"); + return { next: start + number[0].length, duplicate: false }; +} + +function scanJsonObject(text: string, index: number): JsonScanResult { + const keys = new Set(); + let duplicate = false; + let cursor = skipJsonWhitespace(text, index + 1); + if (text[cursor] === "}") return { next: cursor + 1, duplicate: false }; + + while (cursor < text.length) { + const keyEnd = scanJsonStringEnd(text, cursor); + const key = JSON.parse(text.slice(cursor, keyEnd)) as unknown; + if (typeof key !== "string") throw new SyntaxError("invalid JSON object key"); + if (keys.has(key)) duplicate = true; + keys.add(key); + + cursor = skipJsonWhitespace(text, keyEnd); + if (text[cursor] !== ":") throw new SyntaxError("expected JSON object colon"); + const value = scanJsonValue(text, cursor + 1); + duplicate ||= value.duplicate; + cursor = skipJsonWhitespace(text, value.next); + if (text[cursor] === "}") return { next: cursor + 1, duplicate }; + if (text[cursor] !== ",") throw new SyntaxError("expected JSON object separator"); + cursor = skipJsonWhitespace(text, cursor + 1); + } + throw new SyntaxError("unterminated JSON object"); +} + +function scanJsonArray(text: string, index: number): JsonScanResult { + let duplicate = false; + let cursor = skipJsonWhitespace(text, index + 1); + if (text[cursor] === "]") return { next: cursor + 1, duplicate: false }; + + while (cursor < text.length) { + const value = scanJsonValue(text, cursor); + duplicate ||= value.duplicate; + cursor = skipJsonWhitespace(text, value.next); + if (text[cursor] === "]") return { next: cursor + 1, duplicate }; + if (text[cursor] !== ",") throw new SyntaxError("expected JSON array separator"); + cursor = skipJsonWhitespace(text, cursor + 1); + } + throw new SyntaxError("unterminated JSON array"); +} + +function isUnsafeJsonDocument(text: string): boolean { + try { + const result = scanJsonValue(text, 0); + return result.duplicate || skipJsonWhitespace(text, result.next) !== text.length; + } catch { + // Scanner disagreement is untrusted input, just like JSON.parse failure. + return true; + } +} + +function exactResetEligibleCode( + container: Record, +): CodexResetEligibleExhaustionCode | undefined { + const hasCode = hasOwnField(container, "code"); + const hasType = hasOwnField(container, "type"); + if (!hasCode && !hasType) return undefined; + + const code = hasCode ? container.code : undefined; + const type = hasType ? container.type : undefined; + if ((hasCode && typeof code !== "string") || (hasType && typeof type !== "string")) { + return undefined; + } + if (hasCode && hasType && code !== type) return undefined; + + const value = hasCode ? code : type; + if (typeof value !== "string") return undefined; + return RESET_ELIGIBLE_CODES.has(value as CodexResetEligibleExhaustionCode) + ? value as CodexResetEligibleExhaustionCode + : undefined; +} + +function structuredResetEligibleCode(payload: unknown): CodexResetEligibleExhaustionCode | undefined { + if (!payload || typeof payload !== "object" || Array.isArray(payload)) return undefined; + const root = payload as Record; + const hasRootDiscriminator = hasOwnField(root, "code") || hasOwnField(root, "type"); + + if (!hasOwnField(root, "error")) return exactResetEligibleCode(root); + if (hasRootDiscriminator) return undefined; + + const nested = root.error; + if (!nested || typeof nested !== "object" || Array.isArray(nested)) return undefined; + return exactResetEligibleCode(nested as Record); +} + +async function resetEligibleCodeFromResponse( + response: Response, + signal?: AbortSignal, +): Promise { + try { + const body = await readBoundedResponseBody(response.clone(), { signal, fatalUtf8: true }); + if (!body.displaySafe || body.truncated || !body.text.trim()) return undefined; + const payload = JSON.parse(body.text) as unknown; + // JSON.parse silently keeps the last duplicate key, making contradictory + // payloads order-dependent. Reject any duplicate at any object depth. + if (isUnsafeJsonDocument(body.text)) return undefined; + return structuredResetEligibleCode(payload); + } catch { + // Classification must fail closed. A malformed, oversized, consumed, or + // cancelled body cannot authorize an irreversible reset-credit operation. + return undefined; + } +} + +/** + * Classify an upstream Codex rejection before any response event is exposed. + * + * Only an exact structured exhaustion code on HTTP 429/402 is reset-eligible. + * Status alone and message text are intentionally insufficient. The broad + * alternate-account retry remains eligible for 429/402 to preserve #584. + */ +export async function classifyCodexPreStreamRejection( + response: Response, + options: { signal?: AbortSignal } = {}, +): Promise { + const status = response.status; + if (status === 401) return rejection(status, "authentication-error"); + if (status === 403) return rejection(status, "permission-error"); + if (TRANSIENT_SERVER_STATUSES.has(status)) return rejection(status, "transient-server-error"); + if (status !== 429 && status !== 402) return rejection(status, "other"); + + const semanticCode = await resetEligibleCodeFromResponse(response, options.signal); + if (semanticCode) { + return rejection(status, "reset-eligible-exhaustion", { + alternateRetryEligible: true, + semanticCode, + }); + } + return rejection( + status, + status === 429 ? "generic-rate-limit" : "unverified-billing-or-quota", + { alternateRetryEligible: true }, + ); +} diff --git a/src/lib/bounded-body.ts b/src/lib/bounded-body.ts index 745833c79..2cb7ea39b 100644 --- a/src/lib/bounded-body.ts +++ b/src/lib/bounded-body.ts @@ -7,6 +7,12 @@ export const BOUNDED_BODY_TIMEOUT_MS = 5_000; export interface BoundedBodyOptions { /** Abort the read with this signal. Its reason is rethrown by identity. */ signal?: AbortSignal; + /** + * Reject the returned promise with TypeError on malformed or truncated UTF-8 + * instead of replacing invalid bytes, including during timeout-path flushes. + * Reader cancellation and lock release still run. Defaults to false. + */ + fatalUtf8?: boolean; /** * Byte ceiling for retained body data. Defaults to BOUNDED_BODY_MAX_BYTES (64 KiB), * which suits error bodies; callers materializing whole success payloads (e.g. a @@ -75,8 +81,8 @@ function cancelWithoutWaiting(reader: ReadableStreamDefaultReader, r } } -function decodeUtf8(chunks: readonly Uint8Array[]): string { - const decoder = new TextDecoder(); +function decodeUtf8(chunks: readonly Uint8Array[], fatal: boolean): string { + const decoder = new TextDecoder("utf-8", { fatal }); let text = ""; for (const chunk of chunks) text += decoder.decode(chunk, { stream: true }); // Flush an incomplete trailing UTF-8 sequence deterministically. @@ -158,7 +164,7 @@ export async function readBoundedResponseBody( "TimeoutError", ); return { - text: decodeUtf8([retained.subarray(0, retainedBytes)]), + text: decodeUtf8([retained.subarray(0, retainedBytes)], options.fatalUtf8 === true), truncated: true, timedOut: true, totalTimedOut: outcome === TOTAL_TIMEOUT, @@ -171,7 +177,7 @@ export async function readBoundedResponseBody( const { value, done } = outcome as ReadableStreamReadResult; if (done) { return { - text: decodeUtf8([retained.subarray(0, retainedBytes)]), + text: decodeUtf8([retained.subarray(0, retainedBytes)], options.fatalUtf8 === true), truncated: false, timedOut: false, totalTimedOut: false, diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index 845c4a338..e4573637e 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -256,8 +256,8 @@ async function shouldRetryCodexPoolAccountModel400( } /** Pre-stream quota/billing rejections that warrant one alternate-account attempt (#584). */ -function shouldRetryCodexPoolAccountQuota(response: Response): boolean { - return response.status === 429 || response.status === 402; +export function shouldRetryCodexPoolAccountQuota(response: Response): boolean { + return response.status === 402 || response.status === 429; } interface CodexPoolAccountRetryArgs { diff --git a/tests/codex-quota-rejection.test.ts b/tests/codex-quota-rejection.test.ts new file mode 100644 index 000000000..05dc25946 --- /dev/null +++ b/tests/codex-quota-rejection.test.ts @@ -0,0 +1,301 @@ +import { describe, expect, test } from "bun:test"; +import { classifyCodexPreStreamRejection } from "../src/codex/quota-rejection"; +import { BOUNDED_BODY_MAX_BYTES } from "../src/lib/bounded-body"; +import { shouldRetryCodexPoolAccountQuota } from "../src/server/responses/core"; + +function jsonRejection(status: number, error: Record): Response { + return Response.json({ error }, { status }); +} + +function jsonPayload(status: number, payload: Record): Response { + return Response.json(payload, { status }); +} + +describe("Codex pre-stream quota rejection classification", () => { + test.each([ + [402, true], + [429, true], + [400, false], + [503, false], + ])("selects pool-account retries synchronously for HTTP %i", (status, expected) => { + expect(shouldRetryCodexPoolAccountQuota(new Response(null, { status }))).toBe(expected); + }); + + test.each([ + [429, "nested code", { error: { code: "usage_limit_exceeded" } }, "usage_limit_exceeded"], + [429, "nested type", { error: { type: "insufficient_quota" } }, "insufficient_quota"], + [402, "nested code", { error: { code: "insufficient_quota" } }, "insufficient_quota"], + [429, "root code", { code: "usage_limit_exceeded" }, "usage_limit_exceeded"], + [402, "root type", { type: "insufficient_quota" }, "insufficient_quota"], + [429, "matching code and type", { + error: { code: "usage_limit_exceeded", type: "usage_limit_exceeded" }, + }, "usage_limit_exceeded"], + ] as const)("accepts exact %s reset-eligible exhaustion on HTTP %i", async ( + status, + _schema, + payload, + code, + ) => { + const result = await classifyCodexPreStreamRejection(jsonPayload(status, payload)); + expect(result).toEqual({ + kind: "reset-eligible-exhaustion", + status, + alternateRetryEligible: true, + resetCreditEligible: true, + semanticCode: code, + }); + }); + + test.each([ + ["leading and trailing whitespace", { error: { code: " usage_limit_exceeded " } }], + ["uppercase", { error: { code: "USAGE_LIMIT_EXCEEDED" } }], + ["mixed case", { type: "Insufficient_Quota" }], + ["trailing whitespace at the root", { code: "insufficient_quota " }], + ] as const)("rejects the %s near-miss", async (_case, payload) => { + const result = await classifyCodexPreStreamRejection(jsonPayload(429, payload)); + expect(result).toMatchObject({ + kind: "generic-rate-limit", + alternateRetryEligible: true, + resetCreditEligible: false, + }); + expect(result).not.toHaveProperty("semanticCode"); + }); + + test.each([ + ["primitive error with a root code", { + error: "opaque", + code: "usage_limit_exceeded", + }], + ["matching root and nested codes", { + code: "usage_limit_exceeded", + error: { code: "usage_limit_exceeded" }, + }], + ["unknown root and eligible nested codes", { + code: "unknown", + error: { code: "usage_limit_exceeded" }, + }], + ["eligible root code with an empty nested error", { + code: "usage_limit_exceeded", + error: {}, + }], + ] as const)("fails closed for ambiguous schemas: %s", async (_case, payload) => { + const result = await classifyCodexPreStreamRejection(jsonPayload(429, payload)); + expect(result).toMatchObject({ + kind: "generic-rate-limit", + alternateRetryEligible: true, + resetCreditEligible: false, + }); + expect(result).not.toHaveProperty("semanticCode"); + }); + + test.each([ + ["root", '{"code":"rate_limit_error","code":"usage_limit_exceeded"}'], + ["nested", '{"error":{"code":"rate_limit_error","code":"usage_limit_exceeded"}}'], + ] as const)("fails closed for duplicate keys in a %s object", async (_case, body) => { + const response = new Response(body, { + status: 429, + headers: { "content-type": "application/json" }, + }); + const result = await classifyCodexPreStreamRejection(response); + expect(result).toMatchObject({ + kind: "generic-rate-limit", + alternateRetryEligible: true, + resetCreditEligible: false, + }); + expect(result).not.toHaveProperty("semanticCode"); + expect(await response.text()).toBe(body); + }); + + test.each([ + ["nested unknown code and eligible type", { + error: { code: "unknown", type: "insufficient_quota" }, + }], + ["root eligible code and unrelated type", { + code: "usage_limit_exceeded", + type: "rate_limit_error", + }], + ["two different eligible values", { + error: { code: "usage_limit_exceeded", type: "insufficient_quota" }, + }], + ["eligible code and non-string type", { + error: { code: "usage_limit_exceeded", type: null }, + }], + ] as const)("fails closed for code/type disagreement: %s", async (_case, payload) => { + const result = await classifyCodexPreStreamRejection(jsonPayload(429, payload)); + expect(result).toMatchObject({ + kind: "generic-rate-limit", + alternateRetryEligible: true, + resetCreditEligible: false, + }); + expect(result).not.toHaveProperty("semanticCode"); + }); + + test("keeps a generic 429 with Retry-After out of reset-credit eligibility", async () => { + const response = jsonRejection(429, { + type: "rate_limit_error", + code: "rate_limit_exceeded", + message: "try again later", + }); + response.headers.set("retry-after", "60"); + await expect(classifyCodexPreStreamRejection(response)).resolves.toEqual({ + kind: "generic-rate-limit", + status: 429, + alternateRetryEligible: true, + resetCreditEligible: false, + }); + }); + + test("does not trust reset-eligible words found only in a message", async () => { + const result = await classifyCodexPreStreamRejection(jsonRejection(429, { + type: "rate_limit_error", + message: "usage_limit_exceeded: insufficient_quota", + })); + expect(result).toMatchObject({ + kind: "generic-rate-limit", + alternateRetryEligible: true, + resetCreditEligible: false, + }); + }); + + test("fails closed when malformed UTF-8 would otherwise be replaced", async () => { + const prefix = new TextEncoder().encode( + '{"error":{"code":"usage_limit_exceeded","message":"', + ); + const suffix = new TextEncoder().encode('"}}'); + const bytes = new Uint8Array(prefix.length + 1 + suffix.length); + bytes.set(prefix); + bytes[prefix.length] = 0xff; + bytes.set(suffix, prefix.length + 1); + const response = new Response(bytes, { + status: 429, + headers: { "content-type": "application/json" }, + }); + + const result = await classifyCodexPreStreamRejection(response); + + expect(result).toMatchObject({ + kind: "generic-rate-limit", + alternateRetryEligible: true, + resetCreditEligible: false, + }); + expect(result).not.toHaveProperty("semanticCode"); + expect(new Uint8Array(await response.arrayBuffer())).toEqual(bytes); + }); + + test("fails closed for malformed JSON while preserving broad 429 failover", async () => { + const response = new Response('{"error":', { + status: 429, + headers: { "content-type": "application/json" }, + }); + const result = await classifyCodexPreStreamRejection(response); + expect(result).toMatchObject({ + kind: "generic-rate-limit", + alternateRetryEligible: true, + resetCreditEligible: false, + }); + expect(await response.text()).toBe('{"error":'); + }); + + test("fails closed for an empty response body", async () => { + const result = await classifyCodexPreStreamRejection(new Response(null, { status: 429 })); + + expect(result).toMatchObject({ + kind: "generic-rate-limit", + alternateRetryEligible: true, + resetCreditEligible: false, + }); + expect(result).not.toHaveProperty("semanticCode"); + }); + + test("fails closed for an oversized structured body", async () => { + const response = jsonRejection(429, { + code: "usage_limit_exceeded", + padding: "x".repeat(BOUNDED_BODY_MAX_BYTES), + }); + + const result = await classifyCodexPreStreamRejection(response); + + expect(result).toMatchObject({ + kind: "generic-rate-limit", + alternateRetryEligible: true, + resetCreditEligible: false, + }); + expect(result).not.toHaveProperty("semanticCode"); + }); + + test("fails closed when a structured body is truncated by a transport error", async () => { + const response = new Response(new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode('{"error":{"code":"usage_limit_exceeded"')); + controller.error(new TypeError("transport truncated")); + }, + }), { + status: 429, + headers: { "content-type": "application/json" }, + }); + + const result = await classifyCodexPreStreamRejection(response); + + expect(result).toMatchObject({ + kind: "generic-rate-limit", + alternateRetryEligible: true, + resetCreditEligible: false, + }); + expect(result).not.toHaveProperty("semanticCode"); + }); + + test("fails closed when the structured body was already consumed", async () => { + const response = jsonRejection(429, { code: "usage_limit_exceeded" }); + await response.text(); + + const result = await classifyCodexPreStreamRejection(response); + + expect(result).toMatchObject({ + kind: "generic-rate-limit", + alternateRetryEligible: true, + resetCreditEligible: false, + }); + expect(result).not.toHaveProperty("semanticCode"); + }); + + test.each([ + [503, "transient-server-error"], + [401, "authentication-error"], + [403, "permission-error"], + [400, "other"], + ] as const)("separates non-eligible HTTP %i as %s", async (status, kind) => { + const result = await classifyCodexPreStreamRejection(jsonRejection(status, { + code: "usage_limit_exceeded", + })); + expect(result).toMatchObject({ + kind, + alternateRetryEligible: false, + resetCreditEligible: false, + }); + }); + + test("classifies an unverified 402 without authorizing a reset credit", async () => { + await expect(classifyCodexPreStreamRejection(jsonRejection(402, { + code: "billing_error", + }))).resolves.toEqual({ + kind: "unverified-billing-or-quota", + status: 402, + alternateRetryEligible: true, + resetCreditEligible: false, + }); + }); + + test("an aborted body read fails closed and leaves generic failover eligible", async () => { + const controller = new AbortController(); + controller.abort(); + const result = await classifyCodexPreStreamRejection( + jsonRejection(429, { code: "usage_limit_exceeded" }), + { signal: controller.signal }, + ); + expect(result).toMatchObject({ + kind: "generic-rate-limit", + alternateRetryEligible: true, + resetCreditEligible: false, + }); + }); +});