diff --git a/src/codex/routing.ts b/src/codex/routing.ts index 04c5b8b1ab..e26c902b4d 100644 --- a/src/codex/routing.ts +++ b/src/codex/routing.ts @@ -63,6 +63,25 @@ type CodexUpstreamHealth = { lastFailureAt?: number; /** Hard cooldown (quota 429). Survives a later 2xx; blocks auth + selection. */ cooldownUntil?: number; + /** + * How long a quota refusal keeps selection away from this account (or this native quota + * group), as opposed to how long it is hard-blocked. + * + * The two are deliberately different lengths. {@link CODEX_MAX_RESET_DERIVED_COOLDOWN_MS} + * caps the hard cooldown at 15 minutes because a reset announcement is advisory and plan + * quota usually frees up before it — an account must stay reachable so the pool can find + * that out (#433). The window the refusal announced is not 15 minutes, though, so once the + * cooldown lapses the account is selectable again while its burst window is still spent, + * and the strategy picks it straight back: this proxy reads a weekly bar a burst limit never + * touches, so a refused account still scores as the coolest in the pool. Every request then + * earns the same 429 until the process restarts, which is the only thing that drops this map. + * + * So the announcement governs avoidance and the cap still governs blocking. Avoidance is soft + * in the {@link softAvoidUntil} sense: it reorders the pool and releases a bound thread, and + * the last-resort paths still reach the account when nothing else can serve, so one pessimistic + * announcement cannot stall routing. + */ + quotaAvoidUntil?: number; /** When the current cooldown was recorded; origin of the probe interval clock. */ cooldownSince?: number; /** @@ -112,6 +131,12 @@ const CODEX_MAX_QUOTA_COOLDOWN_MS = 24 * 60 * 60_000; * the Retry-After ceiling (#433). */ const CODEX_MAX_RESET_DERIVED_COOLDOWN_MS = 15 * 60_000; +/** + * Ceiling on quota-refusal avoidance. Generous enough to cover a full five-hour burst window, + * tight enough that a weekly or monthly reset four days out cannot take an account out of + * rotation for the {@link CODEX_MAX_QUOTA_COOLDOWN_MS} day the Retry-After ceiling allows. + */ +const CODEX_MAX_QUOTA_AVOID_MS = 6 * 60 * 60_000; /** Minimum gap between probe leases for one cooled-down account. */ export const CODEX_QUOTA_PROBE_INTERVAL_MS = 5 * 60_000; export const CODEX_FAILURE_WINDOW_MS = 5 * 60_000; @@ -546,6 +571,51 @@ export function computeQuotaCooldown(meta: CodexUpstreamOutcomeMeta = {}): { return { until: now + CODEX_DEFAULT_QUOTA_COOLDOWN_MS, source: "default" }; } +/** + * When the pool should stop preferring an account after it refused on quota. + * + * The earliest window the refusal actually announced, bounded by {@link CODEX_MAX_QUOTA_AVOID_MS}, + * and never shorter than the cooldown the same refusal produced — a Retry-After directive that + * outlasts every announcement still governs. + */ +function quotaAvoidUntilFor(meta: CodexUpstreamOutcomeMeta, now: number, cooldownUntil: number): number { + const values = Array.isArray(meta.resetAt) ? meta.resetAt : [meta.resetAt]; + let announced: number | undefined; + for (const value of values) { + const timestamp = resetTimestampMs(value); + if (timestamp === undefined) continue; + const delay = timestamp - now; + if (delay <= 0) continue; + const until = now + Math.min(delay, CODEX_MAX_QUOTA_AVOID_MS); + if (announced === undefined || until < announced) announced = until; + } + return Math.max(cooldownUntil, announced ?? 0); +} + +/** Live quota-refusal avoidance for an account, including the lane the request belongs to. */ +function codexQuotaAvoidUntil( + accountId: string, + quotaScope: CodexQuotaScope | undefined, + now: number, +): number | null { + const live = (value: number | undefined): number | null => + typeof value === "number" && Number.isFinite(value) && value > now ? value : null; + const account = live(upstreamHealth.get(accountId)?.quotaAvoidUntil); + const scoped = quotaScope === undefined + ? null + : live(scopedHealthFor(accountId, quotaScope)?.quotaAvoidUntil); + if (account === null) return scoped; + return scoped === null ? account : Math.max(account, scoped); +} + +function isCodexQuotaAvoided( + accountId: string, + quotaScope: CodexQuotaScope | undefined, + now: number, +): boolean { + return codexQuotaAvoidUntil(accountId, quotaScope, now) !== null; +} + export function computeQuotaCooldownUntil(meta: CodexUpstreamOutcomeMeta = {}): number { return computeQuotaCooldown(meta).until; } @@ -787,6 +857,10 @@ function settleCooldownRecoveryLease(claim: CooldownRecoveryLease, recovered: bo cooldownSource: _source, probeLeaseId: _leaseId, probeLeaseGeneration: _leaseGeneration, + // "The quota window moved" is a statement about the whole refusal, so the avoidance it + // announced goes with the block it produced. Leaving it would make this escape hatch stop + // escaping: the account would still be passed over by every selection it is meant to win. + quotaAvoidUntil: _avoid, ...rest } = health; upstreamHealth.set(claim.accountId, { @@ -914,8 +988,11 @@ export function resetCodexRoutingForManualSelection(accountId: string): void { const current = upstreamHealth.get(accountId); if (!current) return; const preserved = preservedCooldownFields(current); - if (Object.keys(preserved).length === 0) upstreamHealth.delete(accountId); - else upstreamHealth.set(accountId, { consecutiveFailures: 0, ...preserved }); + // Quota avoidance is a preference, like the soft avoid dropped above, and an operator naming + // this account has overruled it. The hard cooldown is the part that survives. + const { quotaAvoidUntil: _avoid, ...retained } = preserved; + if (Object.keys(retained).length === 0) upstreamHealth.delete(accountId); + else upstreamHealth.set(accountId, { consecutiveFailures: 0, ...retained }); } export function getCodexAccountCooldownUntil(accountId: string, now = Date.now()): number | null { @@ -1089,6 +1166,7 @@ function isCodexAccountSelectable( return !isCodexAccountPaused(config, accountId) && !isCodexAccountPlanExcluded(config, accountId) && getCodexQuotaHealthSnapshot(accountId, quotaScope, now) === null + && !isCodexQuotaAvoided(accountId, quotaScope, now) && !isCodexAccountSoftAvoided(accountId, now) && isCodexAccountUsable(config, accountId, selectionOptions); } @@ -1328,6 +1406,7 @@ function getEligiblePoolAccounts( && (!skipFailoverReadyCandidates || !shouldFailover(config, account.id, now))) .filter(account => getCodexQuotaHealthSnapshot(account.id, quotaScope, now) === null) .filter(account => !isCodexAccountSoftAvoided(account.id, now)) + .filter(account => !isCodexQuotaAvoided(account.id, quotaScope, now)) .filter(account => isCodexAccountUsable(config, account.id, selectionOptions)) .map(account => account.id); // The main Codex account is not stored in config.codexAccounts; include it as a @@ -1980,6 +2059,45 @@ export function resolveCodexAccountForThread( return resolution.status === "selected" ? resolution.accountId : null; } +function carriesQuotaRefusal(health: CodexUpstreamHealth | undefined): boolean { + return health?.lastFailureStatus === 429 || health?.lastFailureStatus === 402; +} + +/** + * Has this account refused a request on quota without serving one since? + * + * Thread affinity is a prompt-cache optimization and every rule around it is a preference: + * `autoSwitchThreshold` is a hint that an account is getting busy, and `pool.cacheAffinity` + * deliberately raises that bar further. A refusal is not a preference, and once the account has + * told THIS thread it cannot serve, the binding has nothing left to optimize. + * + * The distinction matters because the cooldown a 429 writes is deliberately short. A reset + * announcement is advisory — plan quota routinely frees up before the advertised instant — so + * {@link CODEX_MAX_RESET_DERIVED_COOLDOWN_MS} caps it at 15 minutes. The five-hour window that + * announcement describes is not capped, so an account whose burst window is spent looks + * selectable again long before it is. For an unbound request that is correct: going back to find + * out is how the pool learns the window moved. For a BOUND thread it is a loop with no exit — + * the cooldown lapses, the account still scores lowest on the only window this proxy has a + * reading for (its weekly bar, untouched by a burst limit), the thread rebinds, and earns the + * identical 429. Cleared affinity does not help: the next request re-derives the same choice. + * From the Codex side that reads exactly as reported — a new session rotates normally while an + * existing one is locked to an exhausted account until the proxy is restarted, because a restart + * is the only thing that drops the binding and the stale health together. + * + * `lastFailureStatus` is the right evidence because of when it ends: {@link preservedCooldownFields} + * strips it from every recovery write, so it survives exactly until the account actually serves a + * request again. Nothing here blocks that — selection is untouched, so unbound traffic still probes + * the account and the first success releases every thread this refused. + * + * Scope follows where the refusal was recorded. An account-wide throttle lands in + * `upstreamHealth` and releases every lane; a reset-derived refusal lands against one native + * quota group, so a spent Spark window still cannot displace the same thread's Terra binding. + */ +function hasUnrecoveredCodexQuotaRefusal(accountId: string, quotaScope?: CodexQuotaScope): boolean { + if (carriesQuotaRefusal(upstreamHealth.get(accountId))) return true; + return quotaScope !== undefined && carriesQuotaRefusal(scopedHealthFor(accountId, quotaScope)); +} + function previewReusableAffinityAccount( entry: ThreadAffinityEntry | undefined, config: OcxConfig, @@ -1992,6 +2110,7 @@ function previewReusableAffinityAccount( || isThreadAffinityExpired(entry, now) || !isThreadAffinityGenerationLive(entry) || !isCodexAccountSelectable(config, entry.accountId, now, quotaScope, selectionOptions) + || hasUnrecoveredCodexQuotaRefusal(entry.accountId, quotaScope) || shouldFailover(config, entry.accountId, now) ) { return null; @@ -2225,6 +2344,7 @@ export function resolveCodexAccountForThreadDetailed( const detourReusable = !isThreadAffinityExpired(detourEntry, now) && isThreadAffinityGenerationLive(detourEntry) && isCodexAccountSelectable(config, detourEntry.accountId, now, quotaScope, selectionOptions) + && !hasUnrecoveredCodexQuotaRefusal(detourEntry.accountId, quotaScope) && !shouldFailover(config, detourEntry.accountId, now); if (detourReusable) { detourEntry.lastUsedAt = now; @@ -2262,11 +2382,16 @@ export function resolveCodexAccountForThreadDetailed( const selectableForRequest = selectableForSharedState && isCodexAccountSelectable(config, entry.accountId, now, quotaScope, selectionOptions); const failoverReady = shouldFailover(config, entry.accountId, now); + // A quota refusal outranks every affinity preference, including `pool.cacheAffinity`: + // the account has already told this thread it cannot serve it. + const quotaRefused = hasUnrecoveredCodexQuotaRefusal(entry.accountId, quotaScope); const healthyForSharedAffinity = selectableForSharedState && hasCodexQuotaHeadroom(config, entry.accountId, sharedSelectionOptions, now) + && !quotaRefused && !failoverReady; if ( selectableForRequest + && !quotaRefused // Affined threads must leave a failing account once the streak trips failover // (soft-avoid covers the first-hit case; this catches post-avoid residual streaks). && !failoverReady @@ -2496,6 +2621,20 @@ export function recordCodexUpstreamOutcome( setScopedHealth(accountId, meta.probeQuotaScope, withProbeLeaseReleased(scopedProbe, now)); } } + // A served request is what ends the refusal marker the quota branch left on this lane. + // The probe contract above owns the scoped COOLDOWN; this owns only the field + // {@link hasUnrecoveredCodexQuotaRefusal} reads, which would otherwise keep threads away + // from an account that is demonstrably serving them again. The account-wide marker needs + // no equivalent: every recovery write below runs it through preservedCooldownFields. + const refusedScope = quotaScope ? scopedHealthFor(accountId, quotaScope) : undefined; + if (quotaScope && refusedScope && carriesQuotaRefusal(refusedScope)) { + const { + lastFailureStatus: _refusal, lastFailureAt: _refusedAt, quotaAvoidUntil: _avoid, ...retained + } = refusedScope; + // A live cooldown and its probe bookkeeping survive; an entry that held nothing else goes. + if (Object.keys(retained).length > 1) setScopedHealth(accountId, quotaScope, retained); + else deleteScopedHealth(accountId, quotaScope); + } const current = upstreamHealth.get(accountId); const cooldownUntil = getCodexAccountCooldownUntil(accountId, now); // A leased probe that is still on its own cooldown generation proves the @@ -2639,6 +2778,7 @@ export function recordCodexUpstreamOutcome( lastFailureStatus, lastFailureAt: now, cooldownUntil: until, + quotaAvoidUntil: quotaAvoidUntilFor(meta, now, until), cooldownSince: now, cooldownSource: source, cooldownGeneration, @@ -2687,6 +2827,7 @@ export function recordCodexUpstreamOutcome( lastFailureStatus, lastFailureAt: now, cooldownUntil: until, + quotaAvoidUntil: quotaAvoidUntilFor(meta, now, until), cooldownSince: now, cooldownSource: source, cooldownGeneration, diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index c34fda778b..155e107799 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -212,6 +212,7 @@ import { fetchWithResetRetry, fetchWithTransientRetry, isNonReplayableResponse, + isTransientUpstreamStatus, prepareSameTarget429Wait, } from "../../lib/upstream-retry"; import { @@ -1146,6 +1147,28 @@ export async function shouldRetryCodexPoolAccountQuota( } } +/** + * A pre-stream upstream 5xx another Codex account may still be able to serve. + * + * `server_is_overloaded` is the shape this exists for. The ChatGPT backend refuses in a few + * hundred milliseconds, the body carries no quota evidence, and nothing in that exchange is + * account health — so the pool keeps choosing the same account and every request fails on it + * while the other accounts sit idle. That is what an operator sees as the pool refusing to move. + * + * The status stays exactly as upstream sent it. `classifyCodexUpstreamOutcome` maps 5xx to the + * transient class, so the account earns an ordinary failure streak and `upstreamFailoverThreshold` + * decides when it is soft-avoided, rather than a quota cooldown it never earned. + * + * Deliberately narrow. {@link isNonReplayableResponse} still refuses: a post-send WebSocket + * gateway status means the body already reached the origin, so sending it from a second account + * could duplicate a turn the origin may still be running. A 5xx whose body confirms quota is not + * routed here either — {@link shouldRetryCodexPoolAccountQuota} classifies that one first and + * carries the cooldown with it. + */ +export function shouldRetryCodexPoolAccountTransient(response: Response): boolean { + return !isNonReplayableResponse(response) && isTransientUpstreamStatus(response.status); +} + interface CodexPoolAccountRetryArgs { /** Sanitized caller input, before any selected Pool credential was materialized. */ callerAuthHeaders: Headers; @@ -1320,6 +1343,21 @@ async function retryCodexPoolOnAlternateAccount( const inboundWire = options.inboundWire ?? "responses"; const entitlementResolver = options.resolveCodexModelEntitlements ?? resolveCodexModelEntitlements; let retryAuthCtx: CodexAuthContext | undefined; + // A transient 5xx must record even when this request cannot move: the ordinary terminal + // recorder only fires for an OK event-stream body, so a pre-stream refusal would otherwise + // leave the account looking healthy no matter how many times it refused, and the pool would + // keep handing it the next request. + const recordUnmovedTransientOutcome = (): void => { + if (!isTransientUpstreamStatus(outcomeStatus)) return; + recordCodexUpstreamOutcome(config, firstAuthCtx.accountId, outcomeStatus, { + threadId: firstAuthCtx.affinityKey, + fixedAccount: firstAuthCtx.fixedAccount, + modelId: route.modelId, + probeLeaseId: codexProbeLeaseId(firstAuthCtx), + probeQuotaScope: codexProbeQuotaScope(firstAuthCtx), + writerGeneration: firstAuthCtx.writerGeneration, + }); + }; if (outcomeStatus === 400 && ACCOUNT_GATED_NATIVE_OPENAI_MODELS.has(route.modelId)) { invalidateCodexModelEntitlementsForAccount(firstAuthCtx.accountId); let refreshed; @@ -1344,6 +1382,7 @@ async function retryCodexPoolOnAlternateAccount( // Exact account selectors may retry the same confirmed account above, but must never resolve // an alternate. Quota failures and a refreshed entitlement miss remain terminal. if (!retryAuthCtx && (firstAuthCtx.fixedAccount || args.sameAccountOnly === true)) { + recordUnmovedTransientOutcome(); return { kind: "no-alternate" }; } try { @@ -1394,6 +1433,7 @@ async function retryCodexPoolOnAlternateAccount( writerGeneration: firstAuthCtx.writerGeneration, }); } + recordUnmovedTransientOutcome(); return { kind: "no-alternate" }; } @@ -5669,6 +5709,10 @@ async function handleResponsesInner( // ChatGPT sometimes wraps quota exhaustion in a generic 5xx. Normalize only // body-confirmed cases to quota evidence so cooldown and rotation both apply. poolRetryOutcome = upstreamResponse.status >= 500 ? 429 : upstreamResponse.status; + } else if (!authCtx.fixedAccount && shouldRetryCodexPoolAccountTransient(upstreamResponse)) { + // A plain transient 5xx the same-account retry layer could not absorb. Keep the real + // status so it records as transient rather than quota. + poolRetryOutcome = upstreamResponse.status; } if (poolRetryOutcome !== undefined) { diff --git a/tests/codex-integration/codex-quota-rejection.test.ts b/tests/codex-integration/codex-quota-rejection.test.ts index f413e6bb4b..f9a44a4748 100644 --- a/tests/codex-integration/codex-quota-rejection.test.ts +++ b/tests/codex-integration/codex-quota-rejection.test.ts @@ -1,7 +1,12 @@ 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 { consumeComboFailure, shouldRetryCodexPoolAccountQuota } from "../../src/server/responses/core"; +import { + consumeComboFailure, + shouldRetryCodexPoolAccountQuota, + shouldRetryCodexPoolAccountTransient, +} from "../../src/server/responses/core"; +import { markResponseNonReplayable } from "../../src/lib/upstream-retry"; function jsonRejection(status: number, error: Record): Response { return Response.json({ error }, { status }); @@ -91,6 +96,38 @@ describe("Codex pre-stream quota rejection classification", () => { await expect(shouldRetryCodexPoolAccountQuota(response)).resolves.toBe(false); }); + test.each([ + [500, true], + [502, true], + [503, true], + [504, true], + [520, true], + [507, false], + [429, false], + [400, false], + [200, false], + ])("selects transient pool-account retries by HTTP %i", (status, expected) => { + expect(shouldRetryCodexPoolAccountTransient(new Response(null, { status }))).toBe(expected); + }); + + test("a server_is_overloaded 503 moves to another account even though it carries no quota evidence", () => { + // The shape that wedged a live pool: the backend refuses in under a second, the body says + // nothing about quota, and the account keeps winning selection because nothing recorded a + // failure against it. + const response = Response.json({ + error: { type: "server_error", code: "server_is_overloaded", message: "server is overloaded" }, + }, { status: 503 }); + expect(shouldRetryCodexPoolAccountTransient(response)).toBe(true); + }); + + test("a non-replayable gateway status is never sent from a second account", () => { + // The body already reached the origin, so a second send could duplicate a turn it may + // still be running. This is the one 5xx that stays put. + const response = new Response(null, { status: 502 }); + markResponseNonReplayable(response); + expect(shouldRetryCodexPoolAccountTransient(response)).toBe(false); + }); + test.each([ ["error.message", { error: { message: "The usage limit has been reached" } }], ["last_error.message", { last_error: { message: "The usage limit has been reached" } }], diff --git a/tests/codex-integration/codex-routing.test.ts b/tests/codex-integration/codex-routing.test.ts index 85ecb3a3fd..5a5d0159a7 100644 --- a/tests/codex-integration/codex-routing.test.ts +++ b/tests/codex-integration/codex-routing.test.ts @@ -948,6 +948,50 @@ describe("codex routing", () => { expect(resolveCodexAccountForThread("quota-next", config)).toBe("b"); }); + test("a bound thread does not return to the quota group that refused it once the capped cooldown lapses", () => { + const config = makeConfig(); + const now = 1_800_000_000_000; + updateAccountQuota("a", 10); + updateAccountQuota("b", 20); + expect(resolveCodexAccountForThread("spark-refused", config, now, "spark")).toBe("a"); + + // The refusal announces a window that reopens in four hours, but a reset-derived cooldown is + // capped at 15 minutes so the account is selectable again long before that window moves. + recordCodexUpstreamOutcome(config, "a", 429, { + now, + threadId: "spark-refused", + modelId: "gpt-5.3-codex-spark", + resetAt: Math.floor((now + 4 * 60 * 60_000) / 1_000), + }); + + // Usage is still the lowest in the pool and well under the threshold, so nothing else would + // move this thread: without the refusal it rebinds to the account that just turned it away. + expect(resolveCodexAccountForThread("spark-refused", config, now + 16 * 60_000, "spark")).toBe("b"); + // The shared lane never refused this thread, so a spent Spark window leaves it alone. + expect(resolveCodexAccountForThread("spark-refused", config, now + 16 * 60_000, "shared")).toBe("a"); + }); + + test("a request the account serves releases the threads its quota refusal moved", () => { + const config = makeConfig(); + const now = 1_800_000_000_000; + updateAccountQuota("a", 10); + updateAccountQuota("b", 20); + expect(resolveCodexAccountForThread("spark-recovered", config, now, "spark")).toBe("a"); + recordCodexUpstreamOutcome(config, "a", 429, { + now, + threadId: "spark-recovered", + modelId: "gpt-5.3-codex-spark", + resetAt: Math.floor((now + 4 * 60 * 60_000) / 1_000), + }); + expect(resolveCodexAccountForThread("spark-recovered", config, now + 16 * 60_000, "spark")).toBe("b"); + + // Selection never stopped offering the account to unbound requests, and the first one it + // serves is what ends the refusal — a rebound thread then keeps it across turns. + recordCodexUpstreamOutcome(config, "a", 200, { now: now + 17 * 60_000, modelId: "gpt-5.3-codex-spark" }); + expect(resolveCodexAccountForThread("spark-rebound", config, now + 18 * 60_000, "spark")).toBe("a"); + expect(resolveCodexAccountForThread("spark-rebound", config, now + 19 * 60_000, "spark")).toBe("a"); + }); + test("shared native reset cooldown clears affinity and rotates the active account", () => { const config = makeConfig(); const now = 1_800_000_000_000;