diff --git a/docs-site/src/content/docs/reference/cli/providers-accounts.md b/docs-site/src/content/docs/reference/cli/providers-accounts.md index dd6c48f2a..ac8047a1a 100644 --- a/docs-site/src/content/docs/reference/cli/providers-accounts.md +++ b/docs-site/src/content/docs/reference/cli/providers-accounts.md @@ -164,6 +164,11 @@ Direct mode keeps using the caller-owned/native main credential. Usage-based pro recovery may later select another eligible Pool account. Those recovery paths remain active when usage-based switching is off. OpenCodex replays the conversation after an account change, but the provider-side prompt cache may be cold. Unknown providers or ids exit 1. +Before substantive output, an upstream model-capacity rejection rotates through each eligible Pool +account once without persisting a cooldown. A rejected account does not create or refresh thread +affinity; only the account whose response is accepted becomes the thread binding. Exhaustion returns +the first capacity error, and the next request starts with a fresh eligible set and no new affinity from +the rejected attempts. Direct mode, exact account selectors, and failures after output do not rotate. On a **401/403**, App login clears that account's process-local affinity and requires reauthentication. On a **429**, opencodex honors `Retry-After`, starts the account cooldown, clears affinity, and may rotate the request to another eligible Pool account. These failure transitions remain active with diff --git a/docs-site/src/content/docs/reference/configuration/providers.md b/docs-site/src/content/docs/reference/configuration/providers.md index 789adce41..e4a81bb44 100644 --- a/docs-site/src/content/docs/reference/configuration/providers.md +++ b/docs-site/src/content/docs/reference/configuration/providers.md @@ -150,6 +150,16 @@ alternate account in the same request, even when usage-based proactive switching changes preserve and replay the conversation context, but provider-side prompt-cache reuse across accounts is not guaranteed and the cache may need to warm again. +Before any text, reasoning, tool call, or other model output reaches the client, a structured +`server_is_overloaded` / `slow_down` model-capacity rejection (or the standard +`Selected model is at capacity. Please try a different model.` response) tries each remaining +eligible Pool account once. This exclusion is request-local: rejected capacity attempts do not write +account cooldown or health. Thread affinity is not created or refreshed until a non-capacity response +is accepted; after rotation, only the accepted account is bound. If every account returns capacity, +the first response is returned and the next request starts with a fresh eligible set, without any new +affinity from the rejected attempts. Direct mode and exact account selectors never use this rotation, +and capacity after substantive output is never replayed. + On a **401/403**, App login clears that account's process-local affinity and requires reauthentication. On a **429**, opencodex honors `Retry-After`, starts the account cooldown, clears affinity, and may rotate the request to another eligible Pool account. These failure transitions remain active with diff --git a/src/codex/auth-context.ts b/src/codex/auth-context.ts index 8f1416715..1fcf5250b 100644 --- a/src/codex/auth-context.ts +++ b/src/codex/auth-context.ts @@ -21,6 +21,8 @@ import { tryAcquireCodexQuotaProbeLease, tryAcquireCodexQuotaScopeProbeLease, pickAlternateCodexAccount, + pickAlternateCodexAccountExcluding, + previewCodexAccountForRequestDetailed, resolveCodexAccountForThreadDetailed, } from "./routing"; import type { CodexCooldownSource, CodexQuotaScope } from "./routing"; @@ -227,10 +229,14 @@ export function shouldMarkAccountNeedsReauthForCodexAuthFailure(cause: unknown): export interface ResolveCodexAuthContextOptions { excludeAccountId?: string; + /** Request-local exclusion set for bounded multi-account recovery. */ + excludeAccountIds?: ReadonlySet; /** Resolve exactly this account without consulting or mutating Pool selection. */ accountId?: string; /** Final native model selected for this request, used to select its quota group. */ modelId?: string; + /** Preview without selection-state writes until the upstream response is accepted. */ + deferSelectionCommit?: boolean; /** Short reservation converted to turn ownership before native `__main__` token materialization. */ beginCodexAccountSelection?: () => CodexAccountSelectionAdmission | undefined; /** Test-only native credential read seams. */ @@ -253,7 +259,9 @@ export async function resolveCodexAuthContext( ): Promise { const writerGeneration = captureConfigGeneration(); const fixedAccountId = options.accountId; - if (fixedAccountId !== undefined && options.excludeAccountId !== undefined) { + const hasExclusions = options.excludeAccountId !== undefined + || (options.excludeAccountIds?.size ?? 0) > 0; + if (fixedAccountId !== undefined && hasExclusions) { throw new Error("Codex auth context cannot select and exclude an account simultaneously"); } // An explicit namespace binding is stronger than the provider's default mode. It must use the @@ -284,19 +292,31 @@ export async function resolveCodexAuthContext( const threadId = headers.get("x-codex-parent-thread-id"); const resolution = fixedAccountId !== undefined ? { status: "selected" as const, accountId: fixedAccountId } - : options.excludeAccountId + : hasExclusions ? (() => { - const selected = pickAlternateCodexAccount( - config, - options.excludeAccountId!, - Date.now(), - quotaScope, - selectionOptions, - ); + const selected = options.excludeAccountIds + ? pickAlternateCodexAccountExcluding( + config, + options.excludeAccountIds, + [...options.excludeAccountIds].at(-1)!, + Date.now(), + quotaScope, + selectionOptions, + false, + ) + : pickAlternateCodexAccount( + config, + options.excludeAccountId!, + Date.now(), + quotaScope, + selectionOptions, + ); return selected ? { status: "selected" as const, accountId: selected } : { status: "none" as const }; })() + : options.deferSelectionCommit + ? previewCodexAccountForRequestDetailed(threadId, config, Date.now(), quotaScope, selectionOptions) : resolveCodexAccountForThreadDetailed(threadId, config, Date.now(), quotaScope, selectionOptions); if (resolution.status === "expired") throw new CodexThreadAffinityExpiredError(resolution.accountId); const selected = resolution.status === "selected" ? resolution.accountId : null; @@ -309,7 +329,7 @@ export async function resolveCodexAuthContext( // temporary fence rather than misclassifying that credential as invalid. // A configured pool retry/exclusion that finds no alternate preserves its // ordinary pool-auth failure instead of being mislabeled as a main fence. - if (nativeMainTrafficBlocked && !options.excludeAccountId) { + if (nativeMainTrafficBlocked && !hasExclusions) { throw new CodexMainProfileDrainingError(); } throw new CodexPoolAuthenticationError(); diff --git a/src/codex/pool-rotation.ts b/src/codex/pool-rotation.ts index d0d032be0..598bd0040 100644 --- a/src/codex/pool-rotation.ts +++ b/src/codex/pool-rotation.ts @@ -210,6 +210,32 @@ export function peekRoundRobinAccount( return pickRoundRobinFromState(eligibleIds, stickyLimit, scratch, false); } +/** + * Record the account whose response was actually accepted after a deferred pick. + * Completion order is authoritative: concurrent requests each contribute one + * smooth-weighted success for the account that served them. + */ +export function commitRoundRobinAccountSuccess( + poolKey: string, + eligibleIds: readonly string[], + accountId: string, + stickyLimit: number, +): boolean { + if (!eligibleIds.includes(accountId)) return false; + const state = getOrCreateState(poolKey); + if (state.activeKey !== accountId) { + delete state.activeKey; + state.successes = 0; + const total = eligibleIds.length; + for (const id of eligibleIds) { + state.currentWeights.set(id, (state.currentWeights.get(id) ?? 0) + 1); + } + state.currentWeights.set(accountId, (state.currentWeights.get(accountId) ?? 0) - total); + } + notePoolRotationSuccess(poolKey, accountId, stickyLimit); + return true; +} + export function notePoolRotationSuccess( poolKey: string, accountId: string, diff --git a/src/codex/routing.ts b/src/codex/routing.ts index ccdb486ee..7b8448590 100644 --- a/src/codex/routing.ts +++ b/src/codex/routing.ts @@ -8,6 +8,7 @@ import { isCodexAccountUsable, type CodexAccountUsabilityOptions } from "./accou import { isAccountNeedsReauth, markAccountNeedsReauth } from "./account-runtime-state"; import { POOL_KEY_CODEX, + commitRoundRobinAccountSuccess, normalizeAccountPoolStickyLimit, normalizeAccountPoolStrategy, notePoolRotationFailure, @@ -900,16 +901,25 @@ function bindThreadAffinity( pruneLruThreadAffinities(); } +type CodexAccountExclusion = string | ReadonlySet | undefined; + +/** Match a single legacy exclusion or a request-local exclusion set. */ +function isExcludedCodexAccount(exclusion: CodexAccountExclusion, accountId: string): boolean { + return typeof exclusion === "string" + ? exclusion === accountId + : exclusion?.has(accountId) === true; +} + function getEligiblePoolAccounts( config: OcxConfig, - excludeId?: string, + exclusion?: CodexAccountExclusion, now = Date.now(), quotaScope?: CodexQuotaScope, selectionOptions?: CodexAccountUsabilityOptions, ): readonly string[] { const ids = (config.codexAccounts ?? []) .filter(account => isSelectableCodexPoolAccount(account) - && account.id !== excludeId + && !isExcludedCodexAccount(exclusion, account.id) && !isCodexAccountPaused(config, account.id) && !isAccountNeedsReauth(account.id)) .filter(account => getCodexQuotaHealthSnapshot(account.id, quotaScope, now) === null) @@ -919,7 +929,7 @@ function getEligiblePoolAccounts( // The main Codex account is not stored in config.codexAccounts; include it as a // first-class rotation candidate when its read-only token is usable (Option A). if ( - excludeId !== MAIN_CODEX_ACCOUNT_ID + !isExcludedCodexAccount(exclusion, MAIN_CODEX_ACCOUNT_ID) && !isCodexAccountPaused(config, MAIN_CODEX_ACCOUNT_ID) && !isAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID) && getCodexQuotaHealthSnapshot(MAIN_CODEX_ACCOUNT_ID, quotaScope, now) === null @@ -1151,6 +1161,26 @@ export function pickAlternateCodexAccount( now = Date.now(), quotaScope?: CodexQuotaScope, selectionOptions?: CodexAccountUsabilityOptions, +): string | null { + return pickAlternateCodexAccountExcluding( + config, + new Set([excludeId]), + excludeId, + now, + quotaScope, + selectionOptions, + ); +} + +/** Strategy-aware alternate that never revisits an account already tried by this request. */ +export function pickAlternateCodexAccountExcluding( + config: OcxConfig, + excludedIds: ReadonlySet, + afterId: string, + now = Date.now(), + quotaScope?: CodexQuotaScope, + selectionOptions?: CodexAccountUsabilityOptions, + commitRoundRobin = true, ): string | null { const strategy = normalizeAccountPoolStrategy(config.accountPoolStrategy); // The exclusion is passed into eligibility rather than post-filtered off its @@ -1158,14 +1188,21 @@ export function pickAlternateCodexAccount( // tier, the tier walk must be free to descend instead of selecting that tier // and then handing back an empty list. if (strategy === "round-robin") { - const eligible = getEligiblePoolAccounts(config, excludeId, now, quotaScope, selectionOptions); - return pickRoundRobinAccount(codexPoolKeyForScope(quotaScope), eligible, stickyLimitForConfig(config)); + const eligible = getEligiblePoolAccounts(config, excludedIds, now, quotaScope, selectionOptions); + const poolKey = codexPoolKeyForScope(quotaScope); + const stickyLimit = stickyLimitForConfig(config); + return commitRoundRobin + ? pickRoundRobinAccount(poolKey, eligible, stickyLimit) + : peekRoundRobinAccount(poolKey, eligible, stickyLimit); } if (strategy === "fill-first") { - const eligible = getEligiblePoolAccounts(config, excludeId, now, quotaScope, selectionOptions); - return pickNextFillFirstCodexAccount(config, excludeId, eligible, now, selectionOptions); + const eligible = getEligiblePoolAccounts(config, excludedIds, now, quotaScope, selectionOptions); + return pickNextFillFirstCodexAccount(config, afterId, eligible, now, selectionOptions); } - return pickLowestUsageCodexAccount(config, excludeId, now, quotaScope, selectionOptions); + return pickLowestUsageAmong( + config, + getEligiblePoolAccounts(config, excludedIds, now, quotaScope, selectionOptions), + ); } /** Effective active: automatic runtime cursor, else operator/persisted selection. */ @@ -1465,6 +1502,80 @@ export function previewCodexAccountForRequest( return active; } +/** Preview Pool selection; expired-affinity cleanup matches live resolve. */ +export function previewCodexAccountForRequestDetailed( + threadId: string | null, + config: OcxConfig, + now = Date.now(), + quotaScope?: CodexQuotaScope, + selectionOptions?: CodexAccountUsabilityOptions, +): CodexThreadResolution { + const entry = threadId ? getThreadAffinity(threadId, quotaScope) : undefined; + if (entry && isThreadAffinityExpired(entry, now)) { + deleteThreadAffinity(threadId!, quotaScope); + return { status: "expired", accountId: entry.accountId }; + } + const accountId = previewCodexAccountForRequest(threadId, config, now, quotaScope, selectionOptions); + return accountId ? { status: "selected", accountId } : { status: "none" }; +} + +/** + * Commit a deferred Pool selection exactly once, after its response is accepted. + * Capacity-rejected and unavailable attempts remain request-local exclusions. + */ +export function commitCodexAcceptedAccountSelection( + config: OcxConfig, + threadId: string | null, + accountId: string, + modelId: string | undefined, + excludedAccountIds: ReadonlySet, + now = Date.now(), +): void { + const quotaScope = codexQuotaScopeForModel(modelId); + if (!isCodexAccountSelectable(config, accountId, now, quotaScope)) return; + if (!isIndependentCodexQuotaScope(quotaScope)) releaseDrainedCodexAccountPin(config); + + const existingAffinity = threadId ? getThreadAffinity(threadId, quotaScope) : undefined; + const reusedAffinity = existingAffinity?.accountId === accountId + && !isThreadAffinityExpired(existingAffinity, now) + && isThreadAffinityGenerationLive(existingAffinity) + && !shouldFailover(config, accountId, now); + + const strategy = normalizeAccountPoolStrategy(config.accountPoolStrategy); + if (reusedAffinity) { + existingAffinity.lastUsedAt = now; + if (strategy === "quota") { + const threshold = config.autoSwitchThreshold ?? 80; + const usage = threshold > 0 + ? computeCodexUsageScore(getAccountQuota(accountId), getPoolAccountPlan(config, accountId)) + : 0; + const overThreshold = threshold > 0 && !isUnknownUsage(usage) && usage >= threshold; + if (overThreshold || now - existingAffinity.lastReevalAt >= CODEX_THREAD_AFFINITY_REEVAL_INTERVAL_MS) { + existingAffinity.lastReevalAt = now; + } + } + return; + } + + if (strategy === "round-robin") { + const eligible = getEligiblePoolAccounts(config, excludedAccountIds, now, quotaScope); + const committed = commitRoundRobinAccountSuccess( + codexPoolKeyForScope(quotaScope), + eligible, + accountId, + stickyLimitForConfig(config), + ); + if (!committed) return; + if (!isIndependentCodexQuotaScope(quotaScope)) rememberActiveCodexAccount(config, accountId); + } else if (strategy === "fill-first") { + if (!isIndependentCodexQuotaScope(quotaScope)) rememberActiveCodexAccount(config, accountId); + } else if (!isIndependentCodexQuotaScope(quotaScope)) { + setActiveCodexAccount(config, accountId); + } + + if (threadId) bindThreadAffinity(threadId, accountId, now, quotaScope); +} + export function resolveCodexAccountForThreadDetailed( threadId: string | null, config: OcxConfig, diff --git a/src/lib/upstream-retry.ts b/src/lib/upstream-retry.ts index 2ba8adb37..3c705c5fb 100644 --- a/src/lib/upstream-retry.ts +++ b/src/lib/upstream-retry.ts @@ -244,6 +244,10 @@ export interface ResetRetryOptions { export interface TransientRetryOptions extends ResetRetryOptions { /** Test seam: per-attempt slow budget override (defaults to TRANSIENT_RETRY_SLOW_ATTEMPT_MS). */ slowAttemptMs?: number; + /** Optional semantic gate run before retrying a transient HTTP response. */ + prepareTransientRetry?: ( + response: Response, + ) => Promise<{ response: Response; retry: boolean }> | { response: Response; retry: boolean }; } export type UpstreamSendRecovery = "connection-reset" | "transient-5xx"; @@ -367,6 +371,11 @@ export async function fetchWithTransientRetry( if (res.ok || !isTransientUpstreamStatus(res.status)) return res; if (opts.abortSignal?.aborted) return res; if (Date.now() - attemptStart > slowAttemptMs) return res; + if (opts.prepareTransientRetry) { + const prepared = await opts.prepareTransientRetry(res); + res = prepared.response; + if (!prepared.retry) return res; + } console.warn( `[upstream-retry] transient ${res.status}${opts.label ? ` (${opts.label})` : ""} — retrying (${attempt + 2}/${attempts})`, ); diff --git a/src/server/responses/codex-capacity.ts b/src/server/responses/codex-capacity.ts new file mode 100644 index 000000000..39faef995 --- /dev/null +++ b/src/server/responses/codex-capacity.ts @@ -0,0 +1,277 @@ +import { readBoundedResponseBody } from "../../lib/bounded-body"; +import { + BoundedSseFrameBuffer, + joinSseFrameBytes, + MAX_CLIENT_SSE_FRAME_BYTES, +} from "../sse-frame-buffer"; +import { sseDataPayload } from "../relay"; + +export const CODEX_MODEL_CAPACITY_MESSAGE = + "Selected model is at capacity. Please try a different model."; + +const CAPACITY_CODES: ReadonlySet = new Set([ + "server_is_overloaded", + "slow_down", +]); + +type JsonRecord = Record; + +export type CodexCapacityInspection = + | { kind: "capacity"; response: Response } + | { kind: "pass"; response: Response }; + +/** Narrow an unknown JSON value to a non-array object. */ +function record(value: unknown): JsonRecord | null { + return value !== null && typeof value === "object" && !Array.isArray(value) + ? value as JsonRecord + : null; +} + +/** Normalize only whitespace and case for the exact-message compatibility check. */ +function normalizeCapacityMessage(value: string): string { + return value.trim().replace(/\s+/gu, " ").toLocaleLowerCase("en-US"); +} + +/** Match one structured capacity signal or the exact standard message. */ +function hasCapacitySignal(value: unknown): boolean { + if (typeof value === "string") return CAPACITY_CODES.has(value.trim().toLocaleLowerCase("en-US")); + const candidate = record(value); + if (!candidate) return false; + const code = typeof candidate.code === "string" + ? candidate.code.trim().toLocaleLowerCase("en-US") + : ""; + if (CAPACITY_CODES.has(code)) return true; + const type = typeof candidate.type === "string" + ? candidate.type.trim().toLocaleLowerCase("en-US") + : ""; + if (CAPACITY_CODES.has(type)) return true; + const message = typeof candidate.message === "string" ? candidate.message : ""; + return normalizeCapacityMessage(message) + === normalizeCapacityMessage(CODEX_MODEL_CAPACITY_MESSAGE); +} + +/** Exact structured capacity signals, plus the one standard message compatibility case. */ +export function isCodexCapacityPayload(payload: unknown): boolean { + const root = record(payload); + if (!root) return false; + const response = record(root.response); + return [ + root, + root.error, + root.last_error, + response, + response?.error, + response?.last_error, + ].some(hasCapacitySignal); +} + +/** Return whether a Responses envelope represents a terminal failure. */ +function isFailedDocument(payload: unknown): boolean { + const root = record(payload); + if (!root) return false; + const response = record(root.response); + return root.type === "response.failed" + || root.type === "error" + || root.status === "failed" + || response?.status === "failed"; +} + +/** Parse an SSE data payload without treating malformed bytes as retryable. */ +function parsePayload(payload: string | null): unknown | undefined { + if (!payload || payload === "[DONE]") return undefined; + try { + return JSON.parse(payload) as unknown; + } catch { + return undefined; + } +} + +/** Allow only lifecycle events that cannot expose model output. */ +function isSafePreOutputLifecycle(payload: unknown): boolean { + const event = record(payload); + if (!event) return false; + return event.type === "response.created" + || event.type === "response.in_progress" + || event.type === "response.queued" + || event.type === "response.heartbeat"; +} + +/** Rebuild a response while preserving status and non-length headers. */ +function responseWithBody(response: Response, body: BodyInit | null): Response { + const headers = new Headers(response.headers); + headers.delete("content-length"); + return new Response(body, { + status: response.status, + statusText: response.statusText, + headers, + }); +} + +/** Replay inspected bytes before resuming the untouched upstream reader. */ +function replayBufferedThenReader( + buffered: readonly Uint8Array[], + reader: ReadableStreamDefaultReader, + terminalError?: unknown, +): ReadableStream { + let prefixIndex = 0; + return new ReadableStream({ + async pull(controller) { + if (prefixIndex < buffered.length) { + controller.enqueue(buffered[prefixIndex++]!); + return; + } + if (terminalError !== undefined) { + controller.error(terminalError); + return; + } + try { + const next = await reader.read(); + if (next.done) controller.close(); + else controller.enqueue(next.value); + } catch (error) { + controller.error(error); + } + }, + cancel(reason) { + return reader.cancel(reason).catch(() => undefined); + }, + }); +} + +/** Return a committed attempt with every inspected byte restored. */ +function passWithBufferedBody( + response: Response, + buffered: readonly Uint8Array[], + reader: ReadableStreamDefaultReader, + terminalError?: unknown, +): CodexCapacityInspection { + return { + kind: "pass", + response: responseWithBody(response, replayBufferedThenReader(buffered, reader, terminalError)), + }; +} + +/** Inspect bounded lifecycle-only SSE frames until capacity or output commits the attempt. */ +async function inspectSseCapacityBeforeOutput( + response: Response, + signal?: AbortSignal, +): Promise { + const body = response.body; + if (!body) return { kind: "pass", response }; + const reader = body.getReader(); + const framer = new BoundedSseFrameBuffer(MAX_CLIENT_SSE_FRAME_BYTES); + const buffered: Uint8Array[] = []; + const completeFrameParts: Uint8Array[] = []; + const decoder = new TextDecoder(); + let bufferedBytes = 0; + let frameBytes = 0; + + try { + while (true) { + if (signal?.aborted) throw signal.reason; + let next: { done: boolean; value?: Uint8Array }; + try { + next = await reader.read(); + } catch (error) { + framer.dispose(); + return passWithBufferedBody(response, buffered, reader, error); + } + if (next.done) { + framer.dispose(); + return passWithBufferedBody(response, buffered, reader); + } + const chunk = next.value!; + buffered.push(chunk); + bufferedBytes += chunk.byteLength; + + let frames: ReturnType; + try { + frames = framer.feed(chunk); + } catch { + framer.dispose(); + return passWithBufferedBody(response, buffered, reader); + } + + for (const frame of frames) { + completeFrameParts.push(frame.block, frame.delimiter); + frameBytes += frame.block.byteLength + frame.delimiter.byteLength; + if (frameBytes > MAX_CLIENT_SSE_FRAME_BYTES) { + framer.dispose(); + return passWithBufferedBody(response, buffered, reader); + } + const payloadText = sseDataPayload(decoder.decode(frame.block)); + if (payloadText === null) continue; + if (payloadText === "[DONE]") { + framer.dispose(); + return passWithBufferedBody(response, buffered, reader); + } + const payload = parsePayload(payloadText); + if (isFailedDocument(payload)) { + if (isCodexCapacityPayload(payload)) { + framer.dispose(); + void reader.cancel("Codex capacity retry").catch(() => undefined); + const capacityBytes = Uint8Array.from(joinSseFrameBytes(completeFrameParts)); + return { + kind: "capacity", + response: responseWithBody(response, capacityBytes), + }; + } + framer.dispose(); + return passWithBufferedBody(response, buffered, reader); + } + if (!isSafePreOutputLifecycle(payload)) { + framer.dispose(); + return passWithBufferedBody(response, buffered, reader); + } + } + + if (bufferedBytes > MAX_CLIENT_SSE_FRAME_BYTES) { + framer.dispose(); + return passWithBufferedBody(response, buffered, reader); + } + } + } catch (error) { + framer.dispose(); + void reader.cancel(error).catch(() => undefined); + throw error; + } +} + +/** + * Inspect a Codex response only while no substantive Responses event has been + * exposed. SSE lifecycle frames are held under the existing client-frame byte + * bound; any unknown or output-bearing event commits the attempt permanently. + */ +export async function inspectCodexCapacityBeforeOutput( + response: Response, + options: { streamRequested: boolean; signal?: AbortSignal }, +): Promise { + const contentType = response.headers.get("content-type")?.toLocaleLowerCase("en-US") ?? ""; + const isEventStream = contentType.includes("text/event-stream") + || (response.ok && !!response.body && !contentType && options.streamRequested); + if (isEventStream) return inspectSseCapacityBeforeOutput(response, options.signal); + if (!response.ok + && response.status !== 402 + && response.status !== 429 + && response.status < 500) { + return { kind: "pass", response }; + } + + try { + const body = await readBoundedResponseBody(response.clone(), { + signal: options.signal, + fatalUtf8: true, + }); + if (!body.displaySafe || body.truncated || !body.text.trim()) { + return { kind: "pass", response }; + } + const payload = JSON.parse(body.text) as unknown; + const rebuilt = responseWithBody(response, body.text); + const capacity = isCodexCapacityPayload(payload) + && (!response.ok || isFailedDocument(payload)); + return capacity ? { kind: "capacity", response: rebuilt } : { kind: "pass", response: rebuilt }; + } catch { + if (options.signal?.aborted) throw options.signal.reason; + return { kind: "pass", response }; + } +} diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index b289f7f99..37d6b8827 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -2,6 +2,7 @@ import type { Server } from "bun"; import { bridgeToResponsesSSE, buildResponseJSON, formatErrorResponse, type ResponsesTerminalStatus } from "../../bridge"; import { formatPassthroughUpstreamError } from "./passthrough-error"; import { describeUpstreamConnectFailure } from "./upstream-error"; +import { inspectCodexCapacityBeforeOutput } from "./codex-capacity"; import { getConfigPath, multiAgentGuidanceEnabled, @@ -81,6 +82,7 @@ import { type CodexAuthContext, } from "../../codex/auth-context"; import { + commitCodexAcceptedAccountSelection, computeQuotaCooldown, formatCodexProviderForLog, previewCodexAccountForRequest, @@ -344,6 +346,9 @@ interface CodexPoolAccountRetryArgs { connectMs: number; passthroughEstimate?: number; stream: boolean; + excludedAccountIds?: Set; + recordRejectedOutcome?: boolean; + preserveRejectedResponse?: boolean; } type CodexPoolAccountRetryResult = @@ -402,49 +407,64 @@ async function retryCodexPoolOnAlternateAccount( if (firstAuthCtx.fixedAccount) return { kind: "no-alternate" }; const inboundWire = options.inboundWire ?? "responses"; let retryAuthCtx: CodexAuthContext | undefined; - try { - retryAuthCtx = await resolveCodexAuthContext( - req.headers, - config, - "pool", - { - excludeAccountId: firstAuthCtx.accountId, - modelId: route.modelId, - beginCodexAccountSelection: codexAccountSelectionForTurn(options.turnAdmissionLease), - }, - ); - } catch (error) { - if ( - !(error instanceof CodexPoolAuthenticationError) - && !(error instanceof CodexAuthContextError) - && !(error instanceof CodexAccountCooldownError) - && !(error instanceof CodexMainProfileDrainingError) - ) throw error; + while (!retryAuthCtx) { + try { + retryAuthCtx = await resolveCodexAuthContext( + req.headers, + config, + "pool", + { + ...(args.excludedAccountIds + ? { excludeAccountIds: args.excludedAccountIds } + : { excludeAccountId: firstAuthCtx.accountId }), + modelId: route.modelId, + beginCodexAccountSelection: codexAccountSelectionForTurn(options.turnAdmissionLease), + }, + ); + } catch (error) { + if ( + args.excludedAccountIds + && (error instanceof CodexAuthContextError || error instanceof CodexAccountCooldownError) + && !args.excludedAccountIds.has(error.accountId) + ) { + args.excludedAccountIds.add(error.accountId); + continue; + } + if ( + !(error instanceof CodexPoolAuthenticationError) + && !(error instanceof CodexAuthContextError) + && !(error instanceof CodexAccountCooldownError) + && !(error instanceof CodexMainProfileDrainingError) + ) throw error; + break; + } } if (retryAuthCtx?.kind !== "pool" && retryAuthCtx?.kind !== "main-pool") { return { kind: "no-alternate" }; } - const quotaMeta = codexQuotaOutcomeMeta(firstResponse); - if (outcomeStatus === 429 || outcomeStatus === 402) { - const { applyAccountQuotaFromUpstreamHeaders } = await import("../../codex/auth-api"); - applyAccountQuotaFromUpstreamHeaders( - firstAuthCtx.accountId, - firstResponse.headers, - firstAuthCtx.writerGeneration, - ); - } - if (!shouldDeferCodexResetDerivedCooldown(firstResponse, options.deferCodexResetDerivedCooldown)) { - recordCodexUpstreamOutcome(config, firstAuthCtx.accountId, outcomeStatus, { - ...quotaMeta, - threadId: req.headers.get("x-codex-parent-thread-id"), - modelId: route.modelId, - probeLeaseId: codexProbeLeaseId(firstAuthCtx), - probeQuotaScope: codexProbeQuotaScope(firstAuthCtx), - writerGeneration: firstAuthCtx.writerGeneration, - // Retry already advanced the RR ring via excludeAccountId — reuse for promotion. - ...(retryAuthCtx.accountId ? { promoteAccountId: retryAuthCtx.accountId } : {}), - }); + if (args.recordRejectedOutcome !== false) { + const quotaMeta = codexQuotaOutcomeMeta(firstResponse); + if (outcomeStatus === 429 || outcomeStatus === 402) { + const { applyAccountQuotaFromUpstreamHeaders } = await import("../../codex/auth-api"); + applyAccountQuotaFromUpstreamHeaders( + firstAuthCtx.accountId, + firstResponse.headers, + firstAuthCtx.writerGeneration, + ); + } + if (!shouldDeferCodexResetDerivedCooldown(firstResponse, options.deferCodexResetDerivedCooldown)) { + recordCodexUpstreamOutcome(config, firstAuthCtx.accountId, outcomeStatus, { + ...quotaMeta, + threadId: req.headers.get("x-codex-parent-thread-id"), + modelId: route.modelId, + probeLeaseId: codexProbeLeaseId(firstAuthCtx), + probeQuotaScope: codexProbeQuotaScope(firstAuthCtx), + writerGeneration: firstAuthCtx.writerGeneration, + // Retry already advanced the RR ring via exclusion — reuse for promotion. + ...(retryAuthCtx.accountId ? { promoteAccountId: retryAuthCtx.accountId } : {}), + }); + } } const retryHeaders = headersForCodexAuthContext(req.headers, retryAuthCtx); @@ -463,7 +483,9 @@ async function retryCodexPoolOnAlternateAccount( }); recordAdapterReasoning(logCtx, request); - await firstResponse.body?.cancel().catch(() => undefined); + if (args.preserveRejectedResponse !== true) { + await firstResponse.body?.cancel().catch(() => undefined); + } options.onCodexAuthContextResolved?.(retryAuthCtx); route.provider = retryProvider; logCtx.provider = formatCodexProviderForLog( @@ -833,6 +855,8 @@ async function resolveResponsesCodexAuth( authCtx = await resolveCodexAuthContext(req.headers, config, route.codexAccountMode, { accountId: route.codexAccountId, modelId: route.modelId, + deferSelectionCommit: route.codexAccountMode === "pool" + && route.codexAccountId === undefined, beginCodexAccountSelection: codexAccountSelectionForTurn(options.turnAdmissionLease), }); options.onCodexAuthContextResolved?.(authCtx); @@ -1961,6 +1985,21 @@ async function handleResponsesInner( : describeUpstreamConnectFailure(err, connectMs); return formatErrorResponse(502, "upstream_error", msg); }; + const passthroughTransientRetryOptions = { + abortSignal: upstream.signal, + label: safeHostLabel(request.url), + ...(usesCodexForwardPoolAuth(authCtx, route.provider) && !authCtx.fixedAccount + ? { + prepareTransientRetry: async (response: Response) => { + const inspected = await inspectCodexCapacityBeforeOutput(response, { + streamRequested: parsed.stream, + signal: options.abortSignal, + }); + return { response: inspected.response, retry: inspected.kind !== "capacity" }; + }, + } + : {}), + }; try { // Transient-5xx pre-stream retry (devlog/_plan/260716_claudecode_hardening/010): // the ChatGPT backend emits transient 502/520s that an immediate retry absorbs. @@ -1981,7 +2020,7 @@ async function handleResponsesInner( return res; }); }, - { abortSignal: upstream.signal, label: safeHostLabel(request.url) }, + passthroughTransientRetryOptions, ); } catch (err) { return transportFailureResponse(err); @@ -2040,14 +2079,104 @@ async function handleResponsesInner( return res; }); }, - { abortSignal: upstream.signal, label: safeHostLabel(request.url) }, + passthroughTransientRetryOptions, ); } catch (err) { return transportFailureResponse(err); } } + let capacityRetryExhausted = false; if (usesCodexForwardPoolAuth(authCtx, route.provider) && !authCtx.fixedAccount) { + const excludedAccountIds = new Set(); + let firstCapacityAttempt: { + response: Response; + authCtx: Extract; + request: typeof request; + provider: OcxProviderConfig; + selectedForwardHeaders: Headers; + logProvider: string; + subagentFallbackAccountId: string | null; + } | undefined; + + while (true) { + const inspected = await inspectCodexCapacityBeforeOutput(upstreamResponse, { + streamRequested: parsed.stream, + signal: options.abortSignal, + }); + upstreamResponse = inspected.response; + if (inspected.kind !== "capacity") { + commitCodexAcceptedAccountSelection( + config, + req.headers.get("x-codex-parent-thread-id"), + authCtx.accountId, + route.modelId, + excludedAccountIds, + ); + if (firstCapacityAttempt) { + await firstCapacityAttempt.response.body?.cancel().catch(() => undefined); + } + break; + } + + excludedAccountIds.add(authCtx.accountId); + firstCapacityAttempt ??= { + response: upstreamResponse, + authCtx, + request, + provider: route.provider, + selectedForwardHeaders, + logProvider: logCtx.provider, + subagentFallbackAccountId, + }; + const retry = await retryCodexPoolOnAlternateAccount({ + req, + config, + route, + parsed, + logCtx, + options, + firstAuthCtx: authCtx, + firstResponse: upstreamResponse, + outcomeStatus: 503, + upstream, + connectMs, + passthroughEstimate, + stream: parsed.stream, + excludedAccountIds, + recordRejectedOutcome: false, + preserveRejectedResponse: upstreamResponse === firstCapacityAttempt.response, + }); + if (retry.kind === "transport") { + await firstCapacityAttempt.response.body?.cancel().catch(() => undefined); + authCtx = retry.authCtx; + return transportFailureResponse(retry.error); + } + if (retry.kind === "no-alternate") { + if (upstreamResponse !== firstCapacityAttempt.response) { + await upstreamResponse.body?.cancel().catch(() => undefined); + } + authCtx = firstCapacityAttempt.authCtx; + request = firstCapacityAttempt.request; + route.provider = firstCapacityAttempt.provider; + selectedForwardHeaders = firstCapacityAttempt.selectedForwardHeaders; + logCtx.provider = firstCapacityAttempt.logProvider; + subagentFallbackAccountId = firstCapacityAttempt.subagentFallbackAccountId; + upstreamResponse = firstCapacityAttempt.response; + options.onCodexAuthContextResolved?.(authCtx); + capacityRetryExhausted = true; + break; + } + + authCtx = retry.authCtx; + request = retry.request; + upstreamResponse = retry.upstreamResponse; + selectedForwardHeaders = retry.selectedForwardHeaders; + subagentFallbackAccountId = retry.authCtx.accountId; + } + } + + if (!capacityRetryExhausted && usesCodexForwardPoolAuth(authCtx, route.provider) && !authCtx.fixedAccount) { let poolRetryOutcome: number | undefined; if (await shouldRetryCodexPoolAccountModel400( upstreamResponse, @@ -2102,17 +2231,19 @@ async function handleResponsesInner( const passthroughCt = headers.get("content-type")?.toLowerCase(); const isEventStream = passthroughCt?.includes("text/event-stream") || (upstreamResponse.ok && !!upstreamResponse.body && !passthroughCt && parsed.stream); - const terminalRecorder = codexForwardTerminalOutcomeRecorder( - config, - authCtx, - route.provider, - route.modelId, - logCtx, - req.headers.get("x-codex-parent-thread-id"), - ); + const terminalRecorder = capacityRetryExhausted + ? undefined + : codexForwardTerminalOutcomeRecorder( + config, + authCtx, + route.provider, + route.modelId, + logCtx, + req.headers.get("x-codex-parent-thread-id"), + ); const terminalBodyWillRecord = !!terminalRecorder && upstreamResponse.ok && isEventStream; // Capture quota from upstream response for multi-account tracking - if (usesCodexForwardPoolAuth(authCtx, route.provider)) { + if (!capacityRetryExhausted && usesCodexForwardPoolAuth(authCtx, route.provider)) { // primary was the 5h window; it now carries weekly data for GPT plans. // Prefer primary when present, fall back to secondary for compatibility. const quotaMeta = codexQuotaOutcomeMeta(upstreamResponse); diff --git a/structure/08_openai-provider-tiers.md b/structure/08_openai-provider-tiers.md index 5a42416ba..9b9e9531b 100644 --- a/structure/08_openai-provider-tiers.md +++ b/structure/08_openai-provider-tiers.md @@ -23,6 +23,14 @@ An explicit `Retry-After` or an unclassified quota 429 is account-wide. A reset- the shared native group (including GPT-5.6 Terra/Luna). This allows a same-account combo to test an independent quota without allowing fallbacks that share the exhausted quota. +Model-capacity rejection is a different recovery class. Before substantive Responses output, exact +structured `server_is_overloaded` / `slow_down` signals and the standard model-capacity sentence +exclude the current account only for that request, then try each remaining eligible Pool account +once. Rejected capacity attempts write no cooldown or health, and thread affinity is committed only +for the account whose non-capacity response is accepted. Exhaustion returns the first rejection; the +next request starts with no capacity exclusions and no new affinity from the rejected attempts. Direct +and exact-account routes never rotate, and a committed output event makes the attempt non-replayable. + `pausedCodexAccountIds` is a persisted Pool eligibility boundary. A paused added account or the stable `__main__` alias remains visible for maintenance and quota reads, but is excluded from new affinity, quota rotation, cooldown probes, transient failover, and manual activation. In-flight diff --git a/tests/codex-capacity-retry.test.ts b/tests/codex-capacity-retry.test.ts new file mode 100644 index 000000000..25ef699e2 --- /dev/null +++ b/tests/codex-capacity-retry.test.ts @@ -0,0 +1,73 @@ +import { describe, expect, test } from "bun:test"; +import { + CODEX_MODEL_CAPACITY_MESSAGE, + inspectCodexCapacityBeforeOutput, + isCodexCapacityPayload, +} from "../src/server/responses/codex-capacity"; + +function sseResponse(chunks: readonly string[]): Response { + const encoder = new TextEncoder(); + return new Response(new ReadableStream({ + start(controller) { + for (const chunk of chunks) controller.enqueue(encoder.encode(chunk)); + controller.close(); + }, + }), { headers: { "content-type": "text/event-stream" } }); +} + +describe("Codex model-capacity classification", () => { + test("accepts structured code/type and only the exact standard message", () => { + expect(isCodexCapacityPayload({ error: { code: "server_is_overloaded" } })).toBe(true); + expect(isCodexCapacityPayload({ error: { type: "slow_down" } })).toBe(true); + expect(isCodexCapacityPayload({ error: { message: CODEX_MODEL_CAPACITY_MESSAGE } })).toBe(true); + expect(isCodexCapacityPayload({ + error: { message: `${CODEX_MODEL_CAPACITY_MESSAGE} retry later` }, + })).toBe(false); + expect(isCodexCapacityPayload({ error: { message: "model unavailable" } })).toBe(false); + }); + + test("recognizes a fragmented SSE error after lifecycle-only frames", async () => { + const source = [ + 'event: response.created\ndata: {"type":"response.created"}\n\n', + 'event: error\ndata: {"type":"error","error":{"message":"Selected model is at ', + 'capacity. Please try a different model."}}\n\n', + ]; + const inspected = await inspectCodexCapacityBeforeOutput(sseResponse(source), { + streamRequested: true, + }); + expect(inspected.kind).toBe("capacity"); + expect(await inspected.response.text()).toBe(source.join("")); + }); + + test("any substantive event commits the stream and forbids capacity replay", async () => { + const source = [ + 'event: response.output_item.added\ndata: {"type":"response.output_item.added","item":{"type":"function_call"}}\n\n', + `event: response.failed\ndata: ${JSON.stringify({ + type: "response.failed", + response: { + status: "failed", + error: { code: "server_is_overloaded", message: "busy" }, + }, + })}\n\n`, + ].join(""); + const inspected = await inspectCodexCapacityBeforeOutput(sseResponse([source]), { + streamRequested: true, + }); + expect(inspected.kind).toBe("pass"); + expect(await inspected.response.text()).toBe(source); + }); + + test("non-capacity transient response stays byte-for-byte readable", async () => { + const body = JSON.stringify({ error: { code: "upstream_server_error", message: "gateway reset" } }); + const inspected = await inspectCodexCapacityBeforeOutput(new Response(body, { + status: 502, + statusText: "Bad Gateway", + headers: { "content-type": "application/json", "x-origin": "kept" }, + }), { streamRequested: false }); + expect(inspected.kind).toBe("pass"); + expect(inspected.response.status).toBe(502); + expect(inspected.response.statusText).toBe("Bad Gateway"); + expect(inspected.response.headers.get("x-origin")).toBe("kept"); + expect(await inspected.response.text()).toBe(body); + }); +}); diff --git a/tests/codex-pool-rotation.test.ts b/tests/codex-pool-rotation.test.ts index e50b10644..6d523ab68 100644 --- a/tests/codex-pool-rotation.test.ts +++ b/tests/codex-pool-rotation.test.ts @@ -1,4 +1,5 @@ import { + commitRoundRobinAccountSuccess, clearPoolRotationState, DEFAULT_ACCOUNT_PRIORITY, normalizeAccountPriority, @@ -308,6 +309,13 @@ describe("pickRoundRobinAccount", () => { expect(peekAfter).not.toBe(picked); expect(pickRoundRobinAccount("codex", ids, 1)).toBe(peekAfter); }); + + test("deferred commit advances from the accepted account, not a rejected preview", () => { + const ids = ["a", "b", "c"]; + expect(peekRoundRobinAccount("codex", ids, 1)).toBe("a"); + expect(commitRoundRobinAccountSuccess("codex", ["b", "c"], "b", 1)).toBe(true); + expect(peekRoundRobinAccount("codex", ids, 1)).toBe("c"); + }); }); describe("accountPoolStrategy new-session routing", () => { diff --git a/tests/server-auth.test.ts b/tests/server-auth.test.ts index f67a1689e..6a0f8dcdd 100644 --- a/tests/server-auth.test.ts +++ b/tests/server-auth.test.ts @@ -13,10 +13,12 @@ import { CODEX_THREAD_AFFINITY_IDLE_TTL_MS, clearCodexUpstreamHealth, clearThreadAccountMap, + getEffectiveActiveCodexAccountId, getCodexUpstreamHealth, isCodexAccountSoftAvoided, recordCodexUpstreamOutcome, } from "../src/codex/routing"; +import { clearPoolRotationState, POOL_KEY_CODEX } from "../src/codex/pool-rotation"; import { loadConfig, saveConfig } from "../src/config"; import { clearUpstreamHostHealth, getUpstreamHostHealth, recordUpstreamHostFailure, upstreamHostHealthKey } from "../src/codex/upstream-host-health"; import { deriveProviderPresets } from "../src/providers/derive"; @@ -40,6 +42,7 @@ import { fakeChatGptJwt } from "./helpers/fake-chatgpt-jwt"; import { installIsolatedCodexHome, type IsolatedCodexHome } from "./helpers/isolated-codex-home"; import { configuredAdminToken } from "../src/lib/admin-secrets"; import { SYSTEM_RESTART_CAPABILITY_VERSION } from "../src/lib/system-restart-contract"; +import { CODEX_MODEL_CAPACITY_MESSAGE } from "../src/server/responses/codex-capacity"; const previousApiToken = process.env.OPENCODEX_API_AUTH_TOKEN; const previousOpencodexHome = process.env.OPENCODEX_HOME; @@ -135,6 +138,7 @@ afterEach(() => { clearThreadAccountMap(); clearAccountNeedsReauth("pool-a"); clearAccountNeedsReauth("pool-b"); + clearAccountNeedsReauth("pool-c"); clearAccountQuota(); if (existsSync(TEST_DIR)) rmSync(TEST_DIR, { recursive: true }); }); @@ -147,6 +151,50 @@ function unsupportedModelBody(model = POOL_RETRY_MODEL): string { }); } +function capacityErrorBody( + code: "server_is_overloaded" | "slow_down" | undefined = undefined, + message = CODEX_MODEL_CAPACITY_MESSAGE, +): string { + return JSON.stringify({ + error: { + type: "server_error", + ...(code ? { code } : {}), + message, + }, + }); +} + +function capacityFailedSse( + code: "server_is_overloaded" | "slow_down" | undefined = undefined, + message = CODEX_MODEL_CAPACITY_MESSAGE, +): string { + const error = { + type: "server_error", + ...(code ? { code } : {}), + message, + }; + return [ + 'event: response.created\ndata: {"type":"response.created","response":{"id":"capacity-attempt","status":"in_progress"}}\n\n', + `event: response.failed\ndata: ${JSON.stringify({ + type: "response.failed", + response: { status: "failed", error, last_error: error }, + })}\n\n`, + ].join(""); +} + +function capacityErrorSse( + code: "server_is_overloaded" | "slow_down" | undefined = undefined, + message = CODEX_MODEL_CAPACITY_MESSAGE, +): string { + return [ + 'event: response.created\ndata: {"type":"response.created","response":{"id":"capacity-attempt","status":"in_progress"}}\n\n', + `event: error\ndata: ${JSON.stringify({ + type: "error", + error: { ...(code ? { code } : {}), message }, + })}\n\n`, + ].join(""); +} + type PoolRetryHarness = { config: OcxConfig; dispatches: string[]; @@ -156,6 +204,7 @@ type PoolRetryHarness = { model?: string; path?: "/v1/responses" | "/v1/responses/compact"; callerBearer?: boolean; + threadId?: string; }) => Promise; restoreFetch: () => void; server: ReturnType; @@ -183,6 +232,7 @@ async function startPoolRetryHarness( reply: (accountId: string, request: Request) => Response | Promise, options: { secondAccount?: boolean; + thirdAccount?: boolean; streamMode?: "legacy-tee" | "eager-relay"; accountMode?: "direct" | "pool"; activeAccountId?: string; @@ -194,6 +244,8 @@ async function startPoolRetryHarness( pausedAccountIds?: string[]; reauthAccountIds?: string[]; omitCredentialAccountIds?: string[]; + accountPoolStrategy?: "quota" | "round-robin" | "fill-first"; + accountPoolStickyLimit?: number; } = {}, ): Promise { await removeTestDirBestEffort(TEST_DIR); @@ -201,10 +253,12 @@ async function startPoolRetryHarness( process.env.OPENCODEX_HOME = TEST_DIR; clearCodexUpstreamHealth(); clearThreadAccountMap(); + clearPoolRotationState(POOL_KEY_CODEX); clearAccountQuota(); clearRequestLogsForTests(); clearAccountNeedsReauth("pool-a"); clearAccountNeedsReauth("pool-b"); + clearAccountNeedsReauth("pool-c"); // The registry is process-global and survives a harness teardown. WS-REBIND-01 // asserts exact per-account socket counts, so a socket leaked by any earlier test // in this file shifts its snapshots and fails it in milliseconds — which reads as @@ -225,6 +279,7 @@ async function startPoolRetryHarness( const redirectedFetch = globalThis.fetch; const secondAccount = options.secondAccount ?? true; + const thirdAccount = options.thirdAccount ?? false; const config = { port: 0, defaultProvider: "openai", @@ -243,8 +298,13 @@ async function startPoolRetryHarness( ...(secondAccount ? [{ id: "pool-b", email: "pool-b@example.test", isMain: false, chatgptAccountId: "acct-pool-b" }] : []), + ...(thirdAccount + ? [{ id: "pool-c", email: "pool-c@example.test", isMain: false, chatgptAccountId: "acct-pool-c" }] + : []), ], activeCodexAccountId: options.activeAccountId ?? "pool-a", + ...(options.accountPoolStrategy ? { accountPoolStrategy: options.accountPoolStrategy } : {}), + ...(options.accountPoolStickyLimit ? { accountPoolStickyLimit: options.accountPoolStickyLimit } : {}), ...(options.accountNamespaces ? { codexAccountNamespaces: options.accountNamespaces } : {}), ...(options.pausedAccountIds ? { pausedCodexAccountIds: options.pausedAccountIds } : {}), ...(options.visionSidecarModel ? { visionSidecar: { model: options.visionSidecarModel } } : {}), @@ -272,6 +332,17 @@ async function startPoolRetryHarness( } updateAccountQuota("pool-b", 20); } + if (thirdAccount) { + if (!options.omitCredentialAccountIds?.includes("pool-c")) { + saveCodexAccountCredential("pool-c", { + accessToken: "pool-c-token", + refreshToken: "pool-c-refresh", + expiresAt: Date.now() + 10 * 60_000, + chatgptAccountId: "acct-pool-c", + }); + } + updateAccountQuota("pool-c", 30); + } for (const accountId of options.reauthAccountIds ?? []) markAccountNeedsReauth(accountId); const server = startServer(0); @@ -289,11 +360,13 @@ async function startPoolRetryHarness( model = POOL_RETRY_MODEL, path = "/v1/responses", callerBearer = true, + threadId, } = {}) => originalGlobalFetch(new URL(path, server.url), { method: "POST", headers: { "content-type": "application/json", ...(callerBearer ? { authorization: "Bearer inbound-token" } : {}), + ...(threadId ? { "x-codex-parent-thread-id": threadId } : {}), }, body: JSON.stringify({ model, input: path.endsWith("/compact") ? [] : "hello", stream }), signal, @@ -2258,6 +2331,220 @@ describe("server local API auth", () => { } }, { timeout: SERVER_BUDGET_MS }); + test("capacity rejection rotates through every eligible pool account once", async () => { + const harness = await startPoolRetryHarness(accountId => { + if (accountId === "acct-pool-a") { + return new Response(capacityErrorBody(undefined), { + status: 502, + headers: { "content-type": "application/json", "x-capacity-attempt": "a" }, + }); + } + if (accountId === "acct-pool-b") { + return new Response(capacityErrorBody("server_is_overloaded", "busy"), { + status: 503, + headers: { "content-type": "application/json", "x-capacity-attempt": "b" }, + }); + } + return Response.json({ id: "capacity-rotation-success", status: "completed", output: [] }); + }, { thirdAccount: true }); + try { + const response = await harness.request(); + expect(response.status).toBe(200); + const responseText = await response.text(); + expect(responseText).toContain("capacity-rotation-success"); + expect(harness.dispatches).toEqual(["acct-pool-a", "acct-pool-b", "acct-pool-c"]); + expect(getCodexUpstreamHealth("pool-a")).toBeNull(); + expect(getCodexUpstreamHealth("pool-b")).toBeNull(); + } finally { + await stopPoolRetryHarness(harness); + } + }); + + test("capacity rotation skips an unavailable intermediate account", async () => { + const harness = await startPoolRetryHarness(accountId => accountId === "acct-pool-a" + ? new Response(capacityErrorBody("server_is_overloaded"), { + status: 503, + headers: { "content-type": "application/json" }, + }) + : Response.json({ id: "available-c", status: "completed", output: [] }), { + thirdAccount: true, + omitCredentialAccountIds: ["pool-b"], + }); + try { + const response = await harness.request(); + expect(response.status).toBe(200); + expect(await response.text()).toContain("available-c"); + expect(harness.dispatches).toEqual(["acct-pool-a", "acct-pool-c"]); + expect(getCodexUpstreamHealth("pool-a")).toBeNull(); + expect(getCodexUpstreamHealth("pool-b")).toBeNull(); + } finally { + await stopPoolRetryHarness(harness); + } + }); + + test("capacity rotation binds thread affinity only to the accepted account", async () => { + const harness = await startPoolRetryHarness(accountId => accountId === "acct-pool-a" + ? new Response(capacityErrorBody("server_is_overloaded"), { + status: 503, + headers: { "content-type": "application/json" }, + }) + : Response.json({ id: "accepted-b", status: "completed", output: [] })); + try { + const first = await harness.request({ threadId: "capacity-thread" }); + expect(first.status).toBe(200); + expect(await first.text()).toContain("accepted-b"); + expect(harness.dispatches).toEqual(["acct-pool-a", "acct-pool-b"]); + + const second = await harness.request({ threadId: "capacity-thread" }); + expect(second.status).toBe(200); + expect(await second.text()).toContain("accepted-b"); + expect(harness.dispatches).toEqual(["acct-pool-a", "acct-pool-b", "acct-pool-b"]); + expect(getCodexUpstreamHealth("pool-a")).toBeNull(); + } finally { + await stopPoolRetryHarness(harness); + } + }); + + test("capacity rotation commits round-robin state only for the accepted account", async () => { + const harness = await startPoolRetryHarness(accountId => accountId === "acct-pool-a" + ? new Response(capacityErrorBody("server_is_overloaded"), { + status: 503, + headers: { "content-type": "application/json" }, + }) + : Response.json({ id: `accepted-${accountId}`, status: "completed", output: [] }), { + thirdAccount: true, + accountPoolStrategy: "round-robin", + accountPoolStickyLimit: 1, + }); + try { + const first = await harness.request({ threadId: "rr-capacity-thread" }); + expect(first.status).toBe(200); + expect(harness.dispatches).toEqual(["acct-pool-a", "acct-pool-b"]); + expect(getEffectiveActiveCodexAccountId(harness.config)).toBe("pool-b"); + + const sameThread = await harness.request({ threadId: "rr-capacity-thread" }); + expect(sameThread.status).toBe(200); + expect(harness.dispatches).toEqual(["acct-pool-a", "acct-pool-b", "acct-pool-b"]); + + const nextThread = await harness.request({ threadId: "rr-next-thread" }); + expect(nextThread.status).toBe(200); + expect(harness.dispatches).toEqual([ + "acct-pool-a", + "acct-pool-b", + "acct-pool-b", + "acct-pool-c", + ]); + expect(getEffectiveActiveCodexAccountId(harness.config)).toBe("pool-c"); + } finally { + await stopPoolRetryHarness(harness); + } + }); + + test("capacity exhaustion preserves the first error and resets on the next request", async () => { + let requestOrdinal = 0; + const harness = await startPoolRetryHarness(accountId => { + if (accountId === "acct-pool-a") requestOrdinal += 1; + if (requestOrdinal === 2 && accountId === "acct-pool-a") { + return Response.json({ id: "fresh-request-success", status: "completed", output: [] }); + } + const attempt = accountId === "acct-pool-a" ? "a" : "b"; + return new Response(capacityErrorBody("slow_down", `capacity-${attempt}`), { + status: 502, + statusText: `Capacity ${attempt.toUpperCase()}`, + headers: { "content-type": "application/json", "x-capacity-attempt": attempt }, + }); + }); + try { + const exhausted = await harness.request(); + expect(exhausted.status).toBe(502); + expect(exhausted.headers.get("x-capacity-attempt")).toBe("a"); + expect(await exhausted.text()).toContain("capacity-a"); + expect(harness.dispatches).toEqual(["acct-pool-a", "acct-pool-b"]); + expect(getCodexUpstreamHealth("pool-a")).toBeNull(); + expect(getCodexUpstreamHealth("pool-b")).toBeNull(); + + const fresh = await harness.request(); + expect(fresh.status).toBe(200); + const freshText = await fresh.text(); + expect(freshText).toContain("fresh-request-success"); + expect(harness.dispatches).toEqual(["acct-pool-a", "acct-pool-b", "acct-pool-a"]); + } finally { + await stopPoolRetryHarness(harness); + } + }); + + test("SSE capacity failure rotates only before substantive output", async () => { + const preOutput = await startPoolRetryHarness(accountId => accountId === "acct-pool-a" + ? new Response(capacityErrorSse(undefined), { + headers: { "content-type": "text/event-stream" }, + }) + : new Response( + 'event: response.completed\ndata: {"type":"response.completed","response":{"id":"from-b","status":"completed","output":[]}}\n\n', + { headers: { "content-type": "text/event-stream" } }, + )); + try { + const text = await (await preOutput.request({ stream: true })).text(); + expect(preOutput.dispatches).toEqual(["acct-pool-a", "acct-pool-b"]); + expect(text).toContain('"id":"from-b"'); + expect(text).not.toContain("capacity-attempt"); + } finally { + await stopPoolRetryHarness(preOutput); + } + + const postOutput = await startPoolRetryHarness(() => new Response([ + 'event: response.created\ndata: {"type":"response.created","response":{"id":"committed","status":"in_progress"}}\n\n', + 'event: response.output_text.delta\ndata: {"type":"response.output_text.delta","delta":"visible"}\n\n', + capacityFailedSse("server_is_overloaded", "busy"), + ].join(""), { headers: { "content-type": "text/event-stream" } })); + try { + const text = await (await postOutput.request({ stream: true })).text(); + expect(postOutput.dispatches).toEqual(["acct-pool-a"]); + expect(text).toContain('"delta":"visible"'); + expect(text).toContain("response.failed"); + } finally { + await stopPoolRetryHarness(postOutput); + } + }); + + test("Direct mode never rotates a capacity rejection", async () => { + const harness = await startPoolRetryHarness(() => new Response(capacityErrorBody("server_is_overloaded"), { + status: 409, + headers: { "content-type": "application/json" }, + }), { accountMode: "direct" }); + try { + const response = await harness.request(); + expect(response.status).toBe(409); + expect(harness.dispatches).toEqual(["missing"]); + } finally { + await stopPoolRetryHarness(harness); + } + }); + + test("exact account selector never rotates a capacity rejection", async () => { + const body = capacityErrorBody("server_is_overloaded"); + const harness = await startPoolRetryHarness(() => new Response(body, { + status: 409, + headers: { "content-type": "application/json", "x-exact-response": "original" }, + }), { + accountMode: "direct", + activeAccountId: "pool-b", + accountNamespaces: { side: "pool-a" }, + }); + try { + const response = await harness.request({ + model: `side/${POOL_RETRY_MODEL}`, + callerBearer: false, + }); + expect(response.status).toBe(409); + expect(response.headers.get("x-exact-response")).toBe("original"); + expect(await response.text()).toBe(body); + expect(harness.dispatches).toEqual(["acct-pool-a"]); + expect(loadConfig().activeCodexAccountId).toBe("pool-b"); + } finally { + await stopPoolRetryHarness(harness); + } + }); + test("#584: pre-stream 429 retries once on another eligible pool account", async () => { const harness = await startPoolRetryHarness(accountId => accountId === "acct-pool-a" ? new Response(JSON.stringify({ error: { message: "rate limited" } }), { diff --git a/tests/upstream-transient-retry.test.ts b/tests/upstream-transient-retry.test.ts index 0f6303813..c15ad9e8c 100644 --- a/tests/upstream-transient-retry.test.ts +++ b/tests/upstream-transient-retry.test.ts @@ -44,6 +44,25 @@ describe("fetchWithTransientRetry", () => { expect(res.status).toBe(400); }); + test("semantic gate can preserve a rebuilt transient response without retrying", async () => { + let calls = 0; + const rebuilt = new Response("capacity", { status: 502, headers: { "x-prepared": "yes" } }); + const res = await fetchWithTransientRetry(async () => { + calls++; + return bodyResponse(502); + }, { + slowAttemptMs: 60_000, + prepareTransientRetry: async response => { + expect(response.status).toBe(502); + return { response: rebuilt, retry: false }; + }, + }); + expect(calls).toBe(1); + expect(res).toBe(rebuilt); + expect(res.headers.get("x-prepared")).toBe("yes"); + expect(await res.text()).toBe("capacity"); + }); + test("honors Retry-After header for the backoff delay", async () => { let calls = 0; const started = Date.now();