diff --git a/scripts/test-layout/layout.json b/scripts/test-layout/layout.json index 4b3ccee9d53..be3124c38da 100644 --- a/scripts/test-layout/layout.json +++ b/scripts/test-layout/layout.json @@ -1052,6 +1052,7 @@ "privacy-mask-account.test.ts": "lib", "privacy-scan-meta-key.test.ts": "ci-workflows", "probe-lease.test.ts": "routing", + "probe-lease-dispatch-wiring.test.ts": "routing", "process-control-graceful.test.ts": "lib", "process-control.test.ts": "lib", "process-state.test.ts": "service", diff --git a/src/codex/auth-context.ts b/src/codex/auth-context.ts index 54431e0b7c8..f476b89f187 100644 --- a/src/codex/auth-context.ts +++ b/src/codex/auth-context.ts @@ -39,7 +39,12 @@ import { pickAlternateCodexAccount, resolveCodexAccountForThreadDetailed, type CodexAffinityDecision, + type CodexThreadResolution, + type TransientProbeGrant, } from "./routing"; +// The half-open TRANSIENT-HOLD lease (#4701). Not the quota-cooldown probe lease imported from +// ./routing above -- different module, different domain, and a request never holds both. +import { releaseTransientProbe } from "../routing/probe-lease"; import { codexConversationIdentity, recordCodexThreadLineage, @@ -202,6 +207,12 @@ export type CodexAuthContext = affinityDecision?: CodexAffinityDecision; /** Scope that owns `probeLeaseId`, when it is a scoped recovery probe. */ probeQuotaScope?: CodexQuotaScope; + /** + * Set when this request is the ONE dispatch admitted to test an account held under a + * transient 5xx hold (#4701). Echo it into the upstream outcome so the trial is settled + * by the request that ran it, and release it on any path that never reaches upstream. + */ + transientProbe?: TransientProbeGrant; } | { // Main Codex account participating in rotation: token injected from ~/.codex/auth.json @@ -222,6 +233,8 @@ export type CodexAuthContext = probeLeaseId?: string; quotaScope?: CodexQuotaScope; probeQuotaScope?: CodexQuotaScope; + /** See `pool.transientProbe`. */ + transientProbe?: TransientProbeGrant; }; /** Probe lease carried by this context, when it holds one. */ @@ -234,11 +247,24 @@ export function codexProbeQuotaScope(ctx: CodexAuthContext | undefined): CodexQu return ctx?.kind === "pool" || ctx?.kind === "main-pool" ? ctx.probeQuotaScope : undefined; } +/** The transient-hold recovery probe carried by this context, when it holds one (#4701). */ +export function codexTransientProbeGrant(ctx: CodexAuthContext | undefined): TransientProbeGrant | undefined { + return ctx?.kind === "pool" || ctx?.kind === "main-pool" ? ctx.transientProbe : undefined; +} + /** * Hand back a probe lease for a request that will not reach upstream. Safe to * call with a context that holds no lease. + * + * BOTH leases, deliberately. A context can carry the quota-cooldown probe or the transient-hold + * probe, and every one of the ~30 call sites that already hands back the first is a path where + * the second would leak too. Releasing them together is what makes those sites correct for the + * new lease without re-deriving the discard set by hand -- the failure mode being avoided is a + * held account nobody may probe because the request that held the trial went away quietly. */ export function releaseCodexAuthContextProbeLease(ctx: CodexAuthContext | undefined): void { + const transientProbe = codexTransientProbeGrant(ctx); + if (transientProbe) releaseTransientProbe(transientProbe.lease); const leaseId = codexProbeLeaseId(ctx); if (!ctx || ctx.kind === "main" || !leaseId) return; if (ctx.probeQuotaScope) releaseCodexQuotaScopeProbeLease(ctx.accountId!, ctx.probeQuotaScope, leaseId); @@ -424,6 +450,44 @@ export class CodexReserveHelperUnsupportedError extends CodexReserveUnavailableE } } +/** + * Every account bound to this conversation is held after upstream failures, the recovery + * budget for this window is spent, and there is no detour left -- so this request is refused + * BEFORE any upstream I/O (#4701). + * + * This is not a quota cooldown, and the message below says so. It subclasses + * {@link CodexAccountCooldownError} for one reason: the deadline-carrying refusal has exactly + * one representation in this codebase, and roughly a dozen transports already map it to a 429 + * with `Retry-After` and treat it as an expected terminal answer rather than a credential + * fault. Introducing a parallel type would mean either re-deriving that handling in every one + * of them or silently falling through to a 500 in the ones that were missed. + * + * What must NOT be inherited is the quota wording -- "cooling down", `ocx account + * clear-cooldown` -- because none of it describes a 5xx hold and following it would do + * nothing. {@link cooldownErrorMessage} therefore returns this class's own message verbatim, + * the same escape hatch {@link CodexMainAccountHardLockError} and + * {@link CodexReserveUnavailableError} already use. + * + * `cooldownUntil` carries the limiter's own change point, which is strictly in the future: + * either the moment the held account may next be probed or the moment the recovery window + * moves, whichever is later. A refusal that answered `now` would busy-loop the caller into + * the same load it just declined. + */ +export class CodexRecoveryWithheldError extends CodexAccountCooldownError { + /** The sibling still remembered for this thread, when one exists but is itself unusable. */ + readonly detourAccountId?: string; + + constructor(accountId: string, retryAt: number, detourAccountId?: string) { + super(accountId, retryAt); + this.name = "CodexRecoveryWithheldError"; + this.detourAccountId = detourAccountId; + this.message = `Codex account (${cooldownAccountLabel(accountId)}) is held after repeated upstream` + + ` failures and the pool's recovery budget for this window is spent, so nothing was sent` + + ` upstream. Retry after ${new Date(retryAt).toISOString()}.` + + " This clears on its own as the account recovers; no cooldown to lift and no account to switch."; + } +} + export type CodexAuthPolicyConfig = Readonly>; @@ -634,7 +698,12 @@ export function cooldownAccountLabel(accountId: string): string { * injected `openai_base_url` in config.toml. */ export function cooldownErrorMessage(err: CodexAccountCooldownError, accountSelector?: string): string { - if (err instanceof CodexMainAccountHardLockError || err instanceof CodexReserveUnavailableError) return err.message; + if (err instanceof CodexMainAccountHardLockError + || err instanceof CodexReserveUnavailableError + // A transient-hold refusal is not a quota cooldown. Its own wording is the only accurate + // one, and the quota recovery advice below would send the operator after a cooldown that + // does not exist (#4701). + || err instanceof CodexRecoveryWithheldError) return err.message; const until = new Date(err.cooldownUntil).toISOString(); const scopeLabels: Record = { shared: "shared native quota", reserve: "Reserve quota", @@ -870,6 +939,15 @@ export async function resolveCodexAuthContext( // Why this request is on this account, carried to the request log so a move reads as an event // instead of something inferred from account labels across lines (#4546). let affinityDecision: CodexAffinityDecision | undefined; + // The half-open trial this request was granted, if it is the one allowed to test a held + // account. Declared out here because the release paths below and the returned context are on + // opposite sides of several throws (#4701). + let transientProbe: TransientProbeGrant | undefined; + const releaseTransientProbeGrant = (): void => { + if (!transientProbe) return; + releaseTransientProbe(transientProbe.lease); + transientProbe = undefined; + }; // Retained startup recovery makes the physical main identity ineligible. Routing // can still preserve service by selecting a healthy configured pool account. A // request-owned bearer likewise cannot inspect or reconcile file-main state. @@ -917,7 +995,11 @@ export async function resolveCodexAuthContext( // and may still route to non-main pool accounts without touching switch state. if (reserve && !nativeMainReadsForbidden && !selectionAdmission) throw new CodexMainProfileDrainingError(); if (!nativeMainReadsForbidden) reconcileMainCodexAccountRuntimeState(); - const resolution = fixedAccountId !== undefined + // Annotated, not inferred. The two literals below carry neither `affinity` nor + // `transientProbe`, so an inferred union makes `"k" in resolution` widen those reads to + // `unknown` and a discriminant narrowing fail outright. Contextually typing every branch to + // the resolver's own union is what lets the reads below stay total. + const resolution: CodexThreadResolution = fixedAccountId !== undefined ? { status: "selected" as const, accountId: fixedAccountId } : options.excludeAccountId ? (() => { @@ -942,8 +1024,17 @@ export async function resolveCodexAuthContext( lineage, ); if (resolution.status === "expired") throw new CodexThreadAffinityExpiredError(resolution.accountId); + // THE REFUSAL. Every candidate is held, the recovery budget is spent, and no detour is + // left -- so this request must not reach upstream at all. Returning the held account here + // is what #4701 is about: under a provider-wide 503 that is every bound request piling + // onto an account already known to be failing. Thrown before any credential is read, so + // nothing is sent and nothing is spent. + if (resolution.status === "withheld") { + throw new CodexRecoveryWithheldError(resolution.accountId, resolution.retryAt, resolution.detourAccountId); + } const selected = resolution.status === "selected" ? resolution.accountId : null; - affinityDecision = "affinity" in resolution ? resolution.affinity : undefined; + affinityDecision = resolution.affinity; + transientProbe = resolution.status === "selected" ? resolution.transientProbe : undefined; if (!selected) { // A retry that excluded a failed Pool account may still use the validated caller-owned // main credential. Treating every exclusion as if main itself had failed strands a healthy @@ -1025,12 +1116,25 @@ export async function resolveCodexAuthContext( throw new CodexPoolAuthenticationError("Selected Codex account is unavailable"); } } + } catch (cause) { + // Selection granted a trial and then a later policy check refused the account. The trial + // never runs, so hand it back instead of leaving the held account unprobeable until its + // deadline lapses (#4701). + releaseTransientProbeGrant(); + throw cause; } finally { selectionAdmission?.release(); } // Legacy selectors may retain an unusable account for actionable errors. A // deferred credential must never become request auth through that fallback. - assertCodexAccountValidationReady(accountId); + try { + assertCodexAccountValidationReady(accountId); + } catch (cause) { + // Nothing will reach upstream, so give the trial back instead of leaving the held account + // unprobeable until the lease deadline lapses (#4701). + releaseTransientProbeGrant(); + throw cause; + } // Lazy prime: if the selected account has no quota yet, the pool is likely // unprimed (dashboard never opened, or startup prime was blocked). Kick a // best-effort prime so the NEXT routing decision has real scores. This never @@ -1049,6 +1153,13 @@ export async function resolveCodexAuthContext( // a literal Retry-After reads very differently to a user than a reset-derived guess. const cooldown = getCodexQuotaHealthSnapshot(accountId, quotaScope); const cooldownUntil = cooldown?.cooldownUntil; + // A transient-hold trial and a quota cooldown cannot both describe this account: + // `isTransientOnlyAffinityBlock` refuses to recognise a transient hold on an account carrying + // quota health, so the cooldown branch below is unreachable while a trial is held. That is + // also why no request pays two recovery permits for one send. The release is defensive -- + // should that invariant ever move, the trial is handed back rather than stranded behind a + // refusal that belongs to the other domain. + if (cooldownUntil && transientProbe) releaseTransientProbeGrant(); // A cooled-down account never sends traffic, so upstream recovery can never be // observed and the cooldown outlives the real limit. Admit one probe per // interval; its outcome decides whether the cooldown ends (#433). @@ -1089,6 +1200,7 @@ export async function resolveCodexAuthContext( if (token) mainQuotaWriter = observeSelectedMainCredential(token, mainQuotaWriter); assertMainAccountPolicy(policy); } catch (cause) { + releaseTransientProbeGrant(); if (probeLeaseId && probeQuotaScope) releaseCodexQuotaScopeProbeLease(accountId, probeQuotaScope, probeLeaseId); else if (probeLeaseId) releaseCodexQuotaProbeLease(accountId, probeLeaseId); if (cause instanceof CodexMainAccountHardLockError) throw cause; @@ -1099,15 +1211,23 @@ export async function resolveCodexAuthContext( } if (!token) { // Nothing will reach upstream, so give the probe back instead of burning it. + releaseTransientProbeGrant(); if (probeLeaseId && probeQuotaScope) releaseCodexQuotaScopeProbeLease(accountId, probeQuotaScope, probeLeaseId); else if (probeLeaseId) releaseCodexQuotaProbeLease(accountId, probeLeaseId); throw new CodexPoolAuthenticationError( fixedAccountId !== undefined ? "Selected Codex account is unavailable" : undefined, ); } - const reserveAuthorization = reserve - ? await authorizeReserveCredential(token, mainQuotaWriter, policy, options.signal, undefined, writerGeneration) - : undefined; + let reserveAuthorization: MainReserveAuthorization | undefined; + try { + reserveAuthorization = reserve + ? await authorizeReserveCredential(token, mainQuotaWriter, policy, options.signal, undefined, writerGeneration) + : undefined; + } catch (cause) { + // A Reserve refusal ends the request here, so the trial it was holding never runs. + releaseTransientProbeGrant(); + throw cause; + } return { kind: "main-pool", accountId, @@ -1121,6 +1241,7 @@ export async function resolveCodexAuthContext( ...(quotaScope ? { quotaScope } : {}), ...(probeLeaseId ? { probeLeaseId } : {}), ...(probeQuotaScope ? { probeQuotaScope } : {}), + ...(transientProbe ? { transientProbe } : {}), }; } @@ -1141,8 +1262,10 @@ export async function resolveCodexAuthContext( ...(probeLeaseId ? { probeLeaseId } : {}), ...(probeQuotaScope ? { probeQuotaScope } : {}), ...(affinityDecision ? { affinityDecision } : {}), + ...(transientProbe ? { transientProbe } : {}), }; } catch (cause) { + releaseTransientProbeGrant(); if (probeLeaseId && probeQuotaScope) releaseCodexQuotaScopeProbeLease(accountId, probeQuotaScope, probeLeaseId); else if (probeLeaseId) releaseCodexQuotaProbeLease(accountId, probeLeaseId); if (!options.signal?.aborted && shouldMarkAccountNeedsReauthForCodexAuthFailure(cause)) { diff --git a/src/codex/routing.ts b/src/codex/routing.ts index 0a9a0b9c3c8..fbae259958c 100644 --- a/src/codex/routing.ts +++ b/src/codex/routing.ts @@ -54,6 +54,13 @@ import { type CodexUpstreamHealth, } from "./routing/health-store"; import { ownsProbeLease, probeMayClearCooldown, withProbeLeaseReleased } from "./routing/probe-lease"; +// `./routing/probe-lease` above is the QUOTA-COOLDOWN lease; the module below owns the +// unrelated TRANSIENT-HOLD trial and the pool-wide recovery bound above it (#4701). +import { + isTransientHoldExpired, + resolveTransientHoldDispatch, + settleTransientProbeForOutcome, +} from "./routing/transient-hold-dispatch"; import { adoptLegacyLineageAffinity, affinityAfterRelease, @@ -181,6 +188,7 @@ export type { CodexAffinityMove, CodexAffinityReason, CodexAffinityDecision, + TransientProbeGrant, } from "./routing/thread-affinity"; export { isCodexAccountPlanExcluded, @@ -276,12 +284,6 @@ function isTransientOnlyAffinityBlock( || isCodexPoolRefreshCooling(entry.accountId, now); } -/** Has a held binding waited longer than a transient failure can reasonably explain? */ -function isTransientHoldExpired(entry: ThreadAffinityEntry, now: number): boolean { - return entry.transientHoldSince !== undefined - && now - entry.transientHoldSince > CODEX_TRANSIENT_AFFINITY_HOLD_MS; -} - /** * Is every pin this thread holds on the failing account past its hold window? * @@ -477,6 +479,8 @@ export function resolveCodexAccountForThread( lineage?: CodexThreadLineage, ): string | null { const resolution = resolveCodexAccountForThreadDetailed(threadId, config, now, quotaScope, undefined, undefined, lineage); + // A WITHHELD dispatch is deliberately not an account here: this wrapper cannot carry a retry + // time, and answering with the held account is the send the hold prevents. Fails closed. return resolution.status === "selected" ? resolution.accountId : null; } @@ -936,15 +940,12 @@ export function resolveCodexAccountForThreadDetailed( const lane = transientDetourAccount(config, detourEntry, now, quotaScope, selectionOptions); detourEntry.transientHoldSince ??= now; detourEntry.lastUsedAt = now; - if (lane !== null && lane !== detourEntry.accountId) { - detourEntry.transientDetourAccountId = lane; - return { status: "selected", accountId: lane, affinity: { move: "detour", reason: "transient" } }; - } // A provider-wide outage soft-avoids every sibling, so there is nowhere to detour. // That is a statement about where this request can go, not about who owns the // conversation: dropping the pin here would rebuild the cold prefix elsewhere for - // exactly the failure mode the hold exists to survive. - return { status: "selected", accountId: detourEntry.accountId, affinity: { move: "held", reason: "transient" } }; + // exactly the failure the hold exists to survive -- nor a licence to send at the + // failing account, which is what the dispatch resolver bounds (#4701). + return resolveTransientHoldDispatch(detourEntry, lane, now); } // Detour expiry or invalidation must not expire the ordinary task. Drop only // this model lane and select from ordinary/shared state below. @@ -1018,16 +1019,11 @@ export function resolveCodexAccountForThreadDetailed( const detour = transientDetourAccount(config, entry, now, quotaScope, selectionOptions); entry.transientHoldSince ??= now; entry.lastUsedAt = now; - if (detour !== null && detour !== entry.accountId) { - entry.transientDetourAccountId = detour; - // Deliberately no promoteActiveCodexAccount and no rebind: this is one request routing - // around a blip, not the pool deciding where the conversation now lives. - return { status: "selected", accountId: detour, affinity: { move: "detour", reason: "transient" } }; - } // No sibling can take it either -- the usual shape of a provider-wide 503. The binding // survives: "cannot send right now" and "forget which account owns this conversation" - // are different answers, and conflating them is what the hold was added to stop. - return { status: "selected", accountId: entry.accountId, affinity: { move: "held", reason: "transient" } }; + // are different answers. So is the third answer this used to give -- "send at the + // failing account" -- now a bounded probe or a typed refusal (#4701). + return resolveTransientHoldDispatch(entry, detour, now); } // A model-only exclusion does not invalidate the shared task binding. Health, // generation, pause, cooldown, and failure evidence still retire it normally. @@ -1277,6 +1273,10 @@ export function recordCodexUpstreamOutcome( recordUpstreamHostFailure(meta.hostKey, { code: meta.lastFailureCode, now: meta.now ?? Date.now() }); } if (!accountId) return; + // Conclude the half-open recovery trial BEFORE the admissibility gate below (#4701): an + // outcome that gate drops still ended this request, and a lease nobody hands back leaves the + // next trial waiting out its deadline. The settle carries its own fences, so this is safe here. + settleTransientProbeForOutcome(accountId, meta, classifyCodexUpstreamOutcome(outcome, meta.denial)); const writerGeneration = meta.writerGeneration ?? captureConfigGeneration(); if (!isHealthAccountAdmissible(accountId, writerGeneration)) return; const now = meta.now ?? Date.now(); diff --git a/src/codex/routing/cooldown-math.ts b/src/codex/routing/cooldown-math.ts index 123da5e3b16..8b2abe65430 100644 --- a/src/codex/routing/cooldown-math.ts +++ b/src/codex/routing/cooldown-math.ts @@ -5,6 +5,7 @@ import { } from "../quota"; import { isThirtyDayOnlyCodexPlan } from "../plan"; import type { CodexQuotaScope } from "./health-store"; +import type { TransientProbeGrant } from "./thread-affinity"; export const CODEX_DEFAULT_QUOTA_COOLDOWN_MS = 60_000; export const CODEX_MAX_QUOTA_COOLDOWN_MS = 24 * 60 * 60_000; @@ -78,6 +79,15 @@ export type CodexUpstreamOutcomeMeta = { probeLeaseId?: string; /** Scope of `probeLeaseId` when it was granted against a model-scoped cooldown. */ probeQuotaScope?: CodexQuotaScope; + /** + * The half-open TRANSIENT-HOLD probe this request was granted, when it was the one request + * admitted to test a held account (#4701). A different lease to `probeLeaseId` above, in a + * different domain: that one governs a quota cooldown, this one governs a 5xx hold. The two + * are mutually exclusive by construction -- `isTransientOnlyAffinityBlock` refuses to + * recognise a transient hold on an account that carries quota health -- so a request never + * holds both and never pays two recovery permits for one send. + */ + transientProbe?: TransientProbeGrant; /** * Already-chosen alternate for same-request 429 retry. When set, promotion * reuses this account instead of calling {@link pickAlternateCodexAccount} diff --git a/src/codex/routing/thread-affinity.ts b/src/codex/routing/thread-affinity.ts index d8d2f5cbfb7..7557ac6537d 100644 --- a/src/codex/routing/thread-affinity.ts +++ b/src/codex/routing/thread-affinity.ts @@ -4,6 +4,8 @@ import { retainedUtf8Bytes } from "../../lib/admission"; import { clearAllCodexPoolRefreshFailures } from "../pool-refresh-backoff"; import type { CodexThreadLineage } from "../lineage"; import type { CodexQuotaScope } from "./health-store"; +import type { TransientProbeLease } from "../../routing/probe-lease"; +import { clearPoolRecoveryState } from "../../routing/probe-lease"; export type ThreadAffinityEntry = { accountId: string; @@ -25,10 +27,51 @@ export type ThreadAffinityEntry = { transientDetourAccountId?: string; }; +/** + * The half-open trial this request was granted against its own held account (#4701). + * + * The lease alone is not enough to settle safely. Its generation is an account-local PROBE + * epoch, while {@link ThreadAffinityEntry.generation} is the selected CREDENTIAL generation, + * and the two move independently: a credential replaced while the probe is in flight leaves + * the probe epoch untouched, so a settle that checked only the lease would write an answer + * about a credential that no longer exists. Capturing the affinity generation here is what + * lets the settle refuse that case. + */ +export interface TransientProbeGrant { + readonly lease: TransientProbeLease; + /** Credential generation the binding held when the probe was granted. */ + readonly affinityGeneration: number; +} + export type CodexThreadResolution = - | { status: "selected"; accountId: string; affinity?: CodexAffinityDecision } + | { + status: "selected"; + accountId: string; + affinity?: CodexAffinityDecision; + /** + * Present only when this request is the single admitted probe of a held account. The + * holder owes the lease a settle or a release; nothing else may act on it. + */ + transientProbe?: TransientProbeGrant; + } | { status: "none"; affinity?: CodexAffinityDecision } - | { status: "expired"; accountId: string; affinity?: CodexAffinityDecision }; + | { status: "expired"; accountId: string; affinity?: CodexAffinityDecision } + /** + * Every candidate for this binding is held and no detour is left, so there is no account + * this request may be sent to. Distinct from `none`: the binding is REMEMBERED and the + * caller is told when to come back, rather than being handed the account already known to + * be failing. Returning `selected` here is the "must not send, sends anyway" defect + * (#4701); the caller must refuse before any upstream I/O. + */ + | { + status: "withheld"; + accountId: string; + /** Earliest moment a recovery dispatch could be admitted. Always strictly in the future. */ + retryAt: number; + /** The remembered detour, when one exists but is itself unusable right now. */ + detourAccountId?: string; + affinity?: CodexAffinityDecision; + }; /** What happened to this thread's binding on this request (#4546). */ export type CodexAffinityMove = @@ -163,6 +206,11 @@ export function clearThreadAccountMap(): void { // A refresh cooldown is per-account runtime state learned alongside these bindings. Leaving it // behind here keeps an account out of selection after the roster it belonged to is gone. clearAllCodexPoolRefreshFailures(); + // Same argument for recovery state (#4701): probe pacing is keyed on account ids this reset + // may have just retired, and the recovery window counts sends made by the roster that is + // going away. A held account nobody may probe because of a lease issued against the previous + // roster is a recovery that never starts. + clearPoolRecoveryState(); conversationStateIssuerMap.clear(); } diff --git a/src/codex/routing/transient-hold-dispatch.ts b/src/codex/routing/transient-hold-dispatch.ts new file mode 100644 index 00000000000..55ce0020828 --- /dev/null +++ b/src/codex/routing/transient-hold-dispatch.ts @@ -0,0 +1,141 @@ +import { isCodexAccountGenerationLive } from "../account-store"; +// From `../account-id`, which declares the constant and imports nothing, rather than from +// `../main-account`, which re-exports it and sits inside the routing/account-lifecycle import +// cycle. Neither reference here runs at module load, but a leaf import keeps this module out of +// that cycle entirely instead of relying on that staying true. +import { MAIN_CODEX_ACCOUNT_ID } from "../account-id"; +import { + invalidateTransientProbe, + releaseTransientProbe, + resolveHeldAccountDispatch, + settleTransientProbe, +} from "../../routing/probe-lease"; +import { + CODEX_TRANSIENT_AFFINITY_HOLD_MS, + type CodexThreadResolution, + type ThreadAffinityEntry, + type TransientProbeGrant, +} from "./thread-affinity"; +import type { CodexUpstreamOutcomeClass, CodexUpstreamOutcomeMeta } from "./cooldown-math"; + +/** + * What a request bound to a HELD account may actually do, and how its trial ends (#4701). + * + * The transient hold (#4546) keeps a thread's binding while its own account serves a 5xx + * streak and detours the request to a healthy sibling. Both of the selector's hold branches + * used to end the same way when no sibling could take it: they returned the held account as + * `selected`, and the caller sent at an account already known to be failing. Under a + * provider-wide 503 that is every bound request at once -- the amplification the hold exists + * to prevent rather than cause. + * + * This module is the seam between the selector and the bounded answer in + * `src/routing/probe-lease.ts`. It is separate from `./probe-lease` in this same directory, + * which is the unrelated QUOTA-COOLDOWN lease; the two govern different domains and must never + * settle each other's probe. + */ + +/** Has a held binding waited longer than a transient failure can reasonably explain? */ +export function isTransientHoldExpired(entry: ThreadAffinityEntry, now: number): boolean { + return entry.transientHoldSince !== undefined + && now - entry.transientHoldSince > CODEX_TRANSIENT_AFFINITY_HOLD_MS; +} + +/** + * Where a request bound to a held account goes this turn. + * + * {@link resolveHeldAccountDispatch} bounds the answer: one probe tests the held account, and a + * caller with nowhere else to go is WITHHELD and told when to come back rather than sent at the + * failure. + * + * A usable detour is taken BEFORE that resolver is consulted, which inverts its own probe-first + * ordering. Deliberately: a healthy sibling is always a better answer for a live request than an + * account carrying a failure streak, and turning the first request after a hold into the trial + * would spend a real user's turn on it. The ordering is not what #4701 bounds -- the defect is + * the third answer the selector used to give, "send at the failing account anyway", and that is + * reached only when no detour exists. Recovery is still discovered there, because that is + * exactly the case where nothing else can find out. + * + * The caller has already committed `transientHoldSince`/`lastUsedAt`; this decides only where + * the request goes. A withheld answer deliberately leaves `transientDetourAccountId` alone: + * being unable to send right now says nothing about which sibling was serving this thread. + */ +export function resolveTransientHoldDispatch( + entry: ThreadAffinityEntry, + detour: string | null, + now: number, +): CodexThreadResolution { + if (detour !== null && detour !== entry.accountId) { + entry.transientDetourAccountId = detour; + // Deliberately no promotion and no rebind: this is one request routing around a blip, not + // the pool deciding where the conversation now lives. + return { status: "selected", accountId: detour, affinity: { move: "detour", reason: "transient" } }; + } + const dispatch = resolveHeldAccountDispatch({ boundAccountId: entry.accountId, now }); + if (dispatch.kind === "probe") { + return { + status: "selected", + accountId: entry.accountId, + affinity: { move: "held", reason: "transient" }, + // The credential generation travels with the lease so a settle can refuse an answer about + // a credential this binding no longer has. See {@link TransientProbeGrant}. + transientProbe: { lease: dispatch.lease, affinityGeneration: entry.generation }, + }; + } + if (dispatch.kind === "withheld") { + return { + status: "withheld", + accountId: dispatch.boundAccountId, + retryAt: dispatch.retryAt, + // The remembered sibling, when there is one. It is unusable right now -- that is why this + // request is refused -- but it is what has been serving this thread, and a refusal that + // dropped it would make the next resolve re-pick cold. + ...(entry.transientDetourAccountId !== undefined + ? { detourAccountId: entry.transientDetourAccountId } + : {}), + affinity: { move: "held", reason: "transient" }, + }; + } + // Unreachable: no detour was handed in, so the resolver has none to hand back. Kept total + // rather than cast away, because the cost of being wrong here is a send at a failing account. + entry.transientDetourAccountId = dispatch.accountId; + return { status: "selected", accountId: dispatch.accountId, affinity: { move: "detour", reason: "transient" } }; +} + +/** Does the credential a probe was granted against still exist at that generation? */ +function transientProbeCredentialLive(accountId: string, generation: number): boolean { + if (accountId === MAIN_CODEX_ACCOUNT_ID) return generation === 0; + return isCodexAccountGenerationLive(accountId, generation); +} + +/** + * Conclude the half-open recovery probe this request was holding. + * + * Three answers, because three things can be true of a probe that just ended: + * + * - The credential moved under it. Its result describes an identity the binding no longer has, + * so the epoch is BURNED instead of settled -- invalidating makes every outstanding lease on + * this account stale at once, which is what stops a late answer from reviving a dead account. + * - The answer says nothing about the account. A 3xx, a 400, or an unclassifiable status is the + * request's problem, not the account's, so the lease is handed back unspent and the next + * request may run a real trial instead of waiting out a recovery nobody observed. + * - Otherwise it is evidence: success means recovered, everything else means still failing. + * + * A no-op when this request held no trial, so the outcome recorder calls it unconditionally. + */ +export function settleTransientProbeForOutcome( + accountId: string, + meta: Pick, + outcomeClass: CodexUpstreamOutcomeClass, +): void { + const grant: TransientProbeGrant | undefined = meta.transientProbe; + if (!grant) return; + if (!transientProbeCredentialLive(accountId, grant.affinityGeneration)) { + invalidateTransientProbe(accountId); + return; + } + if (outcomeClass === "neutral" || outcomeClass === "caller" || outcomeClass === "unknown") { + releaseTransientProbe(grant.lease); + return; + } + settleTransientProbe(grant.lease, outcomeClass === "success" ? "recovered" : "failed", meta.now ?? Date.now()); +} diff --git a/src/routing/probe-lease.ts b/src/routing/probe-lease.ts index 2182bab3eb2..bec9f5ff0b4 100644 --- a/src/routing/probe-lease.ts +++ b/src/routing/probe-lease.ts @@ -542,3 +542,72 @@ export function configureSharedPoolBackpressure(policy: PoolBackpressurePolicy): export function resetSharedPoolBackpressureForTests(): void { sharedLimiter = undefined; } + +/** + * Forget every account's probe pacing AND the shared recovery window. + * + * Called when the pool's routing state is reset wholesale -- a roster change, a config reload, + * an account removal. Both halves describe a pool that no longer exists: pacing is keyed on + * account ids that may be gone, and the window's buckets count sends made by a roster that + * changed underneath them. Keeping either across such a reset lets one context's recovery + * decisions govern the next one, which is also how it leaks between test files. + * + * This is the production reset. The two `ForTests` seams above stay separate because a test + * frequently wants exactly one half of it. + */ +export function clearPoolRecoveryState(): void { + probeStates.clear(); + sharedLimiter = undefined; +} + +/** + * What one physical send IS, as far as the recovery window is concerned. + * + * The window measures recovery traffic against observed demand, so it needs the distinction + * made where the send happens -- and the transport wrapper cannot make it. That layer sees a + * URL and an init; whether this is a conversation's first attempt, its third retry, or the one + * trial admitted against a held account is knowledge only the caller has. So the caller names + * it, and the classification lives here with the window rather than in the transport, which + * owns no routing policy and has an enforced import boundary saying so. + * + * - `initial`: a new request's first send. Recorded, never refused -- it is the denominator, + * and refusing it would make this a throughput cap rather than a recovery bound. + * - `retry`: a re-send of a request that already reached upstream once. Admitted only while + * recovery traffic stays under its ratio of observed demand. + * - `probe`: the half-open trial against a held account. It ALREADY paid at selection, inside + * {@link resolveHeldAccountDispatch}; charging it again would bill one send twice and shrink + * the very budget it was admitted from. + */ +export type PoolRecoveryDispatchClass = "initial" | "retry" | "probe"; + +export interface PoolRecoveryDispatchDecision { + readonly admitted: boolean; + /** + * Earliest moment another recovery dispatch could be admitted. `now` when the send was + * admitted; otherwise a real change point strictly in the future, so a refused caller has + * something to wait on instead of busy-looping against a pool that is already failing. + */ + readonly retryAt: number; +} + +/** + * Admit one physical send against the process-wide recovery window. + * + * Per-request send budgets cannot see a storm: thousands of requests each staying inside their + * own allowance still compose into an unbounded rate against one failing upstream. This is the + * layer above them, and it is shared by construction. + */ +export function classifyPoolRecoveryDispatch( + dispatchClass: PoolRecoveryDispatchClass, + now = Date.now(), + limiter: PoolBackpressureLimiter = sharedPoolBackpressure(), +): PoolRecoveryDispatchDecision { + if (dispatchClass === "initial") { + limiter.recordInitialSend(now); + return { admitted: true, retryAt: now }; + } + if (dispatchClass === "probe") return { admitted: true, retryAt: now }; + return limiter.tryPermitRetryDispatch(now) + ? { admitted: true, retryAt: now } + : { admitted: false, retryAt: limiter.nextRecoveryAt(now) }; +} diff --git a/src/server/responses/compact.ts b/src/server/responses/compact.ts index 4b9e4178185..3d7f6557c75 100644 --- a/src/server/responses/compact.ts +++ b/src/server/responses/compact.ts @@ -53,6 +53,7 @@ import { resolveCodexAuthContext, codexPoolAffinityKey, codexProbeLeaseId, + codexTransientProbeGrant, codexProbeQuotaScope, releaseCodexAuthContextProbeLease, stripCodexRuntimeProviderFields, @@ -878,6 +879,7 @@ export async function handleResponsesCompact( // replacement (#2887). Also covers the replay's own second 401. ...(ctx.kind === "pool" ? { credentialGeneration: ctx.generation } : {}), probeQuotaScope: codexProbeQuotaScope(ctx), + transientProbe: codexTransientProbeGrant(ctx), writerGeneration: ctx.kind === "pool" || ctx.kind === "main-pool" ? ctx.writerGeneration : undefined, diff --git a/src/server/responses/core-codex-account.ts b/src/server/responses/core-codex-account.ts index e1d1758696c..7d25b86501a 100644 --- a/src/server/responses/core-codex-account.ts +++ b/src/server/responses/core-codex-account.ts @@ -17,6 +17,7 @@ import { resetUpstreamHostHealth, } from "../../codex/upstream-host-health"; import { safeOriginLabel, fetchWithHeaderTimeout, providerFetch } from "./fetch-helpers"; +import { classifyPoolRecoveryDispatch } from "../../routing/probe-lease"; import { formatErrorResponse } from "../../bridge"; import { readBoundedResponseBody } from "../../lib/bounded-body"; import { upstreamErrorMessageFromPayload, isRateLimitOrQuotaFailureMessage } from "../../lib/errors"; @@ -40,6 +41,7 @@ import { MAIN_CODEX_ACCOUNT_ID } from "../../codex/main-account"; import { slugsEquivalent } from "../../providers/slug-codec"; import { codexProbeLeaseId, + codexTransientProbeGrant, codexProbeQuotaScope, releaseCodexAuthContextProbeLease, resolveCodexAuthContext, @@ -470,6 +472,7 @@ export async function retryCodexPoolOnAlternateAccount( modelId: route.modelId, probeLeaseId: codexProbeLeaseId(firstAuthCtx), probeQuotaScope: codexProbeQuotaScope(firstAuthCtx), + transientProbe: codexTransientProbeGrant(firstAuthCtx), writerGeneration: firstAuthCtx.writerGeneration, }); }; @@ -579,6 +582,7 @@ export async function retryCodexPoolOnAlternateAccount( modelId: route.modelId, probeLeaseId: codexProbeLeaseId(firstAuthCtx), probeQuotaScope: codexProbeQuotaScope(firstAuthCtx), + transientProbe: codexTransientProbeGrant(firstAuthCtx), writerGeneration: firstAuthCtx.writerGeneration, }); } @@ -610,6 +614,7 @@ export async function retryCodexPoolOnAlternateAccount( modelId: route.modelId, probeLeaseId: codexProbeLeaseId(firstAuthCtx), probeQuotaScope: codexProbeQuotaScope(firstAuthCtx), + transientProbe: codexTransientProbeGrant(firstAuthCtx), writerGeneration: firstAuthCtx.writerGeneration, // Retry already advanced the RR ring via excludeAccountId — reuse for promotion. ...(retryAuthCtx.accountId ? { promoteAccountId: retryAuthCtx.accountId } : {}), @@ -715,9 +720,28 @@ export async function retryCodexPoolOnAlternateAccount( // The same-account gated-model 400 ladder below keeps its own `maxRetrySends` bound and // does not take the reserve again; only the move itself does. if (accountMovePermit) { + // The pool-wide recovery window is consulted BEFORE the request-local permit is used. + // `reserveDispatch` charges at reservation time and `release()` is the only way back, so + // using the permit first and refusing afterwards would spend a send the request never + // made. An account move is recovery traffic like any other: one request's own budget + // cannot see that a thousand other requests are moving at the same moment, which is + // precisely the amplification this window exists to bound (#4701). + // + // A refusal here is not a new failure mode: "no alternate was available" is already the + // outcome when the pool has nowhere to move this request to, and it is handled. + if (!classifyPoolRecoveryDispatch("retry").admitted) { + accountMovePermit.release(); + accountMovePermit = undefined; + // The alternate context was resolved and will not send. Hand back whatever recovery + // lease it is holding rather than leaving that account unprobeable. + releaseCodexAuthContextProbeLease(retryAuthCtx); + recordUnmovedTransientOutcome(); + return { kind: "no-alternate" }; + } const charged = accountMovePermit.use(); accountMovePermit = undefined; if (!charged) { + releaseCodexAuthContextProbeLease(retryAuthCtx); recordUnmovedTransientOutcome(); return { kind: "no-alternate" }; } @@ -849,6 +873,7 @@ export function codexForwardTerminalOutcomeRecorder( modelId, probeLeaseId: codexProbeLeaseId(authCtx), probeQuotaScope: codexProbeQuotaScope(authCtx), + transientProbe: codexTransientProbeGrant(authCtx), writerGeneration: authCtx.writerGeneration, }); return; @@ -871,6 +896,7 @@ export function codexForwardTerminalOutcomeRecorder( modelId, probeLeaseId: codexProbeLeaseId(authCtx), probeQuotaScope: codexProbeQuotaScope(authCtx), + transientProbe: codexTransientProbeGrant(authCtx), writerGeneration: authCtx.writerGeneration, // A mid-stream terminal can carry a semantic 401 long after the credential was // replaced. It is never replayed — the client already saw output — but it must diff --git a/src/server/responses/passthrough-delivery.ts b/src/server/responses/passthrough-delivery.ts index 79e6ca3d468..1c69c0d6d59 100644 --- a/src/server/responses/passthrough-delivery.ts +++ b/src/server/responses/passthrough-delivery.ts @@ -27,7 +27,7 @@ import type { ResponsesTerminalStatus } from "../../bridge"; import { isCodexWsQuotaObservedResponse, isCodexWsUpstreamResponse } from "./ws-upstream"; import { recordSubagentQuotaFailureForThreadSpawn } from "../../codex/subagent-model-fallback"; import { recordCodexUpstreamOutcome } from "../../codex/routing"; -import { codexProbeLeaseId, codexProbeQuotaScope } from "../../codex/auth-context"; +import { codexProbeLeaseId, codexProbeQuotaScope, codexTransientProbeGrant } from "../../codex/auth-context"; import { consumeComboFailure } from "./core-combo-failure"; import { readDisplaySafeErrorText } from "./core-errors"; import { streamingContextOverflowResponse, jsonContextOverflowResponse } from "./context-overflow"; @@ -251,6 +251,7 @@ export async function deliverPassthroughResponse( modelId: route.modelId, probeLeaseId: codexProbeLeaseId(admissionState.authCtx), probeQuotaScope: codexProbeQuotaScope(admissionState.authCtx), + transientProbe: codexTransientProbeGrant(admissionState.authCtx), writerGeneration: admissionState.authCtx.writerGeneration, // Includes a replay's second 401, which is the case that actually retires the // account — fence it on the credential the request was holding. diff --git a/src/server/responses/passthrough-dispatch.ts b/src/server/responses/passthrough-dispatch.ts index 0292f1d77af..274946be983 100644 --- a/src/server/responses/passthrough-dispatch.ts +++ b/src/server/responses/passthrough-dispatch.ts @@ -29,6 +29,7 @@ import { unwrapUpstreamRetryEvidenceError, codexProbeLeaseId, codexProbeQuotaScope, + codexTransientProbeGrant, createCodexReserveDispatchGuard, } from "../../codex/auth-context"; import { @@ -79,6 +80,7 @@ import { safeHostLabel, storedPoolReplayDispatchNotifier, } from "./fetch-helpers"; +import { classifyPoolRecoveryDispatch } from "../../routing/probe-lease"; import { clientCancelledResponse } from "./core-errors"; import { upstreamHostCircuitOpenResponse, @@ -736,6 +738,7 @@ export async function preparePassthroughExchange( modelId: route.modelId, probeLeaseId: codexProbeLeaseId(admissionState.authCtx), probeQuotaScope: codexProbeQuotaScope(admissionState.authCtx), + transientProbe: codexTransientProbeGrant(admissionState.authCtx), writerGeneration: admissionState.authCtx.writerGeneration, }); } @@ -752,6 +755,12 @@ export async function preparePassthroughExchange( // Body is a replayable string; nothing has streamed to the client yet. upstreamResponse = await fetchWithTransientRetry( recovery => { + // The pool-wide recovery window measures recovery traffic against observed demand, + // and this is where demand is observed: `recovery === undefined` is a new request's + // first send, everything after it is the same request trying again. Without this the + // ratio has no denominator and the window collapses to its quiet-pool floor, which + // would throttle recovery on a busy proxy exactly as hard as on an idle one (#4701). + if (recovery === undefined) classifyPoolRecoveryDispatch("initial"); transportState.noteRoutedAttemptSend(passthroughEstimate, recovery); return fetchWithHeaderTimeout(request.url, applyUpstreamRecoveryInit({ method: request.method, diff --git a/structure/catalog.md b/structure/catalog.md index 31221b4b387..427a67d8764 100644 --- a/structure/catalog.md +++ b/structure/catalog.md @@ -253,6 +253,29 @@ Pool mode routes across main plus added Codex credentials. Key rules: candidate is held the caller gets a typed withheld outcome, not a send. Recovery dispatches (retries and probes, never a new request's initial send) sit under a pool-wide ratio ceiling measured over a sliding window (`src/routing/probe-lease.ts`). + + Where that reaches production, because a primitive nobody calls bounds nothing: the two + transient-hold branches of `resolveCodexAccountForThreadDetailed` + (`src/codex/routing.ts`) ask `resolveHeldAccountDispatch` what this request may do and + return a `withheld` resolution instead of selecting the failing account; + `resolveCodexAuthContext` (`src/codex/auth-context.ts`) turns that into + `CodexRecoveryWithheldError` before any upstream I/O, so a refused request reaches the + client as a 429 carrying the limiter's own change point in `Retry-After`. A granted probe + travels on the auth context, is settled by `recordCodexUpstreamOutcome` under the credential + generation the binding held, and is handed back by `releaseCodexAuthContextProbeLease` on + every path that never sends. The pool window observes demand at the initial passthrough send + and gates the alternate-account replay through `classifyPoolRecoveryDispatch`, which lives + with the window itself rather than in the transport: `src/server/responses/fetch-helpers.ts` + owns no routing policy and `tests/responses/responses-fetch-helpers-boundary.test.ts` pins + its runtime imports to three transport modules. Same-account transient retries remain bounded + by the per-request send budget alone: refusing inside the retry helper's thunk would surface a + pool refusal as a 502 transport failure and record a transient outcome against an account + that was never asked, which is worse than the gap. + + This hold and the quota-cooldown probe (`src/codex/routing/probe-lease.ts`) are different + domains one directory apart. They cannot both describe an account at once, because + `isTransientOnlyAffinityBlock` refuses to recognise a transient hold on an account carrying + quota health -- which is why no request ever pays two recovery permits for one send. - **The credential store is generation-guarded.** A refresh takes a lock and persists only if the generation it started from still holds; a lost race raises a generation-conflict error rather than overwriting the newer credential (`src/codex/account-store.ts`). Callers handle that error; diff --git a/tests/fixtures/test-layout-expected.json b/tests/fixtures/test-layout-expected.json index 5f9ab55449d..b50a6f167cb 100644 --- a/tests/fixtures/test-layout-expected.json +++ b/tests/fixtures/test-layout-expected.json @@ -880,6 +880,7 @@ "privacy-mask-account.test.ts": "lib", "privacy-scan-meta-key.test.ts": "ci-workflows", "probe-lease.test.ts": "routing", + "probe-lease-dispatch-wiring.test.ts": "routing", "process-control-graceful.test.ts": "lib", "process-control.test.ts": "lib", "process-state.test.ts": "service", diff --git a/tests/routing/probe-lease-dispatch-wiring.test.ts b/tests/routing/probe-lease-dispatch-wiring.test.ts new file mode 100644 index 00000000000..9ac2157106e --- /dev/null +++ b/tests/routing/probe-lease-dispatch-wiring.test.ts @@ -0,0 +1,359 @@ +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { existsSync, mkdirSync, readFileSync } from "node:fs"; +import { join } from "node:path"; +import { + canAcquireTransientProbe, + classifyPoolRecoveryDispatch, + clearPoolRecoveryState, + createPoolBackpressureLimiter, + transientProbeDiagnostics, + tryAcquireTransientProbe, + TRANSIENT_PROBE_INTERVAL_MS, + TRANSIENT_PROBE_LEASE_MS, +} from "../../src/routing/probe-lease"; +import { + CodexRecoveryWithheldError, + cooldownErrorMessage, + cooldownErrorResponse, + releaseCodexAuthContextProbeLease, +} from "../../src/codex/auth-context"; +import { + clearCodexUpstreamHealth, + clearThreadAccountMap, + recordCodexUpstreamOutcome, + resolveCodexAccountForThread, + resolveCodexAccountForThreadDetailed, +} from "../../src/codex/routing"; +import { clearPoolRotationState } from "../../src/codex/pool-rotation"; +import { saveCodexAccountCredential } from "../../src/codex/account-store"; +import { clearAccountQuota, updateAccountQuota } from "../../src/codex/auth-api"; +import { repoPath } from "../helpers/repo-root"; +import { removeTreeWithRetry } from "../helpers/remove-tree"; +import type { OcxConfig } from "../../src/types"; + +/** + * The pool-wide recovery limiter, wired to the dispatch that actually sends (#4701). + * + * The primitives in `src/routing/probe-lease.ts` were complete and unit-tested before this, and + * bounded nothing: no file under `src/` imported the module, so every hit for + * `resolveHeldAccountDispatch` was its own definition or a direct unit test. An implementation + * nothing calls is indistinguishable from an absent one at runtime, which is the whole of the + * issue -- and it is why the first case here is a source oracle rather than a behaviour. + * + * The defect that reached production lived at the end of both transient-hold branches of + * `resolveCodexAccountForThreadDetailed`: when no sibling could take the request they returned + * the HELD account as `selected`, and the caller sent it at an account already known to be + * failing. Under a provider-wide 503 that is every bound request at once -- the amplification + * the hold exists to prevent rather than cause. + */ + +const TEST_DIR = join(import.meta.dir, ".tmp-probe-lease-dispatch-wiring"); +let previousOpencodexHome: string | undefined; +let previousCodexHome: string | undefined; + +function makeThreeAccountConfig(overrides: Partial = {}): OcxConfig { + const ids = ["a", "b", "c"]; + for (const id of ids) { + saveCodexAccountCredential(id, { + accessToken: `access-${id}`, + refreshToken: `refresh-${id}`, + expiresAt: Date.now() + 5 * 60_000, + chatgptAccountId: `acct-${id}`, + }); + } + return { + providers: {}, + activeCodexAccountId: "a", + autoSwitchThreshold: 80, + accountPoolStrategy: "quota", + upstreamFailoverThreshold: 3, + codexAccounts: ids.map(id => ({ id, email: `${id}@example.test`, isMain: false })), + ...overrides, + } as OcxConfig; +} + +/** Drive one account to the failover threshold this config declares. */ +function streakTransientFailures(config: OcxConfig, accountId: string, now: number): void { + for (let attempt = 0; attempt < 3; attempt += 1) { + recordCodexUpstreamOutcome(config, accountId, 503, { now }); + } +} + +describe("recovery limiter wiring is reachable from production (#4701)", () => { + test("the transient-hold module is imported by the selector and the dispatch boundary", () => { + // Not a style assertion. Before this change the module had complete unit coverage and zero + // production callers, so the suite was green while nothing in a running proxy was bounded. + // If a refactor ever detaches it again, that is the symptom to catch -- the behaviour tests + // below would keep passing against primitives nobody calls. + const holdDispatch = readFileSync( + repoPath("src", "codex", "routing", "transient-hold-dispatch.ts"), "utf8", + ); + expect(holdDispatch).toContain('from "../../routing/probe-lease"'); + expect(holdDispatch).toContain("resolveHeldAccountDispatch"); + + // The selector reaches the bound through that seam, on the production path. + const routing = readFileSync(repoPath("src", "codex", "routing.ts"), "utf8"); + expect(routing).toContain('from "./routing/transient-hold-dispatch"'); + expect(routing).toContain("resolveTransientHoldDispatch"); + + // The physical-send boundary itself owns no routing policy -- `responses-fetch-helpers- + // boundary.test.ts` pins its runtime imports to three transport modules -- so the dispatch + // call sites name their own class instead. + const passthrough = readFileSync(repoPath("src", "server", "responses", "passthrough-dispatch.ts"), "utf8"); + expect(passthrough).toContain('from "../../routing/probe-lease"'); + expect(passthrough).toContain('classifyPoolRecoveryDispatch("initial")'); + + // Two modules are named probe-lease, one directory apart, and they are different domains. + // The selector keeps importing the QUOTA one; merging them would make one settle the + // other's probe. + expect(routing).toContain('from "./routing/probe-lease"'); + }); +}); + +/** + * Module-scoped, not per-describe. Several cases below read the credential store -- the + * settle's generation fence does, through `isCodexAccountGenerationLive` -- and a test that + * reads the operator's real `~/.opencodex` is both non-deterministic and wrong. + */ +beforeEach(() => { + previousOpencodexHome = process.env.OPENCODEX_HOME; + previousCodexHome = process.env.CODEX_HOME; + if (existsSync(TEST_DIR)) removeTreeWithRetry(TEST_DIR); + mkdirSync(TEST_DIR, { recursive: true }); + process.env.OPENCODEX_HOME = TEST_DIR; + process.env.CODEX_HOME = TEST_DIR; + clearThreadAccountMap(); + clearCodexUpstreamHealth(); + clearAccountQuota(); + clearPoolRotationState(); + clearPoolRecoveryState(); +}); + +afterEach(() => { + clearAccountQuota(); + clearCodexUpstreamHealth(); + clearThreadAccountMap(); + clearPoolRotationState(); + clearPoolRecoveryState(); + if (previousOpencodexHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = previousOpencodexHome; + if (previousCodexHome === undefined) delete process.env.CODEX_HOME; + else process.env.CODEX_HOME = previousCodexHome; + if (existsSync(TEST_DIR)) removeTreeWithRetry(TEST_DIR); +}); + +describe("a held binding with nowhere to detour is bounded, not sent", () => { + test("exactly one request probes the held account; the next is withheld with a future retry time", () => { + const config = makeThreeAccountConfig(); + const threadId = "held-dispatch-thread"; + updateAccountQuota("a", 10); + updateAccountQuota("b", 20); + updateAccountQuota("c", 30); + const start = Date.now(); + expect(resolveCodexAccountForThread(threadId, config, start)).toBe("a"); + + // A provider-wide 503 hits every account, so every sibling is soft-avoided and there is + // nowhere to detour. This is the exact state in which the old code returned the failing + // account to every caller. + for (const id of ["a", "b", "c"]) streakTransientFailures(config, id, start); + + const probe = resolveCodexAccountForThreadDetailed(threadId, config, start); + expect(probe.status).toBe("selected"); + if (probe.status !== "selected") throw new Error("unreachable"); + // Somebody has to find out whether the account is back, and the lease guarantees it is + // exactly one somebody. + expect(probe.accountId).toBe("a"); + expect(probe.transientProbe?.lease.accountId).toBe("a"); + + // The second request in the same instant is NOT a second send at the failing account. + const withheld = resolveCodexAccountForThreadDetailed(threadId, config, start); + expect(withheld.status).toBe("withheld"); + if (withheld.status !== "withheld") throw new Error("unreachable"); + expect(withheld.accountId).toBe("a"); + // Strictly in the future: a refusal that answered `now` would busy-loop the caller into the + // same load it just declined, which is the defect L1 fixed in the resolver itself. + expect(withheld.retryAt).toBeGreaterThan(start); + // The binding is REMEMBERED, not released. "Cannot send right now" and "forget which + // account owns this conversation" are different answers. + expect(withheld.affinity).toMatchObject({ move: "held", reason: "transient" }); + + // The simple wrapper has nowhere to carry a retry time, so it fails closed rather than + // handing back the held account. + expect(resolveCodexAccountForThread(threadId, config, start)).toBeNull(); + + // Once the outage clears the thread is still on its own warm account: a refusal costs the + // conversation nothing, which is the whole point of holding the binding. + expect(resolveCodexAccountForThread(threadId, config, start + 6 * 60_000)).toBe("a"); + }); + + test("a usable sibling still wins over the trial", () => { + const config = makeThreeAccountConfig(); + const threadId = "detour-preferred-thread"; + updateAccountQuota("a", 10); + updateAccountQuota("b", 20); + updateAccountQuota("c", 30); + const start = Date.now(); + expect(resolveCodexAccountForThread(threadId, config, start)).toBe("a"); + + // Only the bound account is failing, so a healthy sibling exists. + streakTransientFailures(config, "a", start); + + const detoured = resolveCodexAccountForThreadDetailed(threadId, config, start); + expect(detoured.status).toBe("selected"); + if (detoured.status !== "selected") throw new Error("unreachable"); + expect(detoured.accountId).not.toBe("a"); + expect(detoured.affinity).toMatchObject({ move: "detour", reason: "transient" }); + // A live request is never spent on the trial while something healthy can serve it, so no + // lease is taken and the recovery budget is untouched. + expect(detoured.transientProbe).toBeUndefined(); + expect(transientProbeDiagnostics("a", start).lastProbeAt).toBeUndefined(); + }); + + test("a quota refusal never becomes a transient trial, so no request pays two permits", () => { + const config = makeThreeAccountConfig(); + const threadId = "quota-domain-thread"; + updateAccountQuota("a", 10); + updateAccountQuota("b", 20); + updateAccountQuota("c", 30); + const start = Date.now(); + expect(resolveCodexAccountForThread(threadId, config, start)).toBe("a"); + + // Transient evidence on every account would normally reach the held branch... + for (const id of ["a", "b", "c"]) streakTransientFailures(config, id, start); + // ...but a quota refusal outranks it. `isTransientOnlyAffinityBlock` refuses to recognise a + // transient hold on an account carrying quota health, which is why the quota-cooldown probe + // and this one can never both describe an account, and why nothing is charged twice. + recordCodexUpstreamOutcome(config, "a", 429, { now: start }); + + const resolved = resolveCodexAccountForThreadDetailed(threadId, config, start); + expect(resolved.status).not.toBe("withheld"); + expect(transientProbeDiagnostics("a", start).held).toBe(false); + expect(transientProbeDiagnostics("a", start).lastProbeAt).toBeUndefined(); + }); +}); + +describe("the trial is handed back on every path that does not send", () => { + test("releasing an auth context frees the account for the next trial", () => { + const now = 4_000_000; + const lease = tryAcquireTransientProbe("release-acct", now)!; + expect(lease.accountId).toBe("release-acct"); + // Held: nobody else may probe while the trial is out. + expect(canAcquireTransientProbe("release-acct", now + TRANSIENT_PROBE_INTERVAL_MS)).toBe(false); + + // This is the single function the ~30 existing "resolved a context, never sent" sites + // already call. Teaching it the second lease is what makes all of them correct at once. + releaseCodexAuthContextProbeLease({ + kind: "pool", + accountId: "release-acct", + writerGeneration: 0, + generation: 1, + accessToken: "token", + chatgptAccountId: "chatgpt-acct", + transientProbe: { lease, affinityGeneration: 1 }, + }); + + // Paced by the interval now, not stranded behind the lease deadline. + expect(canAcquireTransientProbe("release-acct", now + TRANSIENT_PROBE_INTERVAL_MS)).toBe(true); + }); + + test("an unreleased trial still cannot block recovery for longer than its deadline", () => { + const now = 5_000_000; + const lease = tryAcquireTransientProbe("leaked-acct", now)!; + expect(lease.accountId).toBe("leaked-acct"); + + // Nothing settles it and nothing releases it -- the request simply vanished. This is the + // worst case, and it is bounded by construction: a leaked lease delays the next trial, it + // can never cancel it. Permanent blockage would be strictly worse than the unlimited + // behaviour this change replaces, so the deadline is the floor under every release path. + expect(canAcquireTransientProbe("leaked-acct", now + TRANSIENT_PROBE_LEASE_MS - 1)).toBe(false); + expect(canAcquireTransientProbe("leaked-acct", now + TRANSIENT_PROBE_LEASE_MS)).toBe(true); + }); + + test("a settle for a credential the binding no longer has is burned, not applied", () => { + const now = 6_000_000; + // No stored credential exists for this id, so ANY captured generation is already dead -- + // the same shape as a credential replaced while its probe was in flight. + const lease = tryAcquireTransientProbe("rotated-acct", now)!; + const grantedGeneration = transientProbeDiagnostics("rotated-acct", now).generation; + + recordCodexUpstreamOutcome(makeThreeAccountConfig(), "rotated-acct", 200, { + transientProbe: { lease, affinityGeneration: 7 }, + now, + }); + + const after = transientProbeDiagnostics("rotated-acct", now); + // Not recorded as a recovery: the probe answered about an identity this binding lost. + expect(after.lastOutcome).toBeUndefined(); + // The epoch moved instead, which makes every outstanding lease on this account stale at + // once rather than waiting for each deadline. + expect(after.generation).toBeGreaterThan(grantedGeneration); + expect(after.held).toBe(false); + }); +}); + +describe("a withheld dispatch reaches the client as a bounded refusal", () => { + test("it answers 429 with Retry-After and never borrows the quota-cooldown wording", () => { + const now = 7_000_000; + const error = new CodexRecoveryWithheldError("acct-held", now + 30_000, "acct-detour"); + expect(error.detourAccountId).toBe("acct-detour"); + + const response = cooldownErrorResponse(error, now); + expect(response.status).toBe(429); + expect(response.headers.get("Retry-After")).toBe("30"); + + // Subclassing the cooldown error buys the transport mapping above. It must not also buy the + // quota advice: there is no cooldown to lift and no account to switch to, so following it + // would waste the operator's time on a fix for a different problem. + expect(cooldownErrorMessage(error)).toBe(error.message); + expect(error.message).not.toContain("cooling down"); + expect(error.message).not.toContain("clear-cooldown"); + expect(error.message).toContain("nothing was sent"); + }); +}); + +describe("the pool-wide window classifies one physical send", () => { + test("demand is counted, a retry is gated, and a probe is never charged twice", () => { + const now = 8_000_000; + const limiter = createPoolBackpressureLimiter({ + windowMs: 10_000, + maxRetryRatio: 0.2, + minRecoveryAllowance: 1, + }); + + // The denominator. Refusing a new request's first send would make this a throughput cap + // rather than a recovery bound. + expect(classifyPoolRecoveryDispatch("initial", now, limiter).admitted).toBe(true); + expect(limiter.state(now).initialSends).toBe(1); + expect(limiter.state(now).recoveryDispatches).toBe(0); + + // A probe already paid at selection, inside `resolveHeldAccountDispatch`. Charging it again + // here would bill one send twice and shrink the budget it was admitted from. + expect(classifyPoolRecoveryDispatch("probe", now, limiter).admitted).toBe(true); + expect(limiter.state(now).recoveryDispatches).toBe(0); + + // One recovery dispatch fits the allowance; the next does not. + expect(classifyPoolRecoveryDispatch("retry", now, limiter).admitted).toBe(true); + expect(limiter.state(now).recoveryDispatches).toBe(1); + + const refused = classifyPoolRecoveryDispatch("retry", now, limiter); + expect(refused.admitted).toBe(false); + // A refusal has to hand back a time, or the caller busy-loops against a pool that is + // already failing -- which is the load this window exists to remove. + expect(refused.retryAt).toBeGreaterThan(now); + expect(limiter.state(now).refusedTotal).toBe(1); + }); + + test("two independent requests draw on one window, not one allowance each", () => { + const now = 9_000_000; + const limiter = createPoolBackpressureLimiter({ + windowMs: 10_000, + maxRetryRatio: 0.2, + minRecoveryAllowance: 1, + }); + // Per-request budgets cannot see a storm: each request staying inside its own allowance + // still composes into an unbounded rate against one failing upstream. Distinct requests + // share this window by construction. + expect(classifyPoolRecoveryDispatch("retry", now, limiter).admitted).toBe(true); + expect(classifyPoolRecoveryDispatch("retry", now, limiter).admitted).toBe(false); + }); +});