From 5e2f1afb4fa7961dd29ed1a199f997d02a3f884a Mon Sep 17 00:00:00 2001 From: luvs01 <27862058+luvs01@users.noreply.github.com> Date: Wed, 26 Aug 2026 07:24:00 +0900 Subject: [PATCH 1/6] fix(codex): close drain routing follow-ups --- src/codex/auth-context.ts | 14 +- src/codex/model-entitlements.ts | 12 + src/codex/routing.ts | 334 ++++++++-- src/codex/subagent-model-fallback.ts | 86 ++- src/server/responses/core.ts | 178 +++++- tests/codex-auth-context.test.ts | 109 +++- tests/codex-routing.test.ts | 405 ++++++++++++ ...subagent-fallback-handle-responses.test.ts | 602 +++++++++++++++++- 8 files changed, 1605 insertions(+), 135 deletions(-) diff --git a/src/codex/auth-context.ts b/src/codex/auth-context.ts index f3c357d4c3..acfb82ccfb 100644 --- a/src/codex/auth-context.ts +++ b/src/codex/auth-context.ts @@ -383,6 +383,8 @@ export async function resolveCodexAuthContext( const nativeMainTrafficBlocked = isNativeMainTrafficBlocked(); const selectionAdmission = options.beginCodexAccountSelection?.(); const nativeMainReadsForbidden = nativeMainTrafficBlocked || selectionAdmission?.mainProfileDraining === true; + const nativeMainSelectionOnly = !nativeMainTrafficBlocked + && selectionAdmission?.mainProfileDraining === true; let accountId: string; const quotaScope = codexQuotaScopeForModel(options.modelId); try { @@ -401,8 +403,7 @@ export async function resolveCodexAuthContext( const selectionOptions = { // Temporary switch drain keeps the candidate until the atomic claim rejects // it. Retained recovery makes main wholly ineligible so pool routing continues. - nativeMainSelectionOnly: !nativeMainTrafficBlocked - && selectionAdmission?.mainProfileDraining === true, + nativeMainSelectionOnly, isMainAccountTokenLive: options.isMainAccountTokenLive, modelEligibleAccountIds, }; @@ -436,12 +437,13 @@ export async function resolveCodexAuthContext( : "Selected Codex account is unavailable", ); } - // Recovery deliberately makes physical main ineligible. If no healthy - // pool route is configured and main is the intended route, report the - // temporary fence rather than misclassifying that credential as invalid. + // Recovery or a turn drain deliberately makes physical main unobservable. + // If no healthy pool route is available, report the temporary fence rather + // than turning a credential we were forbidden to inspect into a permanent + // model-entitlement denial. // 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 (nativeMainReadsForbidden && !options.excludeAccountId) { throw new CodexMainProfileDrainingError(); } throw new CodexPoolAuthenticationError( diff --git a/src/codex/model-entitlements.ts b/src/codex/model-entitlements.ts index 5a649d335a..1bc6f43285 100644 --- a/src/codex/model-entitlements.ts +++ b/src/codex/model-entitlements.ts @@ -46,6 +46,18 @@ export interface CodexModelEntitlementResolveOptions { readonly excludeAccountIds?: ReadonlySet; } +/** + * Explicit request-boundary signal for an operational discovery failure that cannot be + * represented by the ordinary fail-closed `confirmed: false` snapshot. Generic throws are + * programming errors and must retain their original error path. + */ +export class CodexModelEntitlementDiscoveryUnavailableError extends Error { + constructor(cause?: unknown) { + super("Codex model entitlement discovery is temporarily unavailable", { cause }); + this.name = "CodexModelEntitlementDiscoveryUnavailableError"; + } +} + const accountModelsCache = new Map(); const accountModelsFlights = new Map>(); diff --git a/src/codex/routing.ts b/src/codex/routing.ts index 10160fe913..b9ad137390 100644 --- a/src/codex/routing.ts +++ b/src/codex/routing.ts @@ -964,7 +964,7 @@ function getEligiblePoolAccounts( return selectPriorityTier( ids, codexAccountPriorityLookup(config), - id => hasCodexQuotaHeadroom(config, id), + id => hasCodexQuotaHeadroom(config, id, selectionOptions), pinnedCodexAccountId(config), ); } @@ -992,10 +992,17 @@ function stickyLimitForConfig(config: OcxConfig): number { * primed. A genuinely exhausted account 429s into cooldown and leaves * eligibility on its own. */ -function hasCodexQuotaHeadroom(config: OcxConfig, accountId: string): boolean { +function hasCodexQuotaHeadroom( + config: OcxConfig, + accountId: string, + selectionOptions?: CodexAccountUsabilityOptions, +): boolean { const threshold = config.autoSwitchThreshold ?? 80; if (threshold <= 0) return true; - const usage = computeCodexUsageScore(getAccountQuota(accountId), getPoolAccountPlan(config, accountId)); + const usage = computeCodexUsageScore( + getAccountQuota(accountId), + getPoolAccountPlanForSelection(config, accountId, selectionOptions), + ); if (isUnknownUsage(usage)) return true; return usage < threshold; } @@ -1014,7 +1021,7 @@ function pickFillFirstCodexAccount( if (eligible.length === 0) return null; const active = getEffectiveActiveCodexAccountId(config); - if (active && eligible.includes(active) && hasCodexQuotaHeadroom(config, active)) { + if (active && eligible.includes(active) && hasCodexQuotaHeadroom(config, active, selectionOptions)) { return active; } @@ -1034,7 +1041,7 @@ function pickNextFillFirstCodexAccount( if (!afterId) { // Prefer an under-threshold account when starting with no active cursor. for (const id of ordered) { - if (hasCodexQuotaHeadroom(config, id)) return id; + if (hasCodexQuotaHeadroom(config, id, selectionOptions)) return id; } return ordered[0] ?? null; } @@ -1049,7 +1056,7 @@ function pickNextFillFirstCodexAccount( const startIdx = stableAll.indexOf(afterId); if (startIdx < 0) { for (const id of ordered) { - if (hasCodexQuotaHeadroom(config, id)) return id; + if (hasCodexQuotaHeadroom(config, id, selectionOptions)) return id; } return ordered[0] ?? null; } @@ -1060,7 +1067,7 @@ function pickNextFillFirstCodexAccount( const candidate = stableAll[(startIdx + step) % stableAll.length]!; if (!eligible.includes(candidate)) continue; if (!fallback) fallback = candidate; - if (hasCodexQuotaHeadroom(config, candidate)) return candidate; + if (hasCodexQuotaHeadroom(config, candidate, selectionOptions)) return candidate; } return fallback ?? ordered[0] ?? null; } @@ -1069,10 +1076,10 @@ function pickNextFillFirstCodexAccount( * Unbound new-session pick for round-robin / fill-first. Returns null to fall through * to the legacy quota path (or when the strategy is quota). * - * When `commit` is true (resolve path), remembers active in-memory, binds thread affinity, and - * notes RR success. When `commit` is false (preview), returns the same RR/fill-first - * account resolve would pick via a dry-run peek — without mutating ring weights, - * activeKey, sticky counters, config, or affinity. + * When `commit` is true (resolve path), advances RR state. `commitSharedActive` + * and `commitAffinity` independently control the two cross-request side effects: + * model-scoped entitlement selection can bind a new task without replacing an + * existing task binding or global active choice. Preview remains a dry-run peek. * * Automatic strategy picks never sync-write config; only manual selection persists active. * @@ -1087,6 +1094,8 @@ function pickUnboundStrategyAccount( commit: boolean, quotaScope?: CodexQuotaScope, selectionOptions?: CodexAccountUsabilityOptions, + commitSharedActive = commit, + commitAffinity = commit, ): string | null { const strategy = normalizeAccountPoolStrategy(config.accountPoolStrategy); if (strategy === "quota") return null; @@ -1101,8 +1110,10 @@ function pickUnboundStrategyAccount( } picked = pickRoundRobinAccount(poolKey, eligible, limit); if (!picked) return null; - if (!isIndependentCodexQuotaScope(quotaScope)) rememberActiveCodexAccount(config, picked); - if (threadId) bindThreadAffinity(threadId, picked, now, quotaScope); + if (commitSharedActive) { + if (!isIndependentCodexQuotaScope(quotaScope)) rememberActiveCodexAccount(config, picked); + } + if (commitAffinity && threadId) bindThreadAffinity(threadId, picked, now, quotaScope); notePoolRotationSuccess(poolKey, picked, limit); return picked; } @@ -1110,10 +1121,10 @@ function pickUnboundStrategyAccount( if (strategy === "fill-first") { picked = pickFillFirstCodexAccount(config, now, quotaScope, selectionOptions); if (!picked) return null; - if (commit) { + if (commitSharedActive) { if (!isIndependentCodexQuotaScope(quotaScope)) rememberActiveCodexAccount(config, picked); - if (threadId) bindThreadAffinity(threadId, picked, now, quotaScope); } + if (commitAffinity && threadId) bindThreadAffinity(threadId, picked, now, quotaScope); return picked; } @@ -1126,6 +1137,36 @@ export function getPoolAccountPlan(config: OcxConfig, accountId: string): string .find(account => isSelectableCodexPoolAccount(account) && account.id === accountId)?.plan; } +/** Selection-only main routing must not lazily read the fenced native credential for its plan. */ +function getPoolAccountPlanForSelection( + config: OcxConfig, + accountId: string, + selectionOptions?: CodexAccountUsabilityOptions, +): string | undefined { + if (accountId === MAIN_CODEX_ACCOUNT_ID && selectionOptions?.nativeMainSelectionOnly === true) { + return undefined; + } + return getPoolAccountPlan(config, accountId); +} + +/** Shared routing state must ignore a request-scoped entitlement roster. */ +function sharedStateSelectionOptions( + selectionOptions?: CodexAccountUsabilityOptions, +): Pick< + CodexAccountUsabilityOptions, + "nativeMainSelectionOnly" | "isMainAccountTokenLive" +> | undefined { + if (!selectionOptions) return undefined; + return { + ...(selectionOptions.nativeMainSelectionOnly !== undefined + ? { nativeMainSelectionOnly: selectionOptions.nativeMainSelectionOnly } + : {}), + ...(selectionOptions.isMainAccountTokenLive + ? { isMainAccountTokenLive: selectionOptions.isMainAccountTokenLive } + : {}), + }; +} + function pickLowerUsageAccount( config: OcxConfig, active: string, @@ -1137,7 +1178,10 @@ function pickLowerUsageAccount( let best = active; let bestUsage = activeUsage; for (const id of getEligiblePoolAccounts(config, active, now, quotaScope, selectionOptions)) { - const usage = computeCodexUsageScore(getAccountQuota(id), getPoolAccountPlan(config, id)); + const usage = computeCodexUsageScore( + getAccountQuota(id), + getPoolAccountPlanForSelection(config, id, selectionOptions), + ); if (usage < bestUsage) { best = id; bestUsage = usage; @@ -1147,11 +1191,18 @@ function pickLowerUsageAccount( } /** Coolest account in an already-selected candidate list; first index wins ties. */ -function pickLowestUsageAmong(config: OcxConfig, ids: readonly string[]): string | null { +function pickLowestUsageAmong( + config: OcxConfig, + ids: readonly string[], + selectionOptions?: CodexAccountUsabilityOptions, +): string | null { let best: string | null = null; let bestUsage = Number.POSITIVE_INFINITY; for (const id of ids) { - const usage = computeCodexUsageScore(getAccountQuota(id), getPoolAccountPlan(config, id)); + const usage = computeCodexUsageScore( + getAccountQuota(id), + getPoolAccountPlanForSelection(config, id, selectionOptions), + ); if (usage < bestUsage) { best = id; bestUsage = usage; @@ -1167,7 +1218,11 @@ export function pickLowestUsageCodexAccount( quotaScope?: CodexQuotaScope, selectionOptions?: CodexAccountUsabilityOptions, ): string | null { - return pickLowestUsageAmong(config, getEligiblePoolAccounts(config, excludeId, now, quotaScope, selectionOptions)); + return pickLowestUsageAmong( + config, + getEligiblePoolAccounts(config, excludeId, now, quotaScope, selectionOptions), + selectionOptions, + ); } /** @@ -1307,12 +1362,20 @@ function pickPriorityPreemption( // A live pin already lowered the tier ceiling; never preempt past an explicit // operator choice. Same liveness test the tier filter applies, so preview and // resolve agree even before the pin is garbage-collected. - if (pinned !== undefined && eligible.includes(pinned) && hasCodexQuotaHeadroom(config, pinned)) return null; + if ( + pinned !== undefined + && eligible.includes(pinned) + && hasCodexQuotaHeadroom(config, pinned, selectionOptions) + ) return null; const priorityOf = codexAccountPriorityLookup(config); if (priorityOf(eligible[0]!) <= priorityOf(active)) return null; // Members without headroom are in the tier only because a sibling has some; // picking one would hand the request straight back to a drained account. - return pickLowestUsageAmong(config, eligible.filter(id => hasCodexQuotaHeadroom(config, id))); + return pickLowestUsageAmong( + config, + eligible.filter(id => hasCodexQuotaHeadroom(config, id, selectionOptions)), + selectionOptions, + ); } /** @@ -1322,13 +1385,27 @@ function pickPriorityPreemption( * on its own. Clearing the pin also removes the condition, so this writes at * most once per pin. */ -function releaseDrainedCodexAccountPin(config: OcxConfig): void { +function releaseDrainedCodexAccountPin( + config: OcxConfig, + selectionOptions?: Pick< + CodexAccountUsabilityOptions, + "nativeMainSelectionOnly" | "isMainAccountTokenLive" + >, +): void { const pinned = pinnedCodexAccountId(config); if (pinned === undefined) return; - const drained = !isCodexAccountUsable(config, pinned) - || isAccountNeedsReauth(pinned) - || isCodexAccountPaused(config, pinned) - || !hasCodexQuotaHeadroom(config, pinned); + const knownUnavailable = isAccountNeedsReauth(pinned) || isCodexAccountPaused(config, pinned); + if (knownUnavailable) { + clearCodexAccountPin(config); + saveConfigPreservingClaudeCode(config); + return; + } + // Temporary drain deliberately forbids every native-main read. A pin on main + // cannot be classified by credential liveness or quota until the fenced profile + // is readable. Cached reauth and configured pause state were handled above. + if (pinned === MAIN_CODEX_ACCOUNT_ID && selectionOptions?.nativeMainSelectionOnly === true) return; + const drained = !isCodexAccountUsable(config, pinned, selectionOptions) + || !hasCodexQuotaHeadroom(config, pinned, selectionOptions); if (!drained) return; clearCodexAccountPin(config); saveConfigPreservingClaudeCode(config); @@ -1340,18 +1417,24 @@ function applyQuotaAutoSwitch( now: number, quotaScope?: CodexQuotaScope, selectionOptions?: CodexAccountUsabilityOptions, + commitSharedSelection = true, ): string { const threshold = config.autoSwitchThreshold ?? 80; if (threshold <= 0) return active; const quota = getAccountQuota(active); - const activeUsage = computeCodexUsageScore(quota, getPoolAccountPlan(config, active)); + const activeUsage = computeCodexUsageScore( + quota, + getPoolAccountPlanForSelection(config, active, selectionOptions), + ); // Unknown usage is not evidence that a user's explicit selection crossed the // threshold. Wait for quota priming instead of rotating among guesses. if (isUnknownUsage(activeUsage)) return active; if (activeUsage < threshold) return active; const best = pickLowerUsageAccount(config, active, activeUsage, now, quotaScope, selectionOptions); if (best !== active) { - if (!isIndependentCodexQuotaScope(quotaScope)) setActiveCodexAccount(config, best); + if (commitSharedSelection && !isIndependentCodexQuotaScope(quotaScope)) { + setActiveCodexAccount(config, best); + } return best; } @@ -1366,12 +1449,25 @@ function shouldFailover(config: OcxConfig, accountId: string, now: number): bool return !!health && health.consecutiveFailures >= threshold; } +function isHealthySharedCodexSelection( + config: OcxConfig, + accountId: string, + now: number, + quotaScope: CodexQuotaScope | undefined, + selectionOptions: CodexAccountUsabilityOptions | undefined, +): boolean { + return isCodexAccountSelectable(config, accountId, now, quotaScope, selectionOptions) + && hasCodexQuotaHeadroom(config, accountId, selectionOptions) + && !shouldFailover(config, accountId, now); +} + function applyFailureFailover( config: OcxConfig, active: string, now: number, quotaScope?: CodexQuotaScope, selectionOptions?: CodexAccountUsabilityOptions, + commitSharedSelection = true, ): string { if (!shouldFailover(config, active, now)) return active; const best = pickAlternateCodexAccount(config, active, now, quotaScope, selectionOptions); @@ -1382,7 +1478,9 @@ function applyFailureFailover( // the moment of the failure; the streak outlives the soft avoid, so a later // scoped resolve reaches here with the streak still tripped and would otherwise // move the shared cursor after all. - if (!isIndependentCodexQuotaScope(quotaScope)) promoteActiveCodexAccount(config, best); + if (commitSharedSelection && !isIndependentCodexQuotaScope(quotaScope)) { + promoteActiveCodexAccount(config, best); + } return best; } return active; @@ -1429,7 +1527,7 @@ export function previewCodexAccountForRequest( if (threshold > 0) { const usage = computeCodexUsageScore( getAccountQuota(entry.accountId), - getPoolAccountPlan(config, entry.accountId), + getPoolAccountPlanForSelection(config, entry.accountId, selectionOptions), ); if (!isUnknownUsage(usage) && usage >= threshold) { const best = pickLowerUsageAccount( @@ -1476,7 +1574,10 @@ export function previewCodexAccountForRequest( const threshold = config.autoSwitchThreshold ?? 80; if (threshold > 0) { - const usage = computeCodexUsageScore(getAccountQuota(active), getPoolAccountPlan(config, active)); + const usage = computeCodexUsageScore( + getAccountQuota(active), + getPoolAccountPlanForSelection(config, active, selectionOptions), + ); if (!isUnknownUsage(usage) && usage >= threshold) { active = pickLowerUsageAccount(config, active, usage, now, quotaScope, selectionOptions); } @@ -1502,11 +1603,31 @@ export function resolveCodexAccountForThreadDetailed( quotaScope?: CodexQuotaScope, selectionOptions?: CodexAccountUsabilityOptions, ): CodexThreadResolution { + // An entitlement roster constrains only this model request. It must not rewrite + // the operator's shared active/pin choice or the task's ordinary-model affinity. + const modelScopedSelection = selectionOptions?.modelEligibleAccountIds !== undefined; + let preserveExistingModelScopedAffinity = false; + const sharedSelectionOptions: CodexAccountUsabilityOptions | undefined = modelScopedSelection + ? sharedStateSelectionOptions(selectionOptions) ?? {} + : selectionOptions; // Retiring a spent manual pin is independent of affinity: an existing thread // keeps its account below, but the operator's tier ceiling must not silently // revive after quota resets. Independent model scopes must never persist a // change to shared routing state. - if (!isIndependentCodexQuotaScope(quotaScope)) releaseDrainedCodexAccountPin(config); + if (!isIndependentCodexQuotaScope(quotaScope)) { + releaseDrainedCodexAccountPin(config, sharedStateSelectionOptions(selectionOptions)); + } + const sharedActiveBeforeSelection = getEffectiveActiveCodexAccountId(config); + const preserveSharedSelectionForModelDetour = modelScopedSelection && ( + sharedActiveBeforeSelection === undefined + || isHealthySharedCodexSelection( + config, + sharedActiveBeforeSelection, + now, + quotaScope, + sharedSelectionOptions, + ) + ); const entry = threadId ? getThreadAffinity(threadId, quotaScope) : undefined; if (threadId && entry) { @@ -1514,12 +1635,20 @@ export function resolveCodexAccountForThreadDetailed( deleteThreadAffinity(threadId, quotaScope); return { status: "expired", accountId: entry.accountId }; } + const generationLive = isThreadAffinityGenerationLive(entry); + const selectableForSharedState = generationLive + && isCodexAccountSelectable(config, entry.accountId, now, quotaScope, sharedSelectionOptions); + const selectableForRequest = selectableForSharedState + && isCodexAccountSelectable(config, entry.accountId, now, quotaScope, selectionOptions); + const failoverReady = shouldFailover(config, entry.accountId, now); + const healthyForSharedAffinity = selectableForSharedState + && hasCodexQuotaHeadroom(config, entry.accountId, sharedSelectionOptions) + && !failoverReady; if ( - isThreadAffinityGenerationLive(entry) - && isCodexAccountSelectable(config, entry.accountId, now, quotaScope, selectionOptions) + selectableForRequest // 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). - && !shouldFailover(config, entry.accountId, now) + && !failoverReady ) { entry.lastUsedAt = now; // Periodic quota re-eval: a long-lived bound thread must still switch when @@ -1533,11 +1662,11 @@ export function resolveCodexAccountForThreadDetailed( const strategy = normalizeAccountPoolStrategy(config.accountPoolStrategy); if (strategy === "quota") { const threshold = config.autoSwitchThreshold ?? 80; - const usage = threshold > 0 - ? computeCodexUsageScore( - getAccountQuota(entry.accountId), - getPoolAccountPlan(config, entry.accountId), - ) + const usage = threshold > 0 + ? computeCodexUsageScore( + getAccountQuota(entry.accountId), + getPoolAccountPlanForSelection(config, entry.accountId, selectionOptions), + ) : 0; const overThreshold = threshold > 0 && !isUnknownUsage(usage) && usage >= threshold; if (overThreshold || now - entry.lastReevalAt >= CODEX_THREAD_AFFINITY_REEVAL_INTERVAL_MS) { @@ -1545,7 +1674,9 @@ export function resolveCodexAccountForThreadDetailed( if (overThreshold) { const best = pickLowerUsageAccount(config, entry.accountId, usage, now, quotaScope, selectionOptions); if (best !== entry.accountId) { - if (!isIndependentCodexQuotaScope(quotaScope)) setActiveCodexAccount(config, best); + if (!isIndependentCodexQuotaScope(quotaScope)) { + setActiveCodexAccount(config, best); + } bindThreadAffinity(threadId, best, now, quotaScope); // rebinds + resets clocks return { status: "selected", accountId: best }; } @@ -1554,24 +1685,104 @@ export function resolveCodexAccountForThreadDetailed( } return { status: "selected", accountId: entry.accountId }; } - deleteThreadAffinity(threadId, quotaScope); + // A model-only exclusion does not invalidate the shared task binding. Health, + // generation, pause, cooldown, and failure evidence still retire it normally. + if (!modelScopedSelection || !healthyForSharedAffinity) { + deleteThreadAffinity(threadId, quotaScope); + } else { + preserveExistingModelScopedAffinity = true; + } } - const strategyPick = pickUnboundStrategyAccount(config, threadId, now, true, quotaScope, selectionOptions); - if (strategyPick) return { status: "selected", accountId: strategyPick }; + // A request-scoped roster may still contain unhealthy candidates. Non-quota strategies return + // before the quota/failover helpers below, so prefer only shared-healthy roster members here; + // otherwise RR/fill-first can immediately re-pick a known failing account even when another + // entitled account is healthy. If no healthy member exists, the normal fallback path below + // still decides whether the sole eligible candidate must be used. + const strategySelectionOptions = modelScopedSelection + ? { + ...selectionOptions, + modelEligibleAccountIds: new Set( + [...selectionOptions!.modelEligibleAccountIds!].filter(accountId => + isHealthySharedCodexSelection( + config, + accountId, + now, + quotaScope, + sharedSelectionOptions, + ) + ), + ), + } + : selectionOptions; + const strategyPick = pickUnboundStrategyAccount( + config, + threadId, + now, + true, + quotaScope, + strategySelectionOptions, + !modelScopedSelection, + !preserveExistingModelScopedAffinity, + ); + if (strategyPick) { + if ( + modelScopedSelection + && !preserveSharedSelectionForModelDetour + && !isIndependentCodexQuotaScope(quotaScope) + ) { + promoteActiveCodexAccount(config, strategyPick); + } + return { status: "selected", accountId: strategyPick }; + } let active = getEffectiveActiveCodexAccountId(config); if (!active) { const selected = pickLowestUsageCodexAccount(config, undefined, now, quotaScope, selectionOptions); - if (!selected) return { status: "none" }; - if (!isIndependentCodexQuotaScope(quotaScope)) setActiveCodexAccount(config, selected); + if (!selected) { + if ( + selectionOptions?.nativeMainSelectionOnly === true + && selectionOptions.modelEligibleAccountIds !== undefined + ) { + return { status: "selected", accountId: MAIN_CODEX_ACCOUNT_ID }; + } + return { status: "none" }; + } + if (!isIndependentCodexQuotaScope(quotaScope) && !modelScopedSelection) { + setActiveCodexAccount(config, selected); + } active = selected; } + const activeSelectableForSharedState = isCodexAccountSelectable( + config, + active, + now, + quotaScope, + sharedSelectionOptions, + ); + const activeHealthyForSharedSelection = activeSelectableForSharedState + && hasCodexQuotaHeadroom(config, active, sharedSelectionOptions) + && !shouldFailover(config, active, now); if (!isCodexAccountSelectable(config, active, now, quotaScope, selectionOptions)) { const fallback = pickLowestUsageCodexAccount(config, active, now, quotaScope, selectionOptions); if (fallback) { - if (!isIndependentCodexQuotaScope(quotaScope)) setActiveCodexAccount(config, fallback); + const modelOnlyMove = modelScopedSelection + && preserveSharedSelectionForModelDetour + && activeHealthyForSharedSelection; + if (!isIndependentCodexQuotaScope(quotaScope) && !modelOnlyMove) { + setActiveCodexAccount(config, fallback); + } active = fallback; + } else if ( + selectionOptions?.nativeMainSelectionOnly === true + && selectionOptions.modelEligibleAccountIds !== undefined + ) { + // Entitlement discovery intentionally excludes main while a temporary drain + // fences its credential. Once every eligible non-main candidate is unavailable, + // return main only as a non-mutating sentinel so the caller's atomic claim can + // classify maintenance. Do not fall through to the configured-but-ineligible + // active account or persist/bind this synthetic selection. + return { status: "selected", accountId: MAIN_CODEX_ACCOUNT_ID }; } else if ( hasConfiguredPoolAccount(config, active, selectionOptions) && !isCodexAccountPaused(config, active) @@ -1589,11 +1800,30 @@ export function resolveCodexAccountForThreadDetailed( // stays the operator's selection and getEffectiveActiveCodexAccountId is what // surfaces this to the API and dashboard. An independent quota group must not // move the shared cursor at all — its ordering decision is its own. - if (!isIndependentCodexQuotaScope(quotaScope)) rememberActiveCodexAccount(config, preempted); + if ( + !preserveSharedSelectionForModelDetour + && !isIndependentCodexQuotaScope(quotaScope) + ) { + rememberActiveCodexAccount(config, preempted); + } active = preempted; } - active = applyQuotaAutoSwitch(config, active, now, quotaScope, selectionOptions); - active = applyFailureFailover(config, active, now, quotaScope, selectionOptions); + active = applyQuotaAutoSwitch( + config, + active, + now, + quotaScope, + selectionOptions, + !preserveSharedSelectionForModelDetour, + ); + active = applyFailureFailover( + config, + active, + now, + quotaScope, + selectionOptions, + !preserveSharedSelectionForModelDetour, + ); if (!isCodexAccountUsable(config, active, selectionOptions)) { return hasConfiguredPoolAccount(config, active, selectionOptions) ? { status: "selected", accountId: active } @@ -1605,7 +1835,9 @@ export function resolveCodexAccountForThreadDetailed( ? { status: "selected", accountId: active } : { status: "none" }; } - if (threadId) bindThreadAffinity(threadId, active, now, quotaScope); + if (threadId && !preserveExistingModelScopedAffinity) { + bindThreadAffinity(threadId, active, now, quotaScope); + } return { status: "selected", accountId: active }; } diff --git a/src/codex/subagent-model-fallback.ts b/src/codex/subagent-model-fallback.ts index 1f56be2634..5b838f9493 100644 --- a/src/codex/subagent-model-fallback.ts +++ b/src/codex/subagent-model-fallback.ts @@ -39,6 +39,7 @@ import { routeModel, type RouteResult } from "../router"; import { sweepExpiredOnWrite } from "../lib/state-store-sweeper"; import { codexAccountNamespaceForModel } from "./account-namespace-match"; import { ACCOUNT_GATED_NATIVE_OPENAI_MODELS } from "./catalog/native-models"; +import { MAIN_CODEX_ACCOUNT_ID } from "./main-account"; import { getUpstreamHostHealth, normalizeUpstreamHostCircuitThreshold, @@ -276,31 +277,68 @@ export function isSubagentModelUnavailable( poolAccountPreview, candidateAccountUsabilityOptions?.modelEligibleAccountIds, ); - if (isModelHealthBlocked(model, config, resolvedAccountId, now)) return true; - if (!isPoolCodexRoute(route)) return false; - - // Pool candidates need a usable account. Derive requirement from the resolved - // route (canonical openai defaults to pool even when codexAccountMode is omitted). - if (!resolvedAccountId) return true; - if (isCodexAccountPaused(config, resolvedAccountId)) return true; - if (!isCodexAccountUsable(config, resolvedAccountId, candidateAccountUsabilityOptions)) return true; - if (route.codexAccountId !== undefined) { - // An account-qualified route is pinned and cannot consume Pool's recovery-probe - // escape hatch. Honor both account-wide and model-scoped cooldowns so fallback - // advances instead of selecting a candidate that exact auth will reject. - const quotaScope = codexQuotaScopeForModel(route.modelId); - if (getCodexQuotaHealthSnapshot(resolvedAccountId, quotaScope, now) !== null) return true; - } else { - const quotaScope = codexQuotaScopeForModel(route.modelId); - const cooldown = getCodexQuotaHealthSnapshot(resolvedAccountId, quotaScope, now); - if (cooldown !== null) { - const probeAvailable = cooldown.quotaScope - ? canAcquireCodexQuotaScopeProbeLease(resolvedAccountId, cooldown.quotaScope, now) - : canAcquireCodexQuotaProbeLease(resolvedAccountId, now); - if (!probeAvailable) return true; + const accountUnavailable = ( + candidateAccountId: string | null, + usabilityOptions: CodexAccountUsabilityOptions | undefined, + includeQuotaExhaustion: boolean, + ): boolean => { + if (isModelHealthBlocked(model, config, candidateAccountId, now)) return true; + if (!isPoolCodexRoute(route)) return false; + + // Pool candidates need a usable account. Derive requirement from the resolved + // route (canonical openai defaults to pool even when codexAccountMode is omitted). + if (!candidateAccountId) return true; + if (isCodexAccountPaused(config, candidateAccountId)) return true; + if (!isCodexAccountUsable(config, candidateAccountId, usabilityOptions)) return true; + if (route.codexAccountId !== undefined) { + // An account-qualified route is pinned and cannot consume Pool's recovery-probe + // escape hatch. Honor both account-wide and model-scoped cooldowns so fallback + // advances instead of selecting a candidate that exact auth will reject. + const quotaScope = codexQuotaScopeForModel(route.modelId); + if (getCodexQuotaHealthSnapshot(candidateAccountId, quotaScope, now) !== null) return true; + } else { + const quotaScope = codexQuotaScopeForModel(route.modelId); + const cooldown = getCodexQuotaHealthSnapshot(candidateAccountId, quotaScope, now); + if (cooldown !== null) { + const probeAvailable = cooldown.quotaScope + ? canAcquireCodexQuotaScopeProbeLease(candidateAccountId, cooldown.quotaScope, now) + : canAcquireCodexQuotaProbeLease(candidateAccountId, now); + if (!probeAvailable) return true; + } } - } - return isNativeModelQuotaExhausted(model, config, resolvedAccountId, now); + if ( + !includeQuotaExhaustion + || ( + candidateAccountId === MAIN_CODEX_ACCOUNT_ID + && usabilityOptions?.nativeMainSelectionOnly === true + ) + ) return false; + return isNativeModelQuotaExhausted(model, config, candidateAccountId, now); + }; + + // Prefer a genuinely usable entitled pool account. Preview can deliberately return + // the configured active account even when no selectable candidate exists, so a + // null/main-only check is not enough to detect the temporary-drain case. + if (!accountUnavailable(resolvedAccountId, candidateAccountUsabilityOptions, true)) return false; + + // During a temporary native-main drain, entitlement discovery excludes main to + // preserve the credential fence. If no non-main candidate can serve an unqualified + // gated model, retain main only as a read-free sentinel: final auth owns the atomic + // claim and returns maintenance instead of letting a routed fallback bypass it. + const preserveDrainingMainCandidate = route.codexAccountId === undefined + && candidateAccountUsabilityOptions?.nativeMainSelectionOnly === true + && ACCOUNT_GATED_NATIVE_OPENAI_MODELS.has(route.modelId); + if (!preserveDrainingMainCandidate) return true; + const drainingMainUsabilityOptions: CodexAccountUsabilityOptions = { + ...candidateAccountUsabilityOptions, + modelEligibleAccountIds: new Set([ + ...(modelEligibleAccountIds ?? []), + MAIN_CODEX_ACCOUNT_ID, + ]), + }; + // Quota scoring main would lazily read the native credential/plan. Cached health, + // pause, reauth, and cooldown state are safe; defer physical scoring to final auth. + return accountUnavailable(MAIN_CODEX_ACCOUNT_ID, drainingMainUsabilityOptions, false); } export function selectAvailableSubagentModel( diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index c018751cfd..ff6a2dea50 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -150,6 +150,7 @@ import { type CodexAuthContext, } from "../../codex/auth-context"; import { + CodexModelEntitlementDiscoveryUnavailableError, entitledCodexAccountIdsForModel, invalidateCodexModelEntitlementsForAccount, resolveCodexModelEntitlements, @@ -893,6 +894,7 @@ interface CodexPoolAccountRetryArgs { codexWsRuntimeIdentity?: BunRuntimeGateInput; translatorBudget: TranslatorBudget; turnAdmissionLease?: AdmissionLease; + resolveCodexModelEntitlements?: typeof resolveCodexModelEntitlements; }; firstAuthCtx: Extract; firstResponse: Response; @@ -917,6 +919,7 @@ type CodexPoolAccountRetryResult = selectedForwardHeaders: Headers; } | { kind: "no-alternate" } + | { kind: "eligibility-unavailable" } | { kind: "transport"; error: unknown; @@ -1011,10 +1014,21 @@ async function retryCodexPoolOnAlternateAccount( outcomeStatus, upstream, connectMs, passthroughEstimate, stream, } = args; const inboundWire = options.inboundWire ?? "responses"; + const entitlementResolver = options.resolveCodexModelEntitlements ?? resolveCodexModelEntitlements; let retryAuthCtx: CodexAuthContext | undefined; if (outcomeStatus === 400 && ACCOUNT_GATED_NATIVE_OPENAI_MODELS.has(route.modelId)) { invalidateCodexModelEntitlementsForAccount(firstAuthCtx.accountId); - const refreshed = await resolveCodexModelEntitlements(config); + let refreshed; + try { + refreshed = await resolveCodexModelEntitlementsForRequest(entitlementResolver, config); + } catch (error) { + await firstResponse.body?.cancel().catch(() => undefined); + releaseCodexAuthContextProbeLease(firstAuthCtx); + if (error instanceof CodexModelEligibilityUnavailableError) { + return { kind: "eligibility-unavailable" }; + } + throw error; + } if (entitledCodexAccountIdsForModel(refreshed, route.modelId)?.has(firstAuthCtx.accountId)) { // The authenticated roster still grants this exact model. Retry on the same account: // upstream shards can briefly disagree during a gated-model rollout, but a pre-stream 400 @@ -1034,15 +1048,28 @@ async function retryCodexPoolOnAlternateAccount( excludeAccountId: firstAuthCtx.accountId, modelId: route.modelId, beginCodexAccountSelection: codexAccountSelectionForTurn(options.turnAdmissionLease), + resolveCodexModelEntitlements: (entitlementConfig, resolveOptions) => + resolveCodexModelEntitlementsForRequest( + entitlementResolver, + entitlementConfig, + resolveOptions, + ), }, ); } catch (error) { - if ( + const unexpectedRetryError = !(error instanceof CodexPoolAuthenticationError) && !(error instanceof CodexAuthContextError) && !(error instanceof CodexAccountCooldownError) - && !(error instanceof CodexMainProfileDrainingError) - ) throw error; + && !(error instanceof CodexMainProfileDrainingError); + if (unexpectedRetryError) { + await firstResponse.body?.cancel().catch(() => undefined); + releaseCodexAuthContextProbeLease(firstAuthCtx); + if (error instanceof CodexModelEligibilityUnavailableError) { + return { kind: "eligibility-unavailable" }; + } + throw error; + } } if (retryAuthCtx?.kind !== "pool" && retryAuthCtx?.kind !== "main-pool") { return { kind: "no-alternate" }; @@ -1131,24 +1158,30 @@ async function retryCodexPoolOnAlternateAccount( try { while (true) { noteAttemptSend(logCtx.activeAttempt, passthroughEstimate); - upstreamResponse = await fetchWithHeaderTimeout( - request.url, - { - method: request.method, - headers: request.headers, - body: request.body, - }, - upstream.signal, - connectMs, - stream, - providerFetch(route.provider, options.codexWsRuntimeIdentity, { - providerName: route.providerName, - modelId: route.modelId, - }), - // Credential-bearing forward send: never follow a redirect into a - // dead-host rejection after the credential was seen (#914). - route.provider.authMode === "forward", - ); + try { + upstreamResponse = await fetchWithHeaderTimeout( + request.url, + { + method: request.method, + headers: request.headers, + body: request.body, + }, + upstream.signal, + connectMs, + stream, + providerFetch(route.provider, options.codexWsRuntimeIdentity, { + providerName: route.providerName, + modelId: route.modelId, + }), + // Credential-bearing forward send: never follow a redirect into a + // dead-host rejection after the credential was seen (#914). + route.provider.authMode === "forward", + ); + } catch (error) { + // Only the forward send is a transport boundary. Entitlement resolver throws below are + // deliberately outside this catch so programming errors retain their original path. + return { kind: "transport", error, authCtx: retryAuthCtx }; + } retrySendCount += 1; args.onResponse?.(upstreamResponse, retryAuthCtx, request); if (!retrySameConfirmedAccount || retrySendCount >= maxRetrySends) break; @@ -1158,13 +1191,22 @@ async function retryCodexPoolOnAlternateAccount( options.abortSignal, )) break; invalidateCodexModelEntitlementsForAccount(retryAuthCtx.accountId); - const refreshed = await resolveCodexModelEntitlements(config); + let refreshed: Awaited>; + try { + refreshed = await resolveCodexModelEntitlementsForRequest(entitlementResolver, config); + } catch (error) { + await upstreamResponse.body?.cancel().catch(() => undefined); + await firstResponse.body?.cancel().catch(() => undefined); + releaseCodexAuthContextProbeLease(firstAuthCtx); + releaseCodexAuthContextProbeLease(retryAuthCtx); + if (error instanceof CodexModelEligibilityUnavailableError) { + return { kind: "eligibility-unavailable" }; + } + throw error; + } if (!entitledCodexAccountIdsForModel(refreshed, route.modelId)?.has(retryAuthCtx.accountId)) break; await upstreamResponse.body?.cancel().catch(() => undefined); } - } catch (error) { - // Attribute the transport failure to the alternate account (already selected). - return { kind: "transport", error, authCtx: retryAuthCtx }; } finally { request.releaseBodyObservation?.(); } @@ -1602,7 +1644,12 @@ async function resolveResponsesCodexAuth( modelId: route.modelId, substituteMainCredentialForDirect: substituteMainCredential, beginCodexAccountSelection: codexAccountSelectionForTurn(options.turnAdmissionLease), - resolveCodexModelEntitlements: options.resolveCodexModelEntitlements, + resolveCodexModelEntitlements: (entitlementConfig, resolveOptions) => + resolveCodexModelEntitlementsForRequest( + options.resolveCodexModelEntitlements ?? resolveCodexModelEntitlements, + entitlementConfig, + resolveOptions, + ), }); options.onCodexAuthContextResolved?.(authCtx); } else { @@ -1632,6 +1679,9 @@ async function resolveResponsesCodexAuth( if (err instanceof ForwardAdmissionCredentialError) { return { ok: false, response: formatErrorResponse(401, "authentication_error", err.message) }; } + if (err instanceof CodexModelEligibilityUnavailableError) { + return { ok: false, response: codexModelEligibilityUnavailableResponse() }; + } const response = mapCodexAuthContextErrorToResponse(err, { accountSelector: route.codexAccountNamespace, now: Date.now(), @@ -1641,6 +1691,52 @@ async function resolveResponsesCodexAuth( } } +const CODEX_MODEL_ELIGIBILITY_UNAVAILABLE_MESSAGE = + "Codex model eligibility is temporarily unavailable; retry this request"; + +class CodexModelEligibilityUnavailableError extends Error { + constructor() { + super(CODEX_MODEL_ELIGIBILITY_UNAVAILABLE_MESSAGE); + this.name = "CodexModelEligibilityUnavailableError"; + } +} + +/** Return a retryable, redacted failure without letting discovery errors escape the request boundary. */ +function codexModelEligibilityUnavailableResponse(): Response { + return formatErrorResponse(503, "server_error", CODEX_MODEL_ELIGIBILITY_UNAVAILABLE_MESSAGE); +} + +/** Wrap account-roster discovery so both preview and final auth share one fail-closed error contract. */ +async function resolveCodexModelEntitlementsForRequest( + resolver: typeof resolveCodexModelEntitlements, + config: Parameters[0], + options?: Parameters[1], +): ReturnType { + try { + return await resolver(config, options); + } catch (cause) { + if (!(cause instanceof CodexModelEntitlementDiscoveryUnavailableError)) throw cause; + const diagnosticCause = cause.cause ?? cause; + let detail = "unknown error"; + try { + const rawDetail = diagnosticCause instanceof Error + ? `${diagnosticCause.name}: ${diagnosticCause.message}` + : String(diagnosticCause); + detail = sanitizeLogMetadataString(rawDetail, 300) ?? detail; + } catch { + // A hostile thrown value must not replace the fixed retryable response. + } + try { + console.warn( + `[codex-entitlements] model eligibility discovery failed; returning a retryable 503: ${detail}`, + ); + } catch { + // Logging is diagnostic only; the request boundary remains fail-closed below. + } + throw new CodexModelEligibilityUnavailableError(); + } +} + async function resolveSubagentFallbackModelEligibility(args: { config: OcxConfig; fallbackChain: readonly string[] | null; @@ -1651,7 +1747,11 @@ async function resolveSubagentFallbackModelEligibility(args: { const excludeAccountIds = args.nativeMainReadsForbidden ? new Set([MAIN_CODEX_ACCOUNT_ID]) : undefined; - const snapshot = await args.resolver(args.config, { excludeAccountIds }); + const snapshot = await resolveCodexModelEntitlementsForRequest( + args.resolver, + args.config, + { excludeAccountIds }, + ); return (modelId) => { const entitledAccountIds = entitledCodexAccountIdsForModel(snapshot, modelId); return entitledAccountIds @@ -2517,12 +2617,19 @@ async function handleResponsesInner( // "legacy" affinity bucket and never find a binding made under "shared" or a native // model scope, making the preview diverge from the account that actually authenticates. const fallbackChain = initialSubagentFallbackChain; - subagentFallbackModelEligibleAccountIdsForModel = await resolveSubagentFallbackModelEligibility({ - config, - fallbackChain, - nativeMainReadsForbidden, - resolver: options.resolveCodexModelEntitlements ?? resolveCodexModelEntitlements, - }); + try { + subagentFallbackModelEligibleAccountIdsForModel = await resolveSubagentFallbackModelEligibility({ + config, + fallbackChain, + nativeMainReadsForbidden, + resolver: options.resolveCodexModelEntitlements ?? resolveCodexModelEntitlements, + }); + } catch (error) { + if (error instanceof CodexModelEligibilityUnavailableError) { + return codexModelEligibilityUnavailableResponse(); + } + throw error; + } const fallbackNow = Date.now(); subagentFallbackAccountPreview = (modelId, previewNow, modelEligibleAccountIds) => previewCodexAccountForRequest( poolAffinityKey, @@ -3738,6 +3845,9 @@ async function handleResponsesInner( captureAffinityResponse(response, retryAuthCtx, retryRequest, true); }, }); + if (retry.kind === "eligibility-unavailable") { + return codexModelEligibilityUnavailableResponse(); + } if (retry.kind === "transport") { authCtx = retry.authCtx; return transportFailureResponse(retry.error); diff --git a/tests/codex-auth-context.test.ts b/tests/codex-auth-context.test.ts index da4d6cb5d0..d2b4c7d999 100644 --- a/tests/codex-auth-context.test.ts +++ b/tests/codex-auth-context.test.ts @@ -36,7 +36,11 @@ import { saveCodexAccountCredential, } from "../src/codex/account-store"; import { ConfigMutationLockError, getConfigPath } from "../src/config"; -import { MAIN_CODEX_ACCOUNT_ID } from "../src/codex/main-account"; +import { + getMainAccountPlan, + MAIN_CODEX_ACCOUNT_ID, + setMainAccountPlan, +} from "../src/codex/main-account"; import { clearAccountNeedsReauth, clearAccountQuota, @@ -86,6 +90,7 @@ beforeEach(() => { clearThreadAccountMap(); clearCodexUpstreamHealth(); clearAccountQuota(); + setMainAccountPlan(null); __resetGuardianState(); clearAccountNeedsReauth("pool-a"); clearAccountNeedsReauth("pool-b"); @@ -97,6 +102,7 @@ afterEach(() => { clearThreadAccountMap(); clearCodexUpstreamHealth(); clearAccountQuota(); + setMainAccountPlan(null); __resetGuardianState(); clearAccountNeedsReauth("pool-a"); clearAccountNeedsReauth("pool-b"); @@ -122,6 +128,12 @@ function config(): OcxConfig { }; } +function chatgptPlanJwt(plan: string): string { + const header = Buffer.from(JSON.stringify({ alg: "none" })).toString("base64url"); + const body = Buffer.from(JSON.stringify({ chatgpt_plan_type: plan })).toString("base64url"); + return `${header}.${body}.sig`; +} + function guardianConfig(): OcxConfig { const cfg = config(); cfg.defaultProvider = "openai"; @@ -325,12 +337,92 @@ describe("Codex auth context", () => { primeCodexPoolQuotas: async () => { throw new Error("must not prime"); }, })).rejects.toBeInstanceOf(CodexMainProfileDrainingError); expect(nativeReads).toBe(0); + // Selection-only quota scoring must not lazily read/cache the fenced main + // plan before the atomic claim rejects the request. + writeFileSync(join(testDir, "auth.json"), JSON.stringify({ + tokens: { access_token: chatgptPlanJwt("pro"), account_id: "main-account" }, + })); + expect(getMainAccountPlan()).toBe("pro"); + } finally { + turn?.release(); + drain?.release(); + } + }); + + test("gated main selection preserves the temporary drain classification without entitlement reads", async () => { + const cfg = config(); + cfg.activeCodexAccountId = MAIN_CODEX_ACCOUNT_ID; + cfg.activeCodexAccountPinned = MAIN_CODEX_ACCOUNT_ID; + let nativeReads = 0; + let entitlementCalls = 0; + const drain = acquireNativeMainProfileDrain("auth-context-gated-main-test"); + const turn = tryAdmitTurn(); + try { + await expect(resolveCodexAuthContext(new Headers(), cfg, "pool", { + modelId: "gpt-daybreak-blue-latest", + beginCodexAccountSelection: codexAccountSelectionForTurn(turn!), + isMainAccountTokenLive: () => { nativeReads += 1; return true; }, + getMainAccountToken: () => { + nativeReads += 1; + return { accessToken: "main", chatgptAccountId: "main-account" }; + }, + resolveCodexModelEntitlements: async (_config, options) => { + entitlementCalls += 1; + expect(options?.excludeAccountIds?.has(MAIN_CODEX_ACCOUNT_ID)).toBe(true); + return { + modelsByAccount: new Map(), + confirmedAccountIds: new Set(), + credentialIdentities: new Map(), + }; + }, + })).rejects.toBeInstanceOf(CodexMainProfileDrainingError); + expect(entitlementCalls).toBe(1); + expect(nativeReads).toBe(0); + expect(cfg.activeCodexAccountPinned).toBe(MAIN_CODEX_ACCOUNT_ID); + // Pin retirement must not lazily score/read main while the drain fence is up. + // A prior plan read against the intentionally missing auth.json would cache + // `undefined` and make this post-fence JWT fallback fail. + writeFileSync(join(testDir, "auth.json"), JSON.stringify({ + tokens: { access_token: chatgptPlanJwt("pro"), account_id: "main-account" }, + })); + expect(getMainAccountPlan()).toBe("pro"); } finally { turn?.release(); drain?.release(); } }); + test("gated no-active drain reaches the atomic main claim before maintenance classification", async () => { + const cfg = config(); + cfg.activeCodexAccountId = undefined; + let nativeReads = 0; + let claimCalls = 0; + let selectionReleases = 0; + + await expect(resolveCodexAuthContext(new Headers(), cfg, "pool", { + modelId: "gpt-daybreak-blue-latest", + beginCodexAccountSelection: () => ({ + mainProfileDraining: true, + claimMainProfile: () => { claimCalls += 1; return false; }, + release: () => { selectionReleases += 1; }, + }), + isMainAccountTokenLive: () => { nativeReads += 1; return true; }, + getMainAccountToken: () => { + nativeReads += 1; + return { accessToken: "main", chatgptAccountId: "main-account" }; + }, + resolveCodexModelEntitlements: async () => ({ + modelsByAccount: new Map(), + confirmedAccountIds: new Set(), + credentialIdentities: new Map(), + }), + })).rejects.toBeInstanceOf(CodexMainProfileDrainingError); + + expect(claimCalls).toBe(1); + expect(selectionReleases).toBe(1); + expect(nativeReads).toBe(0); + }); + test("direct mode returns caller-owned main context without touching pool selection", async () => { const cfg = { ...config(), activeCodexAccountId: "missing-pool-account" }; await expect(resolveCodexAuthContext(new Headers({ authorization: "Bearer caller" }), cfg, "direct")) @@ -390,6 +482,7 @@ describe("Codex auth context", () => { test("account-gated native routing skips an active account without the model grant", async () => { const cfg = config(); + cfg.activeCodexAccountPinned = "pool-a"; writeFileSync(join(testDir, "auth.json"), JSON.stringify({ tokens: { access_token: "main-token", account_id: "main-account" }, })); @@ -418,6 +511,20 @@ describe("Codex auth context", () => { kind: "main-pool", accountId: MAIN_CODEX_ACCOUNT_ID, }); + expect(cfg.activeCodexAccountId).toBe("pool-a"); + expect(cfg.activeCodexAccountPinned).toBe("pool-a"); + + // The preceding model-only detour must not replace the operator's shared + // selection; a following ordinary request still uses the pinned pool account. + await expect(resolveCodexAuthContext(new Headers(), cfg, "pool", { + modelId: "gpt-5.5", + isMainAccountTokenLive: () => true, + getMainAccountToken: () => ({ accessToken: "main-token", chatgptAccountId: "main-account" }), + primeCodexPoolQuotas: async () => {}, + })).resolves.toMatchObject({ + kind: "pool", + accountId: "pool-a", + }); }); test("exact account-gated routing fails closed for an unentitled account", async () => { diff --git a/tests/codex-routing.test.ts b/tests/codex-routing.test.ts index 4bb678a95f..f5bfcc31d9 100644 --- a/tests/codex-routing.test.ts +++ b/tests/codex-routing.test.ts @@ -1801,6 +1801,411 @@ describe("codex account selection order", () => { expect(getEffectiveActiveCodexAccountId(config)).toBe("a"); }); + test("model eligibility stays request-scoped and preserves shared selection plus affinity", () => { + const config = orderedConfig({ activeCodexAccountPinned: "b" }); + const now = Date.now(); + updateAccountQuota("a", 10); + updateAccountQuota("b", 10); + + expect(resolveCodexAccountForThread("model-gated-task", config, now, "shared")).toBe("b"); + expect(resolveCodexAccountForThreadDetailed( + "model-gated-task", + config, + now + 1, + "shared", + { modelEligibleAccountIds: new Set(["a"]) }, + )).toEqual({ status: "selected", accountId: "a" }); + + expect(config.activeCodexAccountId).toBe("b"); + expect(config.activeCodexAccountPinned).toBe("b"); + expect(getEffectiveActiveCodexAccountId(config)).toBe("b"); + expect(resolveCodexAccountForThread("model-gated-task", config, now + 2, "shared")).toBe("b"); + }); + + test("a gated first request binds its actual account without replacing global active", () => { + const config = makeConfig({ activeCodexAccountId: "b" }); + const now = Date.now(); + updateAccountQuota("a", 10); + updateAccountQuota("b", 10); + + expect(resolveCodexAccountForThreadDetailed( + "gated-first-task", + config, + now, + "shared", + { modelEligibleAccountIds: new Set(["a"]) }, + )).toEqual({ status: "selected", accountId: "a" }); + expect(config.activeCodexAccountId).toBe("b"); + expect(getEffectiveActiveCodexAccountId(config)).toBe("b"); + expect(resolveCodexAccountForThread("gated-first-task", config, now + 1, "shared")).toBe("a"); + }); + + test("model-scoped round-robin advances without replacing shared selection", () => { + const config = makeConfig({ + accountPoolStrategy: "round-robin", + accountPoolStickyLimit: 1, + activeCodexAccountId: "b", + }); + const selectionOptions = { modelEligibleAccountIds: new Set(["a", "b"]) }; + + expect(resolveCodexAccountForThreadDetailed( + null, + config, + Date.now(), + "shared", + selectionOptions, + )).toEqual({ status: "selected", accountId: "a" }); + expect(resolveCodexAccountForThreadDetailed( + null, + config, + Date.now() + 1, + "shared", + selectionOptions, + )).toEqual({ status: "selected", accountId: "b" }); + expect(config.activeCodexAccountId).toBe("b"); + expect(getEffectiveActiveCodexAccountId(config)).toBe("b"); + }); + + test.each(["fill-first", "round-robin"] as const)( + "%s preserves healthy shared active, pin, and affinity during a model-only detour", + (strategy) => { + const now = 1_800_000_000_000; + const threadId = `healthy-model-detour-${strategy}`; + const config = makeConfig({ + accountPoolStrategy: strategy, + accountPoolStickyLimit: 1, + activeCodexAccountId: "b", + activeCodexAccountPinned: "b", + }); + updateAccountQuota("a", 10); + updateAccountQuota("b", 10); + resetCodexRoutingForManualSelection("b"); + expect(resolveCodexAccountForThread(threadId, config, now, "shared")).toBe("b"); + + expect(resolveCodexAccountForThreadDetailed( + threadId, + config, + now + 1, + "shared", + { modelEligibleAccountIds: new Set(["a"]) }, + )).toEqual({ status: "selected", accountId: "a" }); + + expect(config.activeCodexAccountId).toBe("b"); + expect(config.activeCodexAccountPinned).toBe("b"); + expect(getEffectiveActiveCodexAccountId(config)).toBe("b"); + expect(resolveCodexAccountForThread(threadId, config, now + 2, "shared")).toBe("b"); + }, + ); + + test.each(["fill-first", "round-robin"] as const)( + "%s skips a failover-ready detour candidate while preserving healthy shared state", + (strategy) => { + const now = 1_800_000_000_000; + const config = makeConfig({ + accountPoolStrategy: strategy, + accountPoolStickyLimit: 1, + activeCodexAccountId: "c", + activeCodexAccountPinned: "c", + autoSwitchThreshold: 0, + codexAccounts: [ + { id: "a", email: "a@test", isMain: false }, + { id: "b", email: "b@test", isMain: false }, + { id: "c", email: "c@test", isMain: false }, + ], + }); + saveTestCredential("c"); + resetCodexRoutingForManualSelection("c"); + for (let attempt = 0; attempt < 3; attempt += 1) { + recordCodexUpstreamOutcome(config, "a", 503, { + fixedAccount: true, + now: now + attempt + 1, + }); + } + + expect(resolveCodexAccountForThreadDetailed( + null, + config, + now + CODEX_TRANSIENT_SOFT_AVOID_MS + 4, + "shared", + { modelEligibleAccountIds: new Set(["a", "b"]) }, + )).toEqual({ status: "selected", accountId: "b" }); + expect(config.activeCodexAccountId).toBe("c"); + expect(config.activeCodexAccountPinned).toBe("c"); + expect(getEffectiveActiveCodexAccountId(config)).toBe("c"); + }, + ); + + test.each(["fill-first", "round-robin"] as const)( + "%s retires shared state when model ineligibility overlaps quota exhaustion", + (strategy) => { + const now = 1_800_000_000_000; + const threadId = `quota-model-overlap-${strategy}`; + const config = makeConfig({ + accountPoolStrategy: strategy, + accountPoolStickyLimit: 1, + activeCodexAccountId: "b", + activeCodexAccountPinned: "b", + }); + updateAccountQuota("a", 10); + updateAccountQuota("b", 10); + resetCodexRoutingForManualSelection("b"); + expect(resolveCodexAccountForThread(threadId, config, now, "shared")).toBe("b"); + updateAccountQuota("b", 90); + + expect(resolveCodexAccountForThreadDetailed( + threadId, + config, + now + 1, + "shared", + { modelEligibleAccountIds: new Set(["a"]) }, + )).toEqual({ status: "selected", accountId: "a" }); + expect(config.activeCodexAccountId).toBe("b"); + expect(config.activeCodexAccountPinned).toBeUndefined(); + expect(getEffectiveActiveCodexAccountId(config)).toBe("a"); + + updateAccountQuota("b", 10); + expect(resolveCodexAccountForThread(threadId, config, now + 2, "shared")).toBe("a"); + }, + ); + + test.each(["fill-first", "round-robin"] as const)( + "%s retires shared state when model ineligibility overlaps failover", + (strategy) => { + const now = 1_800_000_000_000; + const threadId = `failure-model-overlap-${strategy}`; + const config = makeConfig({ + accountPoolStrategy: strategy, + accountPoolStickyLimit: 1, + activeCodexAccountId: "b", + activeCodexAccountPinned: "b", + autoSwitchThreshold: 0, + }); + resetCodexRoutingForManualSelection("b"); + expect(resolveCodexAccountForThread(threadId, config, now, "shared")).toBe("b"); + for (let attempt = 0; attempt < 3; attempt += 1) { + recordCodexUpstreamOutcome(config, "b", 503, { + fixedAccount: true, + now: now + attempt + 1, + }); + } + const resolveAt = now + CODEX_TRANSIENT_SOFT_AVOID_MS + 4; + + expect(resolveCodexAccountForThreadDetailed( + threadId, + config, + resolveAt, + "shared", + { modelEligibleAccountIds: new Set(["a"]) }, + )).toEqual({ status: "selected", accountId: "a" }); + expect(config.activeCodexAccountId).toBe("b"); + expect(config.activeCodexAccountPinned).toBeUndefined(); + expect(getEffectiveActiveCodexAccountId(config)).toBe("a"); + + clearCodexUpstreamHealthForAccount("b"); + expect(resolveCodexAccountForThread(threadId, config, resolveAt + 1, "shared")).toBe("a"); + }, + ); + + test.each(["fill-first", "round-robin"] as const)( + "%s cannot re-pick a quota-drained shared account that remains model-eligible", + (strategy) => { + const now = 1_800_000_000_000; + const config = makeConfig({ + accountPoolStrategy: strategy, + accountPoolStickyLimit: 1, + activeCodexAccountId: "b", + activeCodexAccountPinned: "b", + }); + updateAccountQuota("a", 10); + updateAccountQuota("b", 90); + resetCodexRoutingForManualSelection("b"); + + expect(resolveCodexAccountForThreadDetailed( + null, + config, + now, + "shared", + { modelEligibleAccountIds: new Set(["a", "b"]) }, + )).toEqual({ status: "selected", accountId: "a" }); + expect(config.activeCodexAccountId).toBe("b"); + expect(config.activeCodexAccountPinned).toBeUndefined(); + expect(getEffectiveActiveCodexAccountId(config)).toBe("a"); + }, + ); + + test.each(["fill-first", "round-robin"] as const)( + "%s cannot re-pick a failover-ready shared account that remains model-eligible", + (strategy) => { + const now = 1_800_000_000_000; + const config = makeConfig({ + accountPoolStrategy: strategy, + accountPoolStickyLimit: 1, + activeCodexAccountId: "b", + activeCodexAccountPinned: "b", + autoSwitchThreshold: 0, + }); + resetCodexRoutingForManualSelection("b"); + for (let attempt = 0; attempt < 3; attempt += 1) { + recordCodexUpstreamOutcome(config, "b", 503, { + fixedAccount: true, + now: now + attempt + 1, + }); + } + const resolveAt = now + CODEX_TRANSIENT_SOFT_AVOID_MS + 4; + + expect(resolveCodexAccountForThreadDetailed( + null, + config, + resolveAt, + "shared", + { modelEligibleAccountIds: new Set(["a", "b"]) }, + )).toEqual({ status: "selected", accountId: "a" }); + expect(config.activeCodexAccountId).toBe("b"); + expect(config.activeCodexAccountPinned).toBeUndefined(); + expect(getEffectiveActiveCodexAccountId(config)).toBe("a"); + }, + ); + + test("temporary main drain preserves unread health but still retires a known paused pin", () => { + const config = makeConfig({ + activeCodexAccountId: MAIN_CODEX_ACCOUNT_ID, + activeCodexAccountPinned: MAIN_CODEX_ACCOUNT_ID, + pausedCodexAccountIds: [MAIN_CODEX_ACCOUNT_ID], + }); + + expect(resolveCodexAccountForThreadDetailed( + null, + config, + Date.now(), + "shared", + { + nativeMainSelectionOnly: true, + modelEligibleAccountIds: new Set(), + }, + )).toEqual({ status: "selected", accountId: MAIN_CODEX_ACCOUNT_ID }); + expect(config.activeCodexAccountPinned).toBeUndefined(); + }); + + test("model-only detour failure does not retire the healthy operator pin", () => { + const now = 1_800_000_000_000; + const config = makeConfig({ + activeCodexAccountId: "b", + activeCodexAccountPinned: "b", + autoSwitchThreshold: 0, + codexAccounts: [ + { id: "a", email: "a@test", isMain: false }, + { id: "b", email: "b@test", isMain: false }, + { id: "c", email: "c@test", isMain: false }, + ], + }); + saveTestCredential("c"); + for (let attempt = 0; attempt < 3; attempt += 1) { + recordCodexUpstreamOutcome(config, "a", 503, { + fixedAccount: true, + now: now + attempt, + }); + } + + expect(resolveCodexAccountForThreadDetailed( + null, + config, + now + CODEX_TRANSIENT_SOFT_AVOID_MS + 3, + "shared", + { modelEligibleAccountIds: new Set(["a", "c"]) }, + )).toEqual({ status: "selected", accountId: "c" }); + expect(config.activeCodexAccountId).toBe("b"); + expect(config.activeCodexAccountPinned).toBe("b"); + expect(getEffectiveActiveCodexAccountId(config)).toBe("b"); + }); + + test("genuine quota transition still retires an exhausted pin during model-scoped selection", () => { + const config = makeConfig({ + activeCodexAccountId: "b", + activeCodexAccountPinned: "b", + }); + updateAccountQuota("a", 10); + updateAccountQuota("b", 90); + + expect(resolveCodexAccountForThreadDetailed( + null, + config, + Date.now(), + "shared", + { modelEligibleAccountIds: new Set(["a", "b"]) }, + )).toEqual({ status: "selected", accountId: "a" }); + expect(config.activeCodexAccountId).toBe("a"); + expect(config.activeCodexAccountPinned).toBeUndefined(); + }); + + test("genuine failure transition still retires a failing pin during model-scoped selection", () => { + const now = 1_800_000_000_000; + const config = makeConfig({ + activeCodexAccountId: "b", + activeCodexAccountPinned: "b", + autoSwitchThreshold: 0, + }); + for (let attempt = 0; attempt < 3; attempt += 1) { + recordCodexUpstreamOutcome(config, "b", 503, { + fixedAccount: true, + now: now + attempt, + }); + } + + expect(resolveCodexAccountForThreadDetailed( + null, + config, + now + CODEX_TRANSIENT_SOFT_AVOID_MS + 3, + "shared", + { modelEligibleAccountIds: new Set(["a", "b"]) }, + )).toEqual({ status: "selected", accountId: "a" }); + expect(config.activeCodexAccountId).toBe("a"); + expect(config.activeCodexAccountPinned).toBeUndefined(); + }); + + test("model ineligibility does not preserve a simultaneously exhausted pin", () => { + const config = makeConfig({ + activeCodexAccountId: "b", + activeCodexAccountPinned: "b", + }); + updateAccountQuota("a", 10); + updateAccountQuota("b", 90); + + expect(resolveCodexAccountForThreadDetailed( + null, + config, + Date.now(), + "shared", + { modelEligibleAccountIds: new Set(["a"]) }, + )).toEqual({ status: "selected", accountId: "a" }); + expect(config.activeCodexAccountId).toBe("a"); + expect(config.activeCodexAccountPinned).toBeUndefined(); + }); + + test("model ineligibility does not preserve a simultaneously failing pin", () => { + const now = 1_800_000_000_000; + const config = makeConfig({ + activeCodexAccountId: "b", + activeCodexAccountPinned: "b", + autoSwitchThreshold: 0, + }); + for (let attempt = 0; attempt < 3; attempt += 1) { + recordCodexUpstreamOutcome(config, "b", 503, { + fixedAccount: true, + now: now + attempt, + }); + } + + expect(resolveCodexAccountForThreadDetailed( + null, + config, + now + CODEX_TRANSIENT_SOFT_AVOID_MS + 3, + "shared", + { modelEligibleAccountIds: new Set(["a"]) }, + )).toEqual({ status: "selected", accountId: "a" }); + expect(config.activeCodexAccountId).toBe("a"); + expect(config.activeCodexAccountPinned).toBeUndefined(); + }); + test("falls through to the lower tier once the higher one is over threshold", () => { const config = orderedConfig(); updateAccountQuota("a", 90); diff --git a/tests/subagent-fallback-handle-responses.test.ts b/tests/subagent-fallback-handle-responses.test.ts index 384368eae7..fc1e4c20f6 100644 --- a/tests/subagent-fallback-handle-responses.test.ts +++ b/tests/subagent-fallback-handle-responses.test.ts @@ -4,10 +4,11 @@ * native effort clamp on final route, pool account preview for native fallback, * encrypted native-only fallback, native passthrough terminal finalization. */ -import { afterEach, beforeEach, describe, expect, setDefaultTimeout, test } from "bun:test"; -import { mkdtempSync, rmSync } from "node:fs"; +import { afterEach, beforeEach, describe, expect, setDefaultTimeout, spyOn, test } from "bun:test"; +import { mkdtempSync, rmSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; +import { fileURLToPath } from "node:url"; import { saveCodexAccountCredential } from "../src/codex/account-store"; import { clearAccountQuota, @@ -29,8 +30,10 @@ import { setSubagentQuotaPrimeForTests, } from "../src/codex/subagent-model-fallback"; import { + CodexModelEntitlementDiscoveryUnavailableError, resetCodexModelEntitlementCacheForTests, } from "../src/codex/model-entitlements"; +import { getMainAccountPlan, setMainAccountPlan } from "../src/codex/main-account"; import { resolveCodexAuthContext, type CodexAuthContext } from "../src/codex/auth-context"; import { handleResponses } from "../src/server/responses"; import { resetAgentTaskRecoveryState } from "../src/server/responses/agent-task-recovery"; @@ -64,6 +67,7 @@ beforeEach(() => { clearAccountQuota(); resetAgentTaskRecoveryState(); resetSubagentModelFallbackStateForTests(); + setMainAccountPlan(null); // Gated-native negative rosters are cached process-wide for 15s; a real-network // miss in one test must not fail-closed the next test's entitlement lookups. resetCodexModelEntitlementCacheForTests(); @@ -77,6 +81,7 @@ afterEach(() => { clearAccountQuota(); resetAgentTaskRecoveryState(); resetSubagentModelFallbackStateForTests(); + setMainAccountPlan(null); rmSync(testDir, { recursive: true, force: true }); if (previousOpencodexHome === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = previousOpencodexHome; @@ -95,6 +100,12 @@ function fernetFixture(ciphertextBytes = 16): string { const FERNET_TASK = fernetFixture(); const GPT56_NATIVE_MODELS = ["gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna"] as const; +function chatgptPlanJwt(plan: string): string { + const header = Buffer.from(JSON.stringify({ alg: "none" })).toString("base64url"); + const body = Buffer.from(JSON.stringify({ chatgpt_plan_type: plan })).toString("base64url"); + return `${header}.${body}.sig`; +} + function encryptedAgentInput(): unknown[] { return [{ type: "agent_message", @@ -251,6 +262,45 @@ async function postSpawn( ); } +async function postDirectCodex( + config: OcxConfig, + body: Record, + options: Parameters[3] = {}, +): Promise { + return handleResponses( + new Request("http://localhost/v1/responses", { + method: "POST", + headers: { + "content-type": "application/json", + authorization: "Bearer caller-codex-token", + }, + body: JSON.stringify(body), + }), + config, + { model: "", provider: "" }, + options, + ); +} + +function unsupportedCodexModelResponse(model: string): Response { + return new Response(JSON.stringify({ + detail: `The '${model}' model is not supported when using Codex with a ChatGPT account.`, + }), { + status: 400, + headers: { "content-type": "application/json" }, + }); +} + +function entitlementSnapshot(grants: Readonly>) { + return { + modelsByAccount: new Map( + Object.entries(grants).map(([accountId, models]) => [accountId, new Set(models)]), + ), + confirmedAccountIds: new Set(Object.keys(grants)), + credentialIdentities: new Map(), + }; +} + describe("subagent fallback without primary auth cooldown failure", () => { test("exact account child bypasses quota priming and fallback on an empty 503", async () => { const now = 1_800_000_000_000; @@ -802,6 +852,8 @@ describe("native fallback account preview", () => { let resolverCalls = 0; let rejectDiscovery!: (reason: Error) => void; const discovery = new Promise((_resolve, reject) => { rejectDiscovery = reject; }); + let signalResolverEntered!: () => void; + const resolverEntered = new Promise((resolve) => { signalResolverEntered = resolve; }); const turnAdmissionLease = { release() {}, beginCodexAccountSelection() { @@ -814,32 +866,131 @@ describe("native fallback account preview", () => { }, } satisfies Pick; let fetchCalls = 0; + const warning = spyOn(console, "warn").mockImplementation(() => {}); + try { + globalThis.fetch = (async () => { + fetchCalls += 1; + throw new Error("must not dispatch"); + }) as typeof fetch; + + const pending = postSpawn( + cfg, + { model: "team/gpt-5.6-sol", input: readableAgentInput(), stream: false }, + { + turnAdmissionLease, + resolveCodexModelEntitlements: async () => { + resolverCalls += 1; + signalResolverEntered(); + return discovery; + }, + }, + ); + await resolverEntered; + + expect(resolverCalls).toBe(1); + expect(beginCount).toBe(1); + expect(releaseCount).toBe(0); + expect(fetchCalls).toBe(0); + + rejectDiscovery(new CodexModelEntitlementDiscoveryUnavailableError(new Error( + "entitlement discovery unavailable sk-secret123456\nforged-record\u2028next", + ))); + const response = await pending; + expect(response.status).toBe(503); + const responseText = await response.text(); + expect(responseText).toContain("Codex model eligibility is temporarily unavailable"); + expect(responseText).not.toContain("entitlement discovery unavailable"); + const warningText = warning.mock.calls.flat().join(" "); + expect(warningText).toContain("model eligibility discovery failed"); + expect(warningText).toContain("[REDACTED]"); + expect(warningText).not.toContain("sk-secret123456"); + expect(warningText).not.toMatch(/[\r\n\u2028\u2029]/); + expect(releaseCount).toBe(1); + expect(fetchCalls).toBe(0); + } finally { + warning.mockRestore(); + } + }); + + test("final auth maps a later entitlement discovery failure to a closed 503 response", async () => { + const now = 1_800_000_000_000; + Date.now = () => now; + installPoolCredential("pool-a", "pool_acc_a", now); + const cfg = poolNativePlusRoutedConfig({ + activeCodexAccountId: "pool-a", + autoSwitchThreshold: 0, + codexAccountNamespaces: { team: "pool-a" }, + subagentModelFallback: ["gpt-daybreak-blue-latest"], + }); + const entitlementSnapshot = { + modelsByAccount: new Map([ + ["pool-a", new Set(["gpt-5.6-sol", "gpt-daybreak-blue-latest"])], + ]), + confirmedAccountIds: new Set(["pool-a"]), + credentialIdentities: new Map(), + }; + let entitlementCalls = 0; + let fetchCalls = 0; + const warning = spyOn(console, "warn").mockImplementation(() => { + throw new Error("logger unavailable"); + }); + try { + globalThis.fetch = (async () => { + fetchCalls += 1; + throw new Error("must not dispatch"); + }) as typeof fetch; + + const response = await postSpawn( + cfg, + { model: "team/gpt-5.6-sol", input: readableAgentInput(), stream: false }, + { + resolveCodexModelEntitlements: async () => { + entitlementCalls += 1; + if (entitlementCalls === 1) return entitlementSnapshot; + throw new CodexModelEntitlementDiscoveryUnavailableError( + new Error("later entitlement discovery unavailable"), + ); + }, + }, + ); + + expect(response.status).toBe(503); + const responseText = await response.text(); + expect(responseText).toContain("Codex model eligibility is temporarily unavailable"); + expect(responseText).not.toContain("later entitlement discovery unavailable"); + expect(entitlementCalls).toBe(2); + expect(fetchCalls).toBe(0); + expect(warning).toHaveBeenCalled(); + } finally { + warning.mockRestore(); + } + }); + + test("programmer errors from entitlement discovery are not mislabeled as retryable 503s", async () => { + const now = 1_800_000_000_000; + Date.now = () => now; + installPoolCredential("pool-a", "pool_acc_a", now); + const cfg = poolNativePlusRoutedConfig({ + activeCodexAccountId: "pool-a", + autoSwitchThreshold: 0, + codexAccountNamespaces: { team: "pool-a" }, + subagentModelFallback: ["gpt-daybreak-blue-latest"], + }); + let fetchCalls = 0; globalThis.fetch = (async () => { fetchCalls += 1; throw new Error("must not dispatch"); }) as typeof fetch; - const pending = postSpawn( + await expect(postSpawn( cfg, { model: "team/gpt-5.6-sol", input: readableAgentInput(), stream: false }, { - turnAdmissionLease, resolveCodexModelEntitlements: async () => { - resolverCalls += 1; - return discovery; + throw new TypeError("programmer sentinel"); }, }, - ); - for (let i = 0; i < 20 && resolverCalls === 0; i += 1) await Promise.resolve(); - - expect(resolverCalls).toBe(1); - expect(beginCount).toBe(1); - expect(releaseCount).toBe(0); - expect(fetchCalls).toBe(0); - - rejectDiscovery(new Error("entitlement discovery unavailable")); - await expect(pending).rejects.toThrow("entitlement discovery unavailable"); - expect(releaseCount).toBe(1); + )).rejects.toThrow("programmer sentinel"); expect(fetchCalls).toBe(0); }); @@ -917,12 +1068,13 @@ describe("native fallback account preview", () => { }; const mainExclusions: boolean[] = []; let selectionReleases = 0; + let claimCalls = 0; const turnAdmissionLease = { release() {}, beginCodexAccountSelection() { return { mainProfileDraining: true, - claimMainProfile: () => false, + claimMainProfile: () => { claimCalls += 1; return false; }, release: () => { selectionReleases += 1; }, }; }, @@ -947,10 +1099,123 @@ describe("native fallback account preview", () => { expect(response.status).toBe(200); expect(mainExclusions).toEqual([true, true]); expect(selectionReleases).toBe(2); + expect(claimCalls).toBe(0); expect(finalAuth).toMatchObject({ kind: "pool", accountId: "pool-b" }); expect(capture.auths[0]).toContain("pool-b_token"); }); + test("temporary drain keeps an unread main-only gated candidate ahead of routed fallback", async () => { + const now = 1_800_000_000_000; + Date.now = () => now; + installPoolCredential("pool-a", "pool_acc_a", now); + const cfg = poolNativePlusRoutedConfig({ + activeCodexAccountId: "pool-a", + autoSwitchThreshold: 0, + codexAccountNamespaces: { team: "pool-a" }, + subagentModelFallback: ["gpt-daybreak-blue-latest", "xai/grok-4.5"], + }); + recordCodexUpstreamOutcome(cfg, "pool-a", 429, { + fixedAccount: true, + modelId: "gpt-5.6-sol", + now, + resetAt: Math.floor((now + 60 * 60_000) / 1_000), + }); + const entitlementSnapshot = { + modelsByAccount: new Map>(), + confirmedAccountIds: new Set(), + credentialIdentities: new Map(), + }; + const mainExclusions: boolean[] = []; + let selectionReleases = 0; + let claimCalls = 0; + const turnAdmissionLease = { + release() {}, + beginCodexAccountSelection() { + return { + mainProfileDraining: true, + claimMainProfile: () => { claimCalls += 1; return false; }, + release: () => { selectionReleases += 1; }, + }; + }, + } satisfies Pick; + let fetchCalls = 0; + globalThis.fetch = (async () => { + fetchCalls += 1; + throw new Error("must not dispatch"); + }) as typeof fetch; + + const response = await postSpawn( + cfg, + { model: "team/gpt-5.6-sol", input: readableAgentInput(), stream: false }, + { + turnAdmissionLease, + resolveCodexModelEntitlements: async (_config, resolveOptions) => { + mainExclusions.push(resolveOptions?.excludeAccountIds?.has("__main__") === true); + return entitlementSnapshot; + }, + }, + ); + + expect(response.status).toBe(503); + expect(await response.text()).toContain("OpenCodex local native-main profile maintenance is active"); + expect(mainExclusions).toEqual([true, true]); + expect(selectionReleases).toBe(2); + expect(claimCalls).toBe(1); + expect(fetchCalls).toBe(0); + // If preview or pin retirement tried to score synthetic main, getMainAccountPlan + // would have consumed the missing auth.json attempt and cached `undefined`. + writeFileSync(join(testDir, "auth.json"), JSON.stringify({ + tokens: { access_token: chatgptPlanJwt("pro"), account_id: "main-account" }, + })); + expect(getMainAccountPlan()).toBe("pro"); + }); + + test("temporary drain keeps ordinary native-main fallback read-free until the final claim", async () => { + const now = 1_800_000_000_000; + Date.now = () => now; + const cfg = poolNativePlusRoutedConfig({ + defaultProvider: "xai", + activeCodexAccountId: "__main__", + subagentModelFallback: ["gpt-5.5"], + }); + const { noteSubagentModelFailure } = await import("../src/codex/subagent-model-fallback"); + noteSubagentModelFailure("xai/grok-4.5", "429", cfg); + noteSubagentModelFailure("grok-4.5", "429", cfg); + let selectionReleases = 0; + let claimCalls = 0; + const turnAdmissionLease = { + release() {}, + beginCodexAccountSelection() { + return { + mainProfileDraining: true, + claimMainProfile: () => { claimCalls += 1; return false; }, + release: () => { selectionReleases += 1; }, + }; + }, + } satisfies Pick; + let fetchCalls = 0; + globalThis.fetch = (async () => { + fetchCalls += 1; + throw new Error("must not dispatch"); + }) as typeof fetch; + + const response = await postSpawn( + cfg, + { model: "xai/grok-4.5", input: readableAgentInput(), stream: false }, + { turnAdmissionLease }, + ); + + expect(response.status).toBe(503); + expect(await response.text()).toContain("OpenCodex local native-main profile maintenance is active"); + expect(selectionReleases).toBe(2); + expect(claimCalls).toBe(1); + expect(fetchCalls).toBe(0); + writeFileSync(join(testDir, "auth.json"), JSON.stringify({ + tokens: { access_token: chatgptPlanJwt("pro"), account_id: "main-account" }, + })); + expect(getMainAccountPlan()).toBe("pro"); + }); + test("Desktop fallback affinity drives the subagent preview and final native account", async () => { const now = 1_800_000_000_000; Date.now = () => now; @@ -1291,7 +1556,7 @@ describe("native fallback account preview", () => { */ test("both fallback preview sites pass the model-eligible account set (#2509)", async () => { const source = await Bun.file( - new URL("../src/server/responses/core.ts", import.meta.url).pathname, + fileURLToPath(new URL("../src/server/responses/core.ts", import.meta.url)), ).text(); const previews = source.match(/subagentFallbackAccountPreview = \([^)]*\)/g) ?? []; @@ -1457,6 +1722,305 @@ describe("native fallback account preview", () => { }); }); +describe("account-gated retry entitlement boundary", () => { + const model = "gpt-daybreak-blue-latest"; + + function retryConfig(secondAccount = false): OcxConfig { + // Keep account selection local to this boundary test. Without known quota, auth performs a + // WHAM prime whose fetch is unrelated to the credential-bearing send count asserted below. + updateAccountQuota("pool-a", 10, undefined, 20); + if (secondAccount) updateAccountQuota("pool-b", 10, undefined, 20); + return poolNativePlusRoutedConfig({ + activeCodexAccountId: "pool-a", + autoSwitchThreshold: 0, + codexAccounts: [ + { id: "main", email: "main@example.test", isMain: true }, + { id: "pool-a", email: "a@example.test", isMain: false, chatgptAccountId: "pool_acc_a" }, + ...(secondAccount + ? [{ id: "pool-b", email: "b@example.test", isMain: false, chatgptAccountId: "pool_acc_b" }] + : []), + ], + }); + } + + test("an explicit discovery outage during the first 400 refresh returns a fixed 503", async () => { + const now = 1_800_000_000_000; + Date.now = () => now; + installPoolCredential("pool-a", "pool_acc_a", now); + const cfg = retryConfig(); + let entitlementCalls = 0; + let fetchCalls = 0; + const upstreamResponses: Response[] = []; + globalThis.fetch = (async () => { + fetchCalls += 1; + const response = unsupportedCodexModelResponse(model); + upstreamResponses.push(response); + return response; + }) as typeof fetch; + const warning = spyOn(console, "warn").mockImplementation(() => {}); + try { + const response = await postDirectCodex( + cfg, + { model, input: "hello", stream: false }, + { + resolveCodexModelEntitlements: async () => { + entitlementCalls += 1; + if (entitlementCalls === 1) return entitlementSnapshot({ "pool-a": [model] }); + throw new CodexModelEntitlementDiscoveryUnavailableError( + new Error("first-refresh-secret"), + ); + }, + }, + ); + + expect(response.status).toBe(503); + const text = await response.text(); + expect(text).toContain("Codex model eligibility is temporarily unavailable"); + expect(text).not.toContain("first-refresh-secret"); + expect(entitlementCalls).toBe(2); + expect(fetchCalls).toBe(1); + expect(upstreamResponses.every(response => response.bodyUsed)).toBe(true); + } finally { + warning.mockRestore(); + } + }); + + test("a first-refresh programmer error cancels the 400 and releases its quota probe", async () => { + const cooldownAt = 1_800_000_000_000; + const probeAt = cooldownAt + CODEX_QUOTA_PROBE_INTERVAL_MS; + Date.now = () => probeAt; + installPoolCredential("pool-a", "pool_acc_a", probeAt); + const cfg = retryConfig(); + recordCodexUpstreamOutcome(cfg, "pool-a", 429, { + fixedAccount: true, + modelId: model, + now: cooldownAt, + resetAt: Math.floor((cooldownAt + 4 * 24 * 60 * 60_000) / 1_000), + }); + let entitlementCalls = 0; + let firstAuth: CodexAuthContext | undefined; + const upstreamResponses: Response[] = []; + globalThis.fetch = (async () => { + const response = unsupportedCodexModelResponse(model); + upstreamResponses.push(response); + return response; + }) as typeof fetch; + + await expect(postDirectCodex( + cfg, + { model, input: "hello", stream: false }, + { + onCodexAuthContextResolved: (ctx) => { firstAuth ??= ctx; }, + resolveCodexModelEntitlements: async () => { + entitlementCalls += 1; + if (entitlementCalls === 1) return entitlementSnapshot({ "pool-a": [model] }); + throw new TypeError("first-refresh programmer sentinel"); + }, + }, + )).rejects.toThrow("first-refresh programmer sentinel"); + + const firstProbeLeaseId = (firstAuth as { probeLeaseId?: string } | undefined)?.probeLeaseId; + expect(firstProbeLeaseId).toBeTruthy(); + expect(upstreamResponses).toHaveLength(1); + expect(upstreamResponses[0]?.bodyUsed).toBe(true); + + Date.now = () => probeAt + CODEX_QUOTA_PROBE_INTERVAL_MS; + const nextProbe = await resolveCodexAuthContext(new Headers(), cfg, "pool", { + modelId: model, + resolveCodexModelEntitlements: async () => entitlementSnapshot({ "pool-a": [model] }), + }); + expect((nextProbe as { probeLeaseId?: string }).probeLeaseId).toBeTruthy(); + expect((nextProbe as { probeLeaseId?: string }).probeLeaseId).not.toBe(firstProbeLeaseId); + }); + + test("an explicit discovery outage while selecting an alternate account returns a fixed 503", async () => { + const now = 1_800_000_000_000; + Date.now = () => now; + installPoolCredential("pool-a", "pool_acc_a", now); + installPoolCredential("pool-b", "pool_acc_b", now); + const cfg = retryConfig(true); + let entitlementCalls = 0; + let fetchCalls = 0; + const upstreamResponses: Response[] = []; + globalThis.fetch = (async () => { + fetchCalls += 1; + const response = unsupportedCodexModelResponse(model); + upstreamResponses.push(response); + return response; + }) as typeof fetch; + const warning = spyOn(console, "warn").mockImplementation(() => {}); + try { + const response = await postDirectCodex( + cfg, + { model, input: "hello", stream: false }, + { + resolveCodexModelEntitlements: async () => { + entitlementCalls += 1; + if (entitlementCalls === 1) { + return entitlementSnapshot({ "pool-a": [model], "pool-b": [model] }); + } + if (entitlementCalls === 2) { + return entitlementSnapshot({ + "pool-a": ["gpt-5.6-sol"], + "pool-b": [model], + }); + } + throw new CodexModelEntitlementDiscoveryUnavailableError( + new Error("alternate-refresh-secret"), + ); + }, + }, + ); + + expect(response.status).toBe(503); + const text = await response.text(); + expect(text).toContain("Codex model eligibility is temporarily unavailable"); + expect(text).not.toContain("alternate-refresh-secret"); + expect(entitlementCalls).toBe(3); + expect(fetchCalls).toBe(1); + expect(upstreamResponses.every(response => response.bodyUsed)).toBe(true); + } finally { + warning.mockRestore(); + } + }); + + test("an alternate-selection programmer error cancels the 400 and releases its quota probe", async () => { + const cooldownAt = 1_800_000_000_000; + const probeAt = cooldownAt + CODEX_QUOTA_PROBE_INTERVAL_MS; + Date.now = () => probeAt; + installPoolCredential("pool-a", "pool_acc_a", probeAt); + installPoolCredential("pool-b", "pool_acc_b", probeAt); + const cfg = retryConfig(true); + recordCodexUpstreamOutcome(cfg, "pool-a", 429, { + fixedAccount: true, + modelId: model, + now: cooldownAt, + resetAt: Math.floor((cooldownAt + 4 * 24 * 60 * 60_000) / 1_000), + }); + let entitlementCalls = 0; + let firstAuth: CodexAuthContext | undefined; + const upstreamResponses: Response[] = []; + globalThis.fetch = (async () => { + const response = unsupportedCodexModelResponse(model); + upstreamResponses.push(response); + return response; + }) as typeof fetch; + + await expect(postDirectCodex( + cfg, + { model, input: "hello", stream: false }, + { + onCodexAuthContextResolved: (ctx) => { firstAuth ??= ctx; }, + resolveCodexModelEntitlements: async () => { + entitlementCalls += 1; + if (entitlementCalls === 1) { + return entitlementSnapshot({ + "pool-a": [model], + "pool-b": ["gpt-5.6-sol"], + }); + } + if (entitlementCalls === 2) { + return entitlementSnapshot({ + "pool-a": ["gpt-5.6-sol"], + "pool-b": [model], + }); + } + throw new TypeError("alternate programmer sentinel"); + }, + }, + )).rejects.toThrow("alternate programmer sentinel"); + + const firstProbeLeaseId = (firstAuth as { probeLeaseId?: string } | undefined)?.probeLeaseId; + expect(firstProbeLeaseId).toBeTruthy(); + expect(upstreamResponses).toHaveLength(1); + expect(upstreamResponses[0]?.bodyUsed).toBe(true); + + Date.now = () => probeAt + CODEX_QUOTA_PROBE_INTERVAL_MS; + const nextProbe = await resolveCodexAuthContext(new Headers(), cfg, "pool", { + modelId: model, + resolveCodexModelEntitlements: async () => entitlementSnapshot({ + "pool-a": [model], + "pool-b": ["gpt-5.6-sol"], + }), + }); + expect((nextProbe as { probeLeaseId?: string }).probeLeaseId).toBeTruthy(); + expect((nextProbe as { probeLeaseId?: string }).probeLeaseId).not.toBe(firstProbeLeaseId); + }); + + test("an explicit discovery outage between bounded same-account retries returns a fixed 503", async () => { + const now = 1_800_000_000_000; + Date.now = () => now; + installPoolCredential("pool-a", "pool_acc_a", now); + const cfg = retryConfig(); + let entitlementCalls = 0; + let fetchCalls = 0; + const upstreamResponses: Response[] = []; + globalThis.fetch = (async () => { + fetchCalls += 1; + const response = unsupportedCodexModelResponse(model); + upstreamResponses.push(response); + return response; + }) as typeof fetch; + const warning = spyOn(console, "warn").mockImplementation(() => {}); + try { + const response = await postDirectCodex( + cfg, + { model, input: "hello", stream: false }, + { + resolveCodexModelEntitlements: async () => { + entitlementCalls += 1; + if (entitlementCalls <= 2) return entitlementSnapshot({ "pool-a": [model] }); + throw new CodexModelEntitlementDiscoveryUnavailableError( + new Error("same-account-refresh-secret"), + ); + }, + }, + ); + + expect(response.status).toBe(503); + const text = await response.text(); + expect(text).toContain("Codex model eligibility is temporarily unavailable"); + expect(text).not.toContain("same-account-refresh-secret"); + expect(entitlementCalls).toBe(3); + expect(fetchCalls).toBe(2); + expect(upstreamResponses.every(response => response.bodyUsed)).toBe(true); + } finally { + warning.mockRestore(); + } + }); + + test("a programmer error between same-account retries keeps its original error path", async () => { + const now = 1_800_000_000_000; + Date.now = () => now; + installPoolCredential("pool-a", "pool_acc_a", now); + const cfg = retryConfig(); + let entitlementCalls = 0; + let fetchCalls = 0; + const upstreamResponses: Response[] = []; + globalThis.fetch = (async () => { + fetchCalls += 1; + const response = unsupportedCodexModelResponse(model); + upstreamResponses.push(response); + return response; + }) as typeof fetch; + + await expect(postDirectCodex( + cfg, + { model, input: "hello", stream: false }, + { + resolveCodexModelEntitlements: async () => { + entitlementCalls += 1; + if (entitlementCalls <= 2) return entitlementSnapshot({ "pool-a": [model] }); + throw new TypeError("retry programmer sentinel"); + }, + }, + )).rejects.toThrow("retry programmer sentinel"); + expect(entitlementCalls).toBe(3); + expect(fetchCalls).toBe(2); + expect(upstreamResponses.every(response => response.bodyUsed)).toBe(true); + }); +}); + describe("encrypted child native-only fallback", () => { test("rejects encrypted routed primary when only routed fallbacks exist", async () => { const cfg = poolNativePlusRoutedConfig({ From 8011ff802982cf39cef3b13aab6e277204c10f46 Mon Sep 17 00:00:00 2001 From: luvs01 <27862058+luvs01@users.noreply.github.com> Date: Wed, 26 Aug 2026 19:51:14 +0900 Subject: [PATCH 2/6] fix(codex): defer unimplemented entitlement outage contract --- src/codex/model-entitlements.ts | 12 - src/server/responses/core.ts | 106 +------ ...subagent-fallback-handle-responses.test.ts | 276 ++++-------------- 3 files changed, 73 insertions(+), 321 deletions(-) diff --git a/src/codex/model-entitlements.ts b/src/codex/model-entitlements.ts index 1bc6f43285..5a649d335a 100644 --- a/src/codex/model-entitlements.ts +++ b/src/codex/model-entitlements.ts @@ -46,18 +46,6 @@ export interface CodexModelEntitlementResolveOptions { readonly excludeAccountIds?: ReadonlySet; } -/** - * Explicit request-boundary signal for an operational discovery failure that cannot be - * represented by the ordinary fail-closed `confirmed: false` snapshot. Generic throws are - * programming errors and must retain their original error path. - */ -export class CodexModelEntitlementDiscoveryUnavailableError extends Error { - constructor(cause?: unknown) { - super("Codex model entitlement discovery is temporarily unavailable", { cause }); - this.name = "CodexModelEntitlementDiscoveryUnavailableError"; - } -} - const accountModelsCache = new Map(); const accountModelsFlights = new Map>(); diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index ff6a2dea50..e88c6b4e97 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -150,7 +150,6 @@ import { type CodexAuthContext, } from "../../codex/auth-context"; import { - CodexModelEntitlementDiscoveryUnavailableError, entitledCodexAccountIdsForModel, invalidateCodexModelEntitlementsForAccount, resolveCodexModelEntitlements, @@ -919,7 +918,6 @@ type CodexPoolAccountRetryResult = selectedForwardHeaders: Headers; } | { kind: "no-alternate" } - | { kind: "eligibility-unavailable" } | { kind: "transport"; error: unknown; @@ -1020,13 +1018,10 @@ async function retryCodexPoolOnAlternateAccount( invalidateCodexModelEntitlementsForAccount(firstAuthCtx.accountId); let refreshed; try { - refreshed = await resolveCodexModelEntitlementsForRequest(entitlementResolver, config); + refreshed = await entitlementResolver(config); } catch (error) { await firstResponse.body?.cancel().catch(() => undefined); releaseCodexAuthContextProbeLease(firstAuthCtx); - if (error instanceof CodexModelEligibilityUnavailableError) { - return { kind: "eligibility-unavailable" }; - } throw error; } if (entitledCodexAccountIdsForModel(refreshed, route.modelId)?.has(firstAuthCtx.accountId)) { @@ -1048,12 +1043,7 @@ async function retryCodexPoolOnAlternateAccount( excludeAccountId: firstAuthCtx.accountId, modelId: route.modelId, beginCodexAccountSelection: codexAccountSelectionForTurn(options.turnAdmissionLease), - resolveCodexModelEntitlements: (entitlementConfig, resolveOptions) => - resolveCodexModelEntitlementsForRequest( - entitlementResolver, - entitlementConfig, - resolveOptions, - ), + resolveCodexModelEntitlements: entitlementResolver, }, ); } catch (error) { @@ -1065,9 +1055,6 @@ async function retryCodexPoolOnAlternateAccount( if (unexpectedRetryError) { await firstResponse.body?.cancel().catch(() => undefined); releaseCodexAuthContextProbeLease(firstAuthCtx); - if (error instanceof CodexModelEligibilityUnavailableError) { - return { kind: "eligibility-unavailable" }; - } throw error; } } @@ -1193,15 +1180,12 @@ async function retryCodexPoolOnAlternateAccount( invalidateCodexModelEntitlementsForAccount(retryAuthCtx.accountId); let refreshed: Awaited>; try { - refreshed = await resolveCodexModelEntitlementsForRequest(entitlementResolver, config); + refreshed = await entitlementResolver(config); } catch (error) { await upstreamResponse.body?.cancel().catch(() => undefined); await firstResponse.body?.cancel().catch(() => undefined); releaseCodexAuthContextProbeLease(firstAuthCtx); releaseCodexAuthContextProbeLease(retryAuthCtx); - if (error instanceof CodexModelEligibilityUnavailableError) { - return { kind: "eligibility-unavailable" }; - } throw error; } if (!entitledCodexAccountIdsForModel(refreshed, route.modelId)?.has(retryAuthCtx.accountId)) break; @@ -1644,12 +1628,7 @@ async function resolveResponsesCodexAuth( modelId: route.modelId, substituteMainCredentialForDirect: substituteMainCredential, beginCodexAccountSelection: codexAccountSelectionForTurn(options.turnAdmissionLease), - resolveCodexModelEntitlements: (entitlementConfig, resolveOptions) => - resolveCodexModelEntitlementsForRequest( - options.resolveCodexModelEntitlements ?? resolveCodexModelEntitlements, - entitlementConfig, - resolveOptions, - ), + resolveCodexModelEntitlements: options.resolveCodexModelEntitlements, }); options.onCodexAuthContextResolved?.(authCtx); } else { @@ -1679,9 +1658,6 @@ async function resolveResponsesCodexAuth( if (err instanceof ForwardAdmissionCredentialError) { return { ok: false, response: formatErrorResponse(401, "authentication_error", err.message) }; } - if (err instanceof CodexModelEligibilityUnavailableError) { - return { ok: false, response: codexModelEligibilityUnavailableResponse() }; - } const response = mapCodexAuthContextErrorToResponse(err, { accountSelector: route.codexAccountNamespace, now: Date.now(), @@ -1691,52 +1667,6 @@ async function resolveResponsesCodexAuth( } } -const CODEX_MODEL_ELIGIBILITY_UNAVAILABLE_MESSAGE = - "Codex model eligibility is temporarily unavailable; retry this request"; - -class CodexModelEligibilityUnavailableError extends Error { - constructor() { - super(CODEX_MODEL_ELIGIBILITY_UNAVAILABLE_MESSAGE); - this.name = "CodexModelEligibilityUnavailableError"; - } -} - -/** Return a retryable, redacted failure without letting discovery errors escape the request boundary. */ -function codexModelEligibilityUnavailableResponse(): Response { - return formatErrorResponse(503, "server_error", CODEX_MODEL_ELIGIBILITY_UNAVAILABLE_MESSAGE); -} - -/** Wrap account-roster discovery so both preview and final auth share one fail-closed error contract. */ -async function resolveCodexModelEntitlementsForRequest( - resolver: typeof resolveCodexModelEntitlements, - config: Parameters[0], - options?: Parameters[1], -): ReturnType { - try { - return await resolver(config, options); - } catch (cause) { - if (!(cause instanceof CodexModelEntitlementDiscoveryUnavailableError)) throw cause; - const diagnosticCause = cause.cause ?? cause; - let detail = "unknown error"; - try { - const rawDetail = diagnosticCause instanceof Error - ? `${diagnosticCause.name}: ${diagnosticCause.message}` - : String(diagnosticCause); - detail = sanitizeLogMetadataString(rawDetail, 300) ?? detail; - } catch { - // A hostile thrown value must not replace the fixed retryable response. - } - try { - console.warn( - `[codex-entitlements] model eligibility discovery failed; returning a retryable 503: ${detail}`, - ); - } catch { - // Logging is diagnostic only; the request boundary remains fail-closed below. - } - throw new CodexModelEligibilityUnavailableError(); - } -} - async function resolveSubagentFallbackModelEligibility(args: { config: OcxConfig; fallbackChain: readonly string[] | null; @@ -1747,11 +1677,7 @@ async function resolveSubagentFallbackModelEligibility(args: { const excludeAccountIds = args.nativeMainReadsForbidden ? new Set([MAIN_CODEX_ACCOUNT_ID]) : undefined; - const snapshot = await resolveCodexModelEntitlementsForRequest( - args.resolver, - args.config, - { excludeAccountIds }, - ); + const snapshot = await args.resolver(args.config, { excludeAccountIds }); return (modelId) => { const entitledAccountIds = entitledCodexAccountIdsForModel(snapshot, modelId); return entitledAccountIds @@ -2617,19 +2543,12 @@ async function handleResponsesInner( // "legacy" affinity bucket and never find a binding made under "shared" or a native // model scope, making the preview diverge from the account that actually authenticates. const fallbackChain = initialSubagentFallbackChain; - try { - subagentFallbackModelEligibleAccountIdsForModel = await resolveSubagentFallbackModelEligibility({ - config, - fallbackChain, - nativeMainReadsForbidden, - resolver: options.resolveCodexModelEntitlements ?? resolveCodexModelEntitlements, - }); - } catch (error) { - if (error instanceof CodexModelEligibilityUnavailableError) { - return codexModelEligibilityUnavailableResponse(); - } - throw error; - } + subagentFallbackModelEligibleAccountIdsForModel = await resolveSubagentFallbackModelEligibility({ + config, + fallbackChain, + nativeMainReadsForbidden, + resolver: options.resolveCodexModelEntitlements ?? resolveCodexModelEntitlements, + }); const fallbackNow = Date.now(); subagentFallbackAccountPreview = (modelId, previewNow, modelEligibleAccountIds) => previewCodexAccountForRequest( poolAffinityKey, @@ -3845,9 +3764,6 @@ async function handleResponsesInner( captureAffinityResponse(response, retryAuthCtx, retryRequest, true); }, }); - if (retry.kind === "eligibility-unavailable") { - return codexModelEligibilityUnavailableResponse(); - } if (retry.kind === "transport") { authCtx = retry.authCtx; return transportFailureResponse(retry.error); diff --git a/tests/subagent-fallback-handle-responses.test.ts b/tests/subagent-fallback-handle-responses.test.ts index fc1e4c20f6..416e1eb3c4 100644 --- a/tests/subagent-fallback-handle-responses.test.ts +++ b/tests/subagent-fallback-handle-responses.test.ts @@ -29,10 +29,7 @@ import { resetSubagentModelFallbackStateForTests, setSubagentQuotaPrimeForTests, } from "../src/codex/subagent-model-fallback"; -import { - CodexModelEntitlementDiscoveryUnavailableError, - resetCodexModelEntitlementCacheForTests, -} from "../src/codex/model-entitlements"; +import { resetCodexModelEntitlementCacheForTests } from "../src/codex/model-entitlements"; import { getMainAccountPlan, setMainAccountPlan } from "../src/codex/main-account"; import { resolveCodexAuthContext, type CodexAuthContext } from "../src/codex/auth-context"; import { handleResponses } from "../src/server/responses"; @@ -838,7 +835,7 @@ describe("native fallback account preview", () => { expect(capture.auths[0]).toContain("pool-a_token"); }); - test("entitlement discovery holds and releases preview admission on rejection", async () => { + test("pending preview entitlement errors release admission after preserving the original path", async () => { const now = 1_800_000_000_000; Date.now = () => now; const cfg = poolNativePlusRoutedConfig({ @@ -866,53 +863,37 @@ describe("native fallback account preview", () => { }, } satisfies Pick; let fetchCalls = 0; - const warning = spyOn(console, "warn").mockImplementation(() => {}); - try { - globalThis.fetch = (async () => { - fetchCalls += 1; - throw new Error("must not dispatch"); - }) as typeof fetch; + globalThis.fetch = (async () => { + fetchCalls += 1; + throw new Error("must not dispatch"); + }) as typeof fetch; - const pending = postSpawn( - cfg, - { model: "team/gpt-5.6-sol", input: readableAgentInput(), stream: false }, - { - turnAdmissionLease, - resolveCodexModelEntitlements: async () => { - resolverCalls += 1; - signalResolverEntered(); - return discovery; - }, + const pending = postSpawn( + cfg, + { model: "team/gpt-5.6-sol", input: readableAgentInput(), stream: false }, + { + turnAdmissionLease, + resolveCodexModelEntitlements: async () => { + resolverCalls += 1; + signalResolverEntered(); + return discovery; }, - ); - await resolverEntered; - - expect(resolverCalls).toBe(1); - expect(beginCount).toBe(1); - expect(releaseCount).toBe(0); - expect(fetchCalls).toBe(0); - - rejectDiscovery(new CodexModelEntitlementDiscoveryUnavailableError(new Error( - "entitlement discovery unavailable sk-secret123456\nforged-record\u2028next", - ))); - const response = await pending; - expect(response.status).toBe(503); - const responseText = await response.text(); - expect(responseText).toContain("Codex model eligibility is temporarily unavailable"); - expect(responseText).not.toContain("entitlement discovery unavailable"); - const warningText = warning.mock.calls.flat().join(" "); - expect(warningText).toContain("model eligibility discovery failed"); - expect(warningText).toContain("[REDACTED]"); - expect(warningText).not.toContain("sk-secret123456"); - expect(warningText).not.toMatch(/[\r\n\u2028\u2029]/); - expect(releaseCount).toBe(1); - expect(fetchCalls).toBe(0); - } finally { - warning.mockRestore(); - } + }, + ); + await resolverEntered; + + expect(resolverCalls).toBe(1); + expect(beginCount).toBe(1); + expect(releaseCount).toBe(0); + expect(fetchCalls).toBe(0); + + rejectDiscovery(new TypeError("preview entitlement programmer sentinel")); + await expect(pending).rejects.toThrow("preview entitlement programmer sentinel"); + expect(releaseCount).toBe(1); + expect(fetchCalls).toBe(0); }); - test("final auth maps a later entitlement discovery failure to a closed 503 response", async () => { + test("final-auth entitlement errors release both selection admissions on their original path", async () => { const now = 1_800_000_000_000; Date.now = () => now; installPoolCredential("pool-a", "pool_acc_a", now); @@ -929,44 +910,46 @@ describe("native fallback account preview", () => { confirmedAccountIds: new Set(["pool-a"]), credentialIdentities: new Map(), }; + let beginCount = 0; + let releaseCount = 0; + const turnAdmissionLease = { + release() {}, + beginCodexAccountSelection() { + beginCount += 1; + return { + mainProfileDraining: false, + claimMainProfile: () => true, + release: () => { releaseCount += 1; }, + }; + }, + } satisfies Pick; let entitlementCalls = 0; let fetchCalls = 0; - const warning = spyOn(console, "warn").mockImplementation(() => { - throw new Error("logger unavailable"); - }); - try { - globalThis.fetch = (async () => { - fetchCalls += 1; - throw new Error("must not dispatch"); - }) as typeof fetch; + globalThis.fetch = (async () => { + fetchCalls += 1; + throw new Error("must not dispatch"); + }) as typeof fetch; - const response = await postSpawn( - cfg, - { model: "team/gpt-5.6-sol", input: readableAgentInput(), stream: false }, - { - resolveCodexModelEntitlements: async () => { - entitlementCalls += 1; - if (entitlementCalls === 1) return entitlementSnapshot; - throw new CodexModelEntitlementDiscoveryUnavailableError( - new Error("later entitlement discovery unavailable"), - ); - }, + await expect(postSpawn( + cfg, + { model: "team/gpt-5.6-sol", input: readableAgentInput(), stream: false }, + { + turnAdmissionLease, + resolveCodexModelEntitlements: async () => { + entitlementCalls += 1; + if (entitlementCalls === 1) return entitlementSnapshot; + throw new TypeError("final-auth entitlement programmer sentinel"); }, - ); + }, + )).rejects.toThrow("final-auth entitlement programmer sentinel"); - expect(response.status).toBe(503); - const responseText = await response.text(); - expect(responseText).toContain("Codex model eligibility is temporarily unavailable"); - expect(responseText).not.toContain("later entitlement discovery unavailable"); - expect(entitlementCalls).toBe(2); - expect(fetchCalls).toBe(0); - expect(warning).toHaveBeenCalled(); - } finally { - warning.mockRestore(); - } + expect(entitlementCalls).toBe(2); + expect(beginCount).toBe(2); + expect(releaseCount).toBe(2); + expect(fetchCalls).toBe(0); }); - test("programmer errors from entitlement discovery are not mislabeled as retryable 503s", async () => { + test("programmer errors from entitlement discovery retain their original path", async () => { const now = 1_800_000_000_000; Date.now = () => now; installPoolCredential("pool-a", "pool_acc_a", now); @@ -1743,48 +1726,6 @@ describe("account-gated retry entitlement boundary", () => { }); } - test("an explicit discovery outage during the first 400 refresh returns a fixed 503", async () => { - const now = 1_800_000_000_000; - Date.now = () => now; - installPoolCredential("pool-a", "pool_acc_a", now); - const cfg = retryConfig(); - let entitlementCalls = 0; - let fetchCalls = 0; - const upstreamResponses: Response[] = []; - globalThis.fetch = (async () => { - fetchCalls += 1; - const response = unsupportedCodexModelResponse(model); - upstreamResponses.push(response); - return response; - }) as typeof fetch; - const warning = spyOn(console, "warn").mockImplementation(() => {}); - try { - const response = await postDirectCodex( - cfg, - { model, input: "hello", stream: false }, - { - resolveCodexModelEntitlements: async () => { - entitlementCalls += 1; - if (entitlementCalls === 1) return entitlementSnapshot({ "pool-a": [model] }); - throw new CodexModelEntitlementDiscoveryUnavailableError( - new Error("first-refresh-secret"), - ); - }, - }, - ); - - expect(response.status).toBe(503); - const text = await response.text(); - expect(text).toContain("Codex model eligibility is temporarily unavailable"); - expect(text).not.toContain("first-refresh-secret"); - expect(entitlementCalls).toBe(2); - expect(fetchCalls).toBe(1); - expect(upstreamResponses.every(response => response.bodyUsed)).toBe(true); - } finally { - warning.mockRestore(); - } - }); - test("a first-refresh programmer error cancels the 400 and releases its quota probe", async () => { const cooldownAt = 1_800_000_000_000; const probeAt = cooldownAt + CODEX_QUOTA_PROBE_INTERVAL_MS; @@ -1833,57 +1774,6 @@ describe("account-gated retry entitlement boundary", () => { expect((nextProbe as { probeLeaseId?: string }).probeLeaseId).not.toBe(firstProbeLeaseId); }); - test("an explicit discovery outage while selecting an alternate account returns a fixed 503", async () => { - const now = 1_800_000_000_000; - Date.now = () => now; - installPoolCredential("pool-a", "pool_acc_a", now); - installPoolCredential("pool-b", "pool_acc_b", now); - const cfg = retryConfig(true); - let entitlementCalls = 0; - let fetchCalls = 0; - const upstreamResponses: Response[] = []; - globalThis.fetch = (async () => { - fetchCalls += 1; - const response = unsupportedCodexModelResponse(model); - upstreamResponses.push(response); - return response; - }) as typeof fetch; - const warning = spyOn(console, "warn").mockImplementation(() => {}); - try { - const response = await postDirectCodex( - cfg, - { model, input: "hello", stream: false }, - { - resolveCodexModelEntitlements: async () => { - entitlementCalls += 1; - if (entitlementCalls === 1) { - return entitlementSnapshot({ "pool-a": [model], "pool-b": [model] }); - } - if (entitlementCalls === 2) { - return entitlementSnapshot({ - "pool-a": ["gpt-5.6-sol"], - "pool-b": [model], - }); - } - throw new CodexModelEntitlementDiscoveryUnavailableError( - new Error("alternate-refresh-secret"), - ); - }, - }, - ); - - expect(response.status).toBe(503); - const text = await response.text(); - expect(text).toContain("Codex model eligibility is temporarily unavailable"); - expect(text).not.toContain("alternate-refresh-secret"); - expect(entitlementCalls).toBe(3); - expect(fetchCalls).toBe(1); - expect(upstreamResponses.every(response => response.bodyUsed)).toBe(true); - } finally { - warning.mockRestore(); - } - }); - test("an alternate-selection programmer error cancels the 400 and releases its quota probe", async () => { const cooldownAt = 1_800_000_000_000; const probeAt = cooldownAt + CODEX_QUOTA_PROBE_INTERVAL_MS; @@ -1947,48 +1837,6 @@ describe("account-gated retry entitlement boundary", () => { expect((nextProbe as { probeLeaseId?: string }).probeLeaseId).not.toBe(firstProbeLeaseId); }); - test("an explicit discovery outage between bounded same-account retries returns a fixed 503", async () => { - const now = 1_800_000_000_000; - Date.now = () => now; - installPoolCredential("pool-a", "pool_acc_a", now); - const cfg = retryConfig(); - let entitlementCalls = 0; - let fetchCalls = 0; - const upstreamResponses: Response[] = []; - globalThis.fetch = (async () => { - fetchCalls += 1; - const response = unsupportedCodexModelResponse(model); - upstreamResponses.push(response); - return response; - }) as typeof fetch; - const warning = spyOn(console, "warn").mockImplementation(() => {}); - try { - const response = await postDirectCodex( - cfg, - { model, input: "hello", stream: false }, - { - resolveCodexModelEntitlements: async () => { - entitlementCalls += 1; - if (entitlementCalls <= 2) return entitlementSnapshot({ "pool-a": [model] }); - throw new CodexModelEntitlementDiscoveryUnavailableError( - new Error("same-account-refresh-secret"), - ); - }, - }, - ); - - expect(response.status).toBe(503); - const text = await response.text(); - expect(text).toContain("Codex model eligibility is temporarily unavailable"); - expect(text).not.toContain("same-account-refresh-secret"); - expect(entitlementCalls).toBe(3); - expect(fetchCalls).toBe(2); - expect(upstreamResponses.every(response => response.bodyUsed)).toBe(true); - } finally { - warning.mockRestore(); - } - }); - test("a programmer error between same-account retries keeps its original error path", async () => { const now = 1_800_000_000_000; Date.now = () => now; From b4310bd18612cc6f0c711e45b899e0bd41fc4ccc Mon Sep 17 00:00:00 2001 From: luvs01 <27862058+luvs01@users.noreply.github.com> Date: Fri, 28 Aug 2026 11:16:05 +0900 Subject: [PATCH 3/6] fix(codex): fence retry entitlement refresh --- src/server/responses/core.ts | 35 +++++++++++++- ...subagent-fallback-handle-responses.test.ts | 48 +++++++++++++++++++ 2 files changed, 81 insertions(+), 2 deletions(-) diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index e88c6b4e97..f3d694f77a 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -924,6 +924,29 @@ type CodexPoolAccountRetryResult = authCtx: Extract; }; +/** Keep retry-stage entitlement snapshots inside the native-main selection fence. */ +async function resolveCodexRetryModelEntitlements( + config: OcxConfig, + resolver: typeof resolveCodexModelEntitlements, + turnAdmissionLease?: AdmissionLease, +): Promise>> { + // The initial auth selection has already released its admission before the first + // response arrives. Re-enter for every refresh so profile switching cannot overlap + // credential discovery, and omit main entirely when a drain or recovery owns it. + const selectionAdmission = codexAccountSelectionForTurn(turnAdmissionLease)?.(); + const nativeMainReadsForbidden = isNativeMainTrafficBlocked() + || selectionAdmission?.mainProfileDraining === true; + try { + return await resolver(config, { + excludeAccountIds: nativeMainReadsForbidden + ? new Set([MAIN_CODEX_ACCOUNT_ID]) + : undefined, + }); + } finally { + selectionAdmission?.release(); + } +} + const CODEX_ACCOUNT_GATED_CANONICAL_WIRE_MODELS: ReadonlyMap = new Map([ // The authenticated catalog currently advertises Daybreak Blue, while successful responses // identify the serving model as gpt-5.6-sol. Sending the selector itself is shard-dependent: @@ -1018,7 +1041,11 @@ async function retryCodexPoolOnAlternateAccount( invalidateCodexModelEntitlementsForAccount(firstAuthCtx.accountId); let refreshed; try { - refreshed = await entitlementResolver(config); + refreshed = await resolveCodexRetryModelEntitlements( + config, + entitlementResolver, + options.turnAdmissionLease, + ); } catch (error) { await firstResponse.body?.cancel().catch(() => undefined); releaseCodexAuthContextProbeLease(firstAuthCtx); @@ -1180,7 +1207,11 @@ async function retryCodexPoolOnAlternateAccount( invalidateCodexModelEntitlementsForAccount(retryAuthCtx.accountId); let refreshed: Awaited>; try { - refreshed = await entitlementResolver(config); + refreshed = await resolveCodexRetryModelEntitlements( + config, + entitlementResolver, + options.turnAdmissionLease, + ); } catch (error) { await upstreamResponse.body?.cancel().catch(() => undefined); await firstResponse.body?.cancel().catch(() => undefined); diff --git a/tests/subagent-fallback-handle-responses.test.ts b/tests/subagent-fallback-handle-responses.test.ts index 416e1eb3c4..3997bc8670 100644 --- a/tests/subagent-fallback-handle-responses.test.ts +++ b/tests/subagent-fallback-handle-responses.test.ts @@ -1726,6 +1726,54 @@ describe("account-gated retry entitlement boundary", () => { }); } + test("temporary main drain fences every retry-stage entitlement refresh", async () => { + const now = 1_800_000_000_000; + Date.now = () => now; + installPoolCredential("pool-a", "pool_acc_a", now); + const cfg = retryConfig(); + let selectionReleases = 0; + const turnAdmissionLease = { + release() {}, + beginCodexAccountSelection() { + return { + mainProfileDraining: true, + claimMainProfile: () => false, + release: () => { selectionReleases += 1; }, + }; + }, + } satisfies Pick; + const mainExclusions: boolean[] = []; + let fetchCalls = 0; + globalThis.fetch = (async () => { + fetchCalls += 1; + if (fetchCalls <= 2) return unsupportedCodexModelResponse(model); + return Response.json({ + id: "resp_retry_fenced", + object: "response", + status: "completed", + model, + output: [], + }); + }) as typeof fetch; + + const response = await postDirectCodex( + cfg, + { model, input: "hello", stream: false }, + { + turnAdmissionLease, + resolveCodexModelEntitlements: async (_config, resolveOptions) => { + mainExclusions.push(resolveOptions?.excludeAccountIds?.has("__main__") === true); + return entitlementSnapshot({ "pool-a": [model] }); + }, + }, + ); + + expect(response.status).toBe(200); + expect(fetchCalls).toBe(3); + expect(mainExclusions).toEqual([true, true, true]); + expect(selectionReleases).toBe(3); + }); + test("a first-refresh programmer error cancels the 400 and releases its quota probe", async () => { const cooldownAt = 1_800_000_000_000; const probeAt = cooldownAt + CODEX_QUOTA_PROBE_INTERVAL_MS; From 28a20456222c2c79b98bea5cbe75be699661dcc7 Mon Sep 17 00:00:00 2001 From: luvs01 Date: Fri, 28 Aug 2026 20:44:00 +0900 Subject: [PATCH 4/6] fix(codex): preserve main claims and model detours --- src/codex/auth-context.ts | 64 +++- src/codex/routing.ts | 330 +++++++++++++----- src/server/responses/core.ts | 24 +- .../bearer-admission-routed-provider.test.ts | 118 ++++++- tests/codex-auth-context.test.ts | 170 ++++++++- tests/codex-routing.test.ts | 296 +++++++++++++++- ...subagent-fallback-handle-responses.test.ts | 4 +- 7 files changed, 901 insertions(+), 105 deletions(-) diff --git a/src/codex/auth-context.ts b/src/codex/auth-context.ts index acfb82ccfb..dcf9bb88df 100644 --- a/src/codex/auth-context.ts +++ b/src/codex/auth-context.ts @@ -361,21 +361,50 @@ export async function resolveCodexAuthContext( // selected stored credential even while the canonical OpenAI provider is globally Direct. if (mode === "direct" && fixedAccountId === undefined) { if (!hasCallerCodexBearer(headers)) throw new CodexDirectAuthenticationError(); - if (options.modelId && ACCOUNT_GATED_NATIVE_OPENAI_MODELS.has(options.modelId)) { - const entitled = options.substituteMainCredentialForDirect - ? entitledCodexAccountIdsForModel( - await (options.resolveCodexModelEntitlements ?? resolveCodexModelEntitlements)(config), - options.modelId, - )?.has(MAIN_CODEX_ACCOUNT_ID) === true - : await (options.isDirectCallerEntitledToCodexModel ?? isDirectCallerEntitledToCodexModel)( - headers, - options.modelId, - ); - if (!entitled) { - throw new CodexPoolAuthenticationError("The selected ChatGPT account does not support this model"); + const substituteStoredMain = options.substituteMainCredentialForDirect === true; + if (!substituteStoredMain) { + if (options.modelId && ACCOUNT_GATED_NATIVE_OPENAI_MODELS.has(options.modelId)) { + const entitled = await ( + options.isDirectCallerEntitledToCodexModel ?? isDirectCallerEntitledToCodexModel + )(headers, options.modelId); + if (!entitled) { + throw new CodexPoolAuthenticationError("The selected ChatGPT account does not support this model"); + } } + return { kind: "main", accountId: null }; + } + + // Admission-bearer Direct requests later replace the proxy secret with the stored + // native-main credential. Reserve and claim that physical profile before entitlement + // discovery or materialization can read it; caller-owned Direct credentials never enter + // this branch. A missing turn admission must fail closed instead of recreating an + // untracked native-main read. + if (isNativeMainTrafficBlocked()) throw new CodexMainProfileDrainingError(); + const directSelectionAdmission = options.beginCodexAccountSelection?.(); + if (!directSelectionAdmission) throw new CodexMainProfileDrainingError(); + try { + if ( + directSelectionAdmission.mainProfileDraining + || !directSelectionAdmission.claimMainProfile() + || isNativeMainTrafficBlocked() + ) { + throw new CodexMainProfileDrainingError(); + } + if (options.modelId && ACCOUNT_GATED_NATIVE_OPENAI_MODELS.has(options.modelId)) { + const entitled = entitledCodexAccountIdsForModel( + await (options.resolveCodexModelEntitlements ?? resolveCodexModelEntitlements)(config), + options.modelId, + )?.has(MAIN_CODEX_ACCOUNT_ID) === true; + if (!entitled) { + throw new CodexPoolAuthenticationError("The selected ChatGPT account does not support this model"); + } + } + return { kind: "main", accountId: null }; + } finally { + // The short selector reservation ends here. A successful claim remains owned by + // the enclosing turn lease until the request or transferred stream settles. + directSelectionAdmission.release(); } - return { kind: "main", accountId: null }; } const affinityKey = fixedAccountId === undefined ? codexPoolAffinityKey(headers) : undefined; // Retained startup recovery makes the physical main identity ineligible. Routing @@ -426,7 +455,14 @@ export async function resolveCodexAuthContext( ? { status: "selected" as const, accountId: selected } : { status: "none" as const }; })() - : resolveCodexAccountForThreadDetailed(affinityKey ?? null, config, Date.now(), quotaScope, selectionOptions); + : resolveCodexAccountForThreadDetailed( + affinityKey ?? null, + config, + Date.now(), + quotaScope, + selectionOptions, + options.modelId, + ); if (resolution.status === "expired") throw new CodexThreadAffinityExpiredError(resolution.accountId); const selected = resolution.status === "selected" ? resolution.accountId : null; if (!selected) { diff --git a/src/codex/routing.ts b/src/codex/routing.ts index b9ad137390..100a7e0a9e 100644 --- a/src/codex/routing.ts +++ b/src/codex/routing.ts @@ -158,22 +158,21 @@ export type CodexQuotaRecoveryProbeProof = { * affinity so a Spark failover cannot displace the same thread's Terra/Luna * account (and vice versa). */ -type ThreadAffinityScope = CodexQuotaScope | "legacy"; +type BaseThreadAffinityScope = CodexQuotaScope | "legacy"; +type ModelDetourAffinityScope = `model-detour:${BaseThreadAffinityScope}:${string}`; +type ThreadAffinityScope = BaseThreadAffinityScope | ModelDetourAffinityScope; const LEGACY_THREAD_AFFINITY_SCOPE = "legacy" as const; const threadAccountMap = new Map>(); +let threadAffinityEntryTotal = 0; + +function isModelDetourAffinityScope(scope: ThreadAffinityScope): scope is ModelDetourAffinityScope { + return scope.startsWith("model-detour:"); +} const NATIVE_MODEL_QUOTA_SCOPES: Readonly> = { "gpt-5.3-codex-spark": "spark", }; -// A thread can have one legacy binding plus one binding for each known scope. -// This upper-bound guard avoids an exact map scan until it can be over capacity. -const MAX_THREAD_AFFINITY_SCOPES = new Set([ - LEGACY_THREAD_AFFINITY_SCOPE, - "shared", - ...Object.values(NATIVE_MODEL_QUOTA_SCOPES), -]).size; - export function codexQuotaScopeForModel(modelId: string | undefined): CodexQuotaScope | undefined { if (!modelId?.trim()) return undefined; return NATIVE_MODEL_QUOTA_SCOPES[modelId.trim().toLowerCase()] ?? "shared"; @@ -253,12 +252,15 @@ export function listLiveCodexAccountIds(config: OcxConfig): ReadonlySet export function clearThreadAccountMap(): void { threadAccountMap.clear(); + threadAffinityEntryTotal = 0; } export function clearThreadAccountMapForAccount(accountId: string): void { for (const [threadId, affinities] of threadAccountMap) { for (const [scope, entry] of affinities) { - if (entry.accountId === accountId) affinities.delete(scope); + if (entry.accountId === accountId && affinities.delete(scope)) { + threadAffinityEntryTotal = Math.max(0, threadAffinityEntryTotal - 1); + } } if (affinities.size === 0) threadAccountMap.delete(threadId); } @@ -830,7 +832,7 @@ function isCodexAccountSelectable( && isCodexAccountUsable(config, accountId, selectionOptions); } -function threadAffinityScope(quotaScope?: CodexQuotaScope): ThreadAffinityScope { +function threadAffinityScope(quotaScope?: CodexQuotaScope): BaseThreadAffinityScope { return quotaScope ?? LEGACY_THREAD_AFFINITY_SCOPE; } @@ -838,34 +840,74 @@ function admissibleAffinityComponent(value: string): boolean { return retainedUtf8Bytes(value) <= MAX_AFFINITY_COMPONENT_BYTES; } -function getThreadAffinity(threadId: string, quotaScope?: CodexQuotaScope): ThreadAffinityEntry | undefined { +function modelDetourAffinityScope( + modelId: string | undefined, + quotaScope?: CodexQuotaScope, +): ModelDetourAffinityScope | undefined { + const canonicalModelId = modelId?.trim().toLowerCase(); + if (!canonicalModelId || !admissibleAffinityComponent(canonicalModelId)) return undefined; + return `model-detour:${threadAffinityScope(quotaScope)}:${canonicalModelId}`; +} + +function getThreadAffinityForScope( + threadId: string, + scope: ThreadAffinityScope, +): ThreadAffinityEntry | undefined { if (!admissibleAffinityComponent(threadId)) return undefined; - return threadAccountMap.get(threadId)?.get(threadAffinityScope(quotaScope)); + return threadAccountMap.get(threadId)?.get(scope); } -function deleteThreadAffinity(threadId: string, quotaScope?: CodexQuotaScope): void { +function getThreadAffinity(threadId: string, quotaScope?: CodexQuotaScope): ThreadAffinityEntry | undefined { + return getThreadAffinityForScope(threadId, threadAffinityScope(quotaScope)); +} + +function getModelDetourAffinity( + threadId: string, + modelId: string | undefined, + quotaScope?: CodexQuotaScope, +): ThreadAffinityEntry | undefined { + const scope = modelDetourAffinityScope(modelId, quotaScope); + return scope ? getThreadAffinityForScope(threadId, scope) : undefined; +} + +function deleteThreadAffinityForScope(threadId: string, scope: ThreadAffinityScope): void { if (!admissibleAffinityComponent(threadId)) return; const affinities = threadAccountMap.get(threadId); if (!affinities) return; - affinities.delete(threadAffinityScope(quotaScope)); + if (affinities.delete(scope)) { + threadAffinityEntryTotal = Math.max(0, threadAffinityEntryTotal - 1); + } if (affinities.size === 0) threadAccountMap.delete(threadId); } +function deleteThreadAffinity(threadId: string, quotaScope?: CodexQuotaScope): void { + deleteThreadAffinityForScope(threadId, threadAffinityScope(quotaScope)); +} + +function deleteModelDetourAffinity( + threadId: string, + modelId: string | undefined, + quotaScope?: CodexQuotaScope, +): void { + const scope = modelDetourAffinityScope(modelId, quotaScope); + if (scope) deleteThreadAffinityForScope(threadId, scope); +} + /** Remove only the matching failed account's affinities for one thread. */ function deleteThreadAffinitiesForAccount(threadId: string, accountId: string): void { if (!admissibleAffinityComponent(threadId) || !admissibleAffinityComponent(accountId)) return; const affinities = threadAccountMap.get(threadId); if (!affinities) return; for (const [scope, entry] of affinities) { - if (entry.accountId === accountId) affinities.delete(scope); + if (entry.accountId === accountId && affinities.delete(scope)) { + threadAffinityEntryTotal = Math.max(0, threadAffinityEntryTotal - 1); + } } if (affinities.size === 0) threadAccountMap.delete(threadId); } function threadAffinityEntryCount(): number { - let count = 0; - for (const affinities of threadAccountMap.values()) count += affinities.size; - return count; + return threadAffinityEntryTotal; } function isThreadAffinityExpired(entry: ThreadAffinityEntry, now: number): boolean { @@ -880,43 +922,50 @@ function isThreadAffinityGenerationLive(entry: ThreadAffinityEntry): boolean { function pruneExpiredThreadAffinities(now: number): void { for (const [threadId, affinities] of threadAccountMap) { for (const [scope, entry] of affinities) { - if (isThreadAffinityExpired(entry, now)) affinities.delete(scope); + if (isThreadAffinityExpired(entry, now) && affinities.delete(scope)) { + threadAffinityEntryTotal = Math.max(0, threadAffinityEntryTotal - 1); + } } if (affinities.size === 0) threadAccountMap.delete(threadId); } } function pruneLruThreadAffinities(): void { - if (threadAccountMap.size * MAX_THREAD_AFFINITY_SCOPES <= CODEX_THREAD_AFFINITY_MAX_ENTRIES) return; + if (threadAffinityEntryCount() <= CODEX_THREAD_AFFINITY_MAX_ENTRIES) return; while (threadAffinityEntryCount() > CODEX_THREAD_AFFINITY_MAX_ENTRIES) { let oldestThreadId: string | null = null; let oldestScope: ThreadAffinityScope | null = null; let oldestLastUsedAt = Number.POSITIVE_INFINITY; + let oldestIsDetour = false; for (const [threadId, affinities] of threadAccountMap) { for (const [scope, entry] of affinities) { - if (entry.lastUsedAt < oldestLastUsedAt) { + const candidateIsDetour = isModelDetourAffinityScope(scope); + if ( + (candidateIsDetour && !oldestIsDetour) + || (candidateIsDetour === oldestIsDetour && entry.lastUsedAt < oldestLastUsedAt) + ) { oldestThreadId = threadId; oldestScope = scope; oldestLastUsedAt = entry.lastUsedAt; + oldestIsDetour = candidateIsDetour; } } } if (!oldestThreadId || !oldestScope) return; - deleteThreadAffinity(oldestThreadId, oldestScope === LEGACY_THREAD_AFFINITY_SCOPE ? undefined : oldestScope); + deleteThreadAffinityForScope(oldestThreadId, oldestScope); } } -function bindThreadAffinity( +function bindThreadAffinityForScope( threadId: string, accountId: string, now: number, - quotaScope?: CodexQuotaScope, + scope: ThreadAffinityScope, ): void { if (!admissibleAffinityComponent(threadId) || !admissibleAffinityComponent(accountId)) return; const record = accountId === MAIN_CODEX_ACCOUNT_ID ? undefined : readCodexAccountRecord(accountId); if (accountId !== MAIN_CODEX_ACCOUNT_ID && (!record?.credential || record.deletedAt != null)) return; pruneExpiredThreadAffinities(now); - const scope = threadAffinityScope(quotaScope); const affinities = threadAccountMap.get(threadId) ?? new Map(); const previous = affinities.get(scope); affinities.set(scope, { @@ -926,10 +975,31 @@ function bindThreadAffinity( lastUsedAt: now, lastReevalAt: now, }); + if (!previous) threadAffinityEntryTotal += 1; threadAccountMap.set(threadId, affinities); pruneLruThreadAffinities(); } +function bindThreadAffinity( + threadId: string, + accountId: string, + now: number, + quotaScope?: CodexQuotaScope, +): void { + bindThreadAffinityForScope(threadId, accountId, now, threadAffinityScope(quotaScope)); +} + +function bindModelDetourAffinity( + threadId: string, + accountId: string, + now: number, + modelId: string | undefined, + quotaScope?: CodexQuotaScope, +): void { + const scope = modelDetourAffinityScope(modelId, quotaScope); + if (scope) bindThreadAffinityForScope(threadId, accountId, now, scope); +} + function getEligiblePoolAccounts( config: OcxConfig, excludeId?: string, @@ -1461,6 +1531,30 @@ function isHealthySharedCodexSelection( && !shouldFailover(config, accountId, now); } +function strategySelectionOptionsForModelDetour( + config: OcxConfig, + now: number, + quotaScope: CodexQuotaScope | undefined, + selectionOptions: CodexAccountUsabilityOptions | undefined, +): CodexAccountUsabilityOptions | undefined { + if (selectionOptions?.modelEligibleAccountIds === undefined) return selectionOptions; + const sharedSelectionOptions = sharedStateSelectionOptions(selectionOptions) ?? {}; + return { + ...selectionOptions, + modelEligibleAccountIds: new Set( + [...selectionOptions.modelEligibleAccountIds].filter(accountId => + isHealthySharedCodexSelection( + config, + accountId, + now, + quotaScope, + sharedSelectionOptions, + ) + ), + ), + }; +} + function applyFailureFailover( config: OcxConfig, active: string, @@ -1496,6 +1590,47 @@ export function resolveCodexAccountForThread( return resolution.status === "selected" ? resolution.accountId : null; } +function previewReusableAffinityAccount( + entry: ThreadAffinityEntry | undefined, + config: OcxConfig, + now: number, + quotaScope?: CodexQuotaScope, + selectionOptions?: CodexAccountUsabilityOptions, +): string | null { + if ( + !entry + || isThreadAffinityExpired(entry, now) + || !isThreadAffinityGenerationLive(entry) + || !isCodexAccountSelectable(config, entry.accountId, now, quotaScope, selectionOptions) + || shouldFailover(config, entry.accountId, now) + ) { + return null; + } + // Quota strategy only: non-quota strategies keep affinity for ongoing threads + // (new-session-only rotation — docs / affinity policy A). + if (normalizeAccountPoolStrategy(config.accountPoolStrategy) === "quota") { + const threshold = config.autoSwitchThreshold ?? 80; + if (threshold > 0) { + const usage = computeCodexUsageScore( + getAccountQuota(entry.accountId), + getPoolAccountPlanForSelection(config, entry.accountId, selectionOptions), + ); + if (!isUnknownUsage(usage) && usage >= threshold) { + const best = pickLowerUsageAccount( + config, + entry.accountId, + usage, + now, + quotaScope, + selectionOptions, + ); + if (best !== entry.accountId) return best; + } + } + } + return entry.accountId; +} + /** * Side-effect-free preview of the Codex pool account native routing would prefer. * Used for subagent fallback quota decisions before final auth. @@ -1510,42 +1645,31 @@ export function previewCodexAccountForRequest( now = Date.now(), quotaScope?: CodexQuotaScope, selectionOptions?: CodexAccountUsabilityOptions, + modelId?: string, ): string | null { - const entry = threadId ? getThreadAffinity(threadId, quotaScope) : undefined; - if (threadId && entry) { - if ( - !isThreadAffinityExpired(entry, now) - && isThreadAffinityGenerationLive(entry) - && isCodexAccountSelectable(config, entry.accountId, now, quotaScope, selectionOptions) - && !shouldFailover(config, entry.accountId, now) - ) { - // Quota strategy only: non-quota strategies keep affinity for ongoing threads - // (new-session-only rotation — docs / affinity policy A). - const strategy = normalizeAccountPoolStrategy(config.accountPoolStrategy); - if (strategy === "quota") { - const threshold = config.autoSwitchThreshold ?? 80; - if (threshold > 0) { - const usage = computeCodexUsageScore( - getAccountQuota(entry.accountId), - getPoolAccountPlanForSelection(config, entry.accountId, selectionOptions), - ); - if (!isUnknownUsage(usage) && usage >= threshold) { - const best = pickLowerUsageAccount( - config, - entry.accountId, - usage, - now, - quotaScope, - selectionOptions, - ); - if (best !== entry.accountId) return best; - } - } - } - return entry.accountId; - } - // Stale/unusable affinity is ignored for preview (no map mutation). + // A request-scoped model detour keeps its own serving-account affinity. Preview + // reads it before the ordinary lane, but never repairs or deletes it. Roster + // expansion therefore preserves the already-serving account, and preview mirrors + // final resolution even when the ordinary lane was independently retired. + if (threadId && selectionOptions?.modelEligibleAccountIds !== undefined) { + const detourPreview = previewReusableAffinityAccount( + getModelDetourAffinity(threadId, modelId, quotaScope), + config, + now, + quotaScope, + selectionOptions, + ); + if (detourPreview) return detourPreview; } + const entry = threadId ? getThreadAffinity(threadId, quotaScope) : undefined; + const ordinaryPreview = previewReusableAffinityAccount( + entry, + config, + now, + quotaScope, + selectionOptions, + ); + if (ordinaryPreview) return ordinaryPreview; const strategyPick = pickUnboundStrategyAccount( config, @@ -1553,7 +1677,7 @@ export function previewCodexAccountForRequest( now, false, quotaScope, - selectionOptions, + strategySelectionOptionsForModelDetour(config, now, quotaScope, selectionOptions), ); if (strategyPick) return strategyPick; @@ -1602,6 +1726,7 @@ export function resolveCodexAccountForThreadDetailed( now = Date.now(), quotaScope?: CodexQuotaScope, selectionOptions?: CodexAccountUsabilityOptions, + modelId?: string, ): CodexThreadResolution { // An entitlement roster constrains only this model request. It must not rewrite // the operator's shared active/pin choice or the task's ordinary-model affinity. @@ -1629,6 +1754,56 @@ export function resolveCodexAccountForThreadDetailed( ) ); + if (threadId && modelScopedSelection) { + const detourEntry = getModelDetourAffinity(threadId, modelId, quotaScope); + if (detourEntry) { + const detourReusable = !isThreadAffinityExpired(detourEntry, now) + && isThreadAffinityGenerationLive(detourEntry) + && isCodexAccountSelectable(config, detourEntry.accountId, now, quotaScope, selectionOptions) + && !shouldFailover(config, detourEntry.accountId, now); + if (detourReusable) { + detourEntry.lastUsedAt = now; + // Model detours follow the same affinity policy as ordinary bindings: + // RR/fill-first stay sticky, while quota strategy may re-evaluate an + // over-threshold account without changing the ordinary lane. + if (normalizeAccountPoolStrategy(config.accountPoolStrategy) === "quota") { + const threshold = config.autoSwitchThreshold ?? 80; + const usage = threshold > 0 + ? computeCodexUsageScore( + getAccountQuota(detourEntry.accountId), + getPoolAccountPlanForSelection(config, detourEntry.accountId, selectionOptions), + ) + : 0; + const overThreshold = threshold > 0 && !isUnknownUsage(usage) && usage >= threshold; + if ( + overThreshold + || now - detourEntry.lastReevalAt >= CODEX_THREAD_AFFINITY_REEVAL_INTERVAL_MS + ) { + detourEntry.lastReevalAt = now; + if (overThreshold) { + const best = pickLowerUsageAccount( + config, + detourEntry.accountId, + usage, + now, + quotaScope, + selectionOptions, + ); + if (best !== detourEntry.accountId) { + bindModelDetourAffinity(threadId, best, now, modelId, quotaScope); + return { status: "selected", accountId: best }; + } + } + } + } + return { status: "selected", accountId: detourEntry.accountId }; + } + // Detour expiry or invalidation must not expire the ordinary task. Drop only + // this model lane and select from ordinary/shared state below. + deleteModelDetourAffinity(threadId, modelId, quotaScope); + } + } + const entry = threadId ? getThreadAffinity(threadId, quotaScope) : undefined; if (threadId && entry) { if (isThreadAffinityExpired(entry, now)) { @@ -1699,22 +1874,12 @@ export function resolveCodexAccountForThreadDetailed( // otherwise RR/fill-first can immediately re-pick a known failing account even when another // entitled account is healthy. If no healthy member exists, the normal fallback path below // still decides whether the sole eligible candidate must be used. - const strategySelectionOptions = modelScopedSelection - ? { - ...selectionOptions, - modelEligibleAccountIds: new Set( - [...selectionOptions!.modelEligibleAccountIds!].filter(accountId => - isHealthySharedCodexSelection( - config, - accountId, - now, - quotaScope, - sharedSelectionOptions, - ) - ), - ), - } - : selectionOptions; + const strategySelectionOptions = strategySelectionOptionsForModelDetour( + config, + now, + quotaScope, + selectionOptions, + ); const strategyPick = pickUnboundStrategyAccount( config, threadId, @@ -1726,6 +1891,9 @@ export function resolveCodexAccountForThreadDetailed( !preserveExistingModelScopedAffinity, ); if (strategyPick) { + if (threadId && preserveExistingModelScopedAffinity) { + bindModelDetourAffinity(threadId, strategyPick, now, modelId, quotaScope); + } if ( modelScopedSelection && !preserveSharedSelectionForModelDetour @@ -1835,8 +2003,12 @@ export function resolveCodexAccountForThreadDetailed( ? { status: "selected", accountId: active } : { status: "none" }; } - if (threadId && !preserveExistingModelScopedAffinity) { - bindThreadAffinity(threadId, active, now, quotaScope); + if (threadId) { + if (preserveExistingModelScopedAffinity) { + bindModelDetourAffinity(threadId, active, now, modelId, quotaScope); + } else { + bindThreadAffinity(threadId, active, now, quotaScope); + } } return { status: "selected", accountId: active }; } diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index f3d694f77a..561cd72d70 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -214,7 +214,13 @@ import { import { shouldAttemptImageTierRetry } from "../image-retry"; import { isXaiResponsesDestination, resolveProviderTransport } from "../../providers/xai-transport"; import type { WsData } from "../ws-bridge"; -import { codexAccountSelectionForTurn, registerTurn, trackStreamLifetime, unregisterTurn } from "../lifecycle"; +import { + codexAccountSelectionForTurn, + registerTurn, + trackStreamLifetime, + tryClaimNativeMainProfileForTurn, + unregisterTurn, +} from "../lifecycle"; import { redactSecretString, sanitizeLogMetadataString } from "../../lib/redact"; import { readBoundedResponseBody } from "../../lib/bounded-body"; import type { AdmissionLease } from "../../lib/admission"; @@ -1663,6 +1669,20 @@ async function resolveResponsesCodexAuth( }); options.onCodexAuthContextResolved?.(authCtx); } else { + // A custom-named canonical-forward provider has no Codex account mode, but an + // admission bearer still substitutes the stored main credential below. Claim the + // same physical profile before synthesizing the main context so transport-based + // substitution cannot bypass a switch drain. + if ( + substituteMainCredential + && ( + isNativeMainTrafficBlocked() + || !tryClaimNativeMainProfileForTurn(options.turnAdmissionLease) + || isNativeMainTrafficBlocked() + ) + ) { + throw new CodexMainProfileDrainingError(); + } authCtx = { kind: "main", accountId: null }; options.onCodexAuthContextResolved?.(undefined); } @@ -2587,6 +2607,7 @@ async function handleResponsesInner( previewNow, codexQuotaScopeForModel(modelId), { ...previewSelectionOptions, modelEligibleAccountIds }, + modelId, ); const previewAccountId = route.codexAccountId ?? subagentFallbackAccountPreview( route.modelId, @@ -2708,6 +2729,7 @@ async function handleResponsesInner( previewNow, codexQuotaScopeForModel(modelId), { ...recoverySelectionOptions, modelEligibleAccountIds }, + modelId, ); const recoveryPreviewAccountId = subagentFallbackAccountPreview( parsed.modelId, diff --git a/tests/bearer-admission-routed-provider.test.ts b/tests/bearer-admission-routed-provider.test.ts index 51f1fd99ee..82a8c75400 100644 --- a/tests/bearer-admission-routed-provider.test.ts +++ b/tests/bearer-admission-routed-provider.test.ts @@ -4,7 +4,15 @@ import { tmpdir } from "node:os"; import { join } from "node:path"; import { saveConfig } from "../src/config"; import { startServer } from "../src/server"; +import { + acquireNativeMainProfileDrain, + getNativeMainProfileRequestCount, +} from "../src/server/lifecycle"; +import { waitForNativeMainStartupGate } from "../src/codex/native-profile-startup"; +import { handleNativeProfileAPI } from "../src/codex/native-profile-api"; +import type { NativeProfileManager } from "../src/codex/native-profile-manager"; import type { OcxConfig } from "../src/types"; +import { ownedServiceHomeInspection } from "./helpers/owned-service-home-inspection"; /** * Issue #2132: bearer admission must not require a stored ChatGPT credential. @@ -33,6 +41,7 @@ let nativeAuth: Array = []; const ADMISSION_SECRET = "ocx_data_2132secret"; const ROUTED_KEY = "sk-routed-provider-key"; +const inspectNativeCodexOwnership = ownedServiceHomeInspection("bearer admission routed provider test"); /** A JWT whose `exp` is far in the future, so a stored main token reads as live. */ function liveJwt(): string { @@ -131,7 +140,7 @@ describe("#2132 bearer admission does not require a ChatGPT credential for route // The reported install: no ChatGPT login was ever performed. writeFileSync(join(codexHome, "auth.json"), JSON.stringify({ tokens: {} })); - const server = startServer(0); + const server = startServer(0, { inspectNativeCodexOwnership }); try { const response = await postResponses(server.url, "gateway/gateway-model"); @@ -151,8 +160,9 @@ describe("#2132 bearer admission does not require a ChatGPT credential for route saveConfig(mixedConfig()); writeFileSync(join(codexHome, "auth.json"), JSON.stringify({ tokens: {} })); - const server = startServer(0); + const server = startServer(0, { inspectNativeCodexOwnership }); try { + await waitForNativeMainStartupGate(); const response = await postResponses(server.url, "gpt-5.5"); // This is the #1686 guarantee and it must survive: a native route genuinely needs the @@ -173,8 +183,9 @@ describe("#2132 bearer admission does not require a ChatGPT credential for route JSON.stringify({ tokens: { access_token: stored, account_id: "stored_main_acc" } }), ); - const server = startServer(0); + const server = startServer(0, { inspectNativeCodexOwnership }); try { + await waitForNativeMainStartupGate(); const response = await postResponses(server.url, "gpt-5.5"); expect(response.status).toBe(200); @@ -223,8 +234,9 @@ describe("an admission bearer never reaches a canonical ChatGPT transport, whate saveConfig(customNamedCanonicalConfig()); writeFileSync(join(codexHome, "auth.json"), JSON.stringify({ tokens: {} })); - const server = startServer(0); + const server = startServer(0, { inspectNativeCodexOwnership }); try { + await waitForNativeMainStartupGate(); const response = await postResponses(server.url, "mirror/gpt-5.5"); // Fail-before-I/O is the contract (src/codex/auth-context.ts): the only two acceptable @@ -245,8 +257,9 @@ describe("an admission bearer never reaches a canonical ChatGPT transport, whate JSON.stringify({ tokens: { access_token: stored, account_id: "stored_main_acc" } }), ); - const server = startServer(0); + const server = startServer(0, { inspectNativeCodexOwnership }); try { + await waitForNativeMainStartupGate(); await postResponses(server.url, "mirror/gpt-5.5"); expect(nativeAuth.join("|")).not.toContain(ADMISSION_SECRET); @@ -255,4 +268,99 @@ describe("an admission bearer never reaches a canonical ChatGPT transport, whate await server.stop(true); } }); + + test("stored-main substitution respects a native-main drain", async () => { + saveConfig(customNamedCanonicalConfig()); + const stored = liveJwt(); + writeFileSync( + join(codexHome, "auth.json"), + JSON.stringify({ tokens: { access_token: stored, account_id: "stored_main_acc" } }), + ); + + const server = startServer(0, { inspectNativeCodexOwnership }); + await waitForNativeMainStartupGate(); + const drain = acquireNativeMainProfileDrain("custom-forward-substitution"); + expect(drain).not.toBeNull(); + try { + const response = await postResponses(server.url, "mirror/gpt-5.5"); + + expect(response.status).toBe(503); + expect(response.headers.get("retry-after")).toBe("1"); + expect(nativeAuth).toHaveLength(0); + expect(routedAuth).toHaveLength(0); + } finally { + drain?.release(); + await server.stop(true); + } + }); + + test("stored-main substitution holds ownership until the upstream request settles", async () => { + saveConfig(customNamedCanonicalConfig()); + const stored = liveJwt(); + writeFileSync( + join(codexHome, "auth.json"), + JSON.stringify({ tokens: { access_token: stored, account_id: "stored_main_acc" } }), + ); + let signalUpstreamStarted!: () => void; + const upstreamStarted = new Promise((resolve) => { signalUpstreamStarted = resolve; }); + let releaseUpstream!: () => void; + const upstreamGate = new Promise((resolve) => { releaseUpstream = resolve; }); + globalThis.fetch = (async (input, init) => { + const raw = input instanceof Request ? input.url : String(input); + const headers = new Headers(input instanceof Request ? input.headers : init?.headers); + if (new URL(raw).hostname === "chatgpt.com") { + nativeAuth.push(headers.get("authorization")); + signalUpstreamStarted(); + await upstreamGate; + return Response.json({ id: "resp_2132_held", object: "response", status: "completed", output: [] }); + } + return originalFetch(input, init); + }) as typeof fetch; + + const server = startServer(0, { inspectNativeCodexOwnership }); + await waitForNativeMainStartupGate(); + const pending = postResponses(server.url, "mirror/gpt-5.5"); + const switchUrl = new URL("http://localhost/api/native-main-profiles/switch"); + const switchRequest = () => new Request(switchUrl, { + method: "POST", + body: JSON.stringify({ target: "target", confirmedStopped: true }), + }); + let switches = 0; + const manager = { + switch: async () => { + switches += 1; + return { ok: true }; + }, + } as unknown as NativeProfileManager; + try { + await upstreamStarted; + expect(getNativeMainProfileRequestCount()).toBe(1); + const blocked = await handleNativeProfileAPI( + switchRequest(), + switchUrl, + {} as OcxConfig, + { manager, drainTimeoutMs: 0 }, + ); + expect(blocked?.status).toBe(409); + expect(switches).toBe(0); + + releaseUpstream(); + const response = await pending; + expect(response.status).toBe(200); + expect(nativeAuth).toEqual([`Bearer ${stored}`]); + expect(getNativeMainProfileRequestCount()).toBe(0); + const switched = await handleNativeProfileAPI( + switchRequest(), + switchUrl, + {} as OcxConfig, + { manager, drainTimeoutMs: 0 }, + ); + expect(switched?.status).toBe(200); + expect(switches).toBe(1); + } finally { + releaseUpstream(); + await pending.catch(() => {}); + await server.stop(true); + } + }); }); diff --git a/tests/codex-auth-context.test.ts b/tests/codex-auth-context.test.ts index d2b4c7d999..6135066cdb 100644 --- a/tests/codex-auth-context.test.ts +++ b/tests/codex-auth-context.test.ts @@ -55,6 +55,7 @@ import { clearCodexUpstreamHealth, clearThreadAccountMap, recordCodexUpstreamOutcome, + resetCodexRoutingForManualSelection, } from "../src/codex/routing"; import type { OcxConfig, OcxProviderConfig } from "../src/types"; import { setIcaclsRunnerForTests } from "../src/lib/windows-secret-acl"; @@ -425,7 +426,11 @@ describe("Codex auth context", () => { test("direct mode returns caller-owned main context without touching pool selection", async () => { const cfg = { ...config(), activeCodexAccountId: "missing-pool-account" }; - await expect(resolveCodexAuthContext(new Headers({ authorization: "Bearer caller" }), cfg, "direct")) + await expect(resolveCodexAuthContext(new Headers({ authorization: "Bearer caller" }), cfg, "direct", { + beginCodexAccountSelection: () => { + throw new Error("caller-owned Direct must not reserve stored main"); + }, + })) .resolves.toEqual({ kind: "main", accountId: null }); }); test("direct mode fails locally without a caller bearer", async () => { @@ -463,6 +468,9 @@ describe("Codex auth context", () => { credentialIdentities: new Map(), }; let callerChecks = 0; + let claimCalls = 0; + let selectionReleases = 0; + let claimed = false; await expect(resolveCodexAuthContext( new Headers({ authorization: "Bearer ocx-admission" }), config(), @@ -470,7 +478,19 @@ describe("Codex auth context", () => { { modelId: "gpt-daybreak-blue-latest", substituteMainCredentialForDirect: true, - resolveCodexModelEntitlements: async () => entitledMain, + beginCodexAccountSelection: () => ({ + mainProfileDraining: false, + claimMainProfile: () => { + claimCalls += 1; + claimed = true; + return true; + }, + release: () => { selectionReleases += 1; }, + }), + resolveCodexModelEntitlements: async () => { + expect(claimed).toBe(true); + return entitledMain; + }, isDirectCallerEntitledToCodexModel: async () => { callerChecks += 1; return false; @@ -478,6 +498,86 @@ describe("Codex auth context", () => { }, )).resolves.toEqual({ kind: "main", accountId: null }); expect(callerChecks).toBe(0); + expect(claimCalls).toBe(1); + expect(selectionReleases).toBe(1); + }); + + test("Direct admission-bearer substitution fails closed during native-main drain", async () => { + let entitlementCalls = 0; + let claimCalls = 0; + let selectionReleases = 0; + await expect(resolveCodexAuthContext( + new Headers({ authorization: "Bearer ocx-admission" }), + config(), + "direct", + { + modelId: "gpt-5.5", + substituteMainCredentialForDirect: true, + beginCodexAccountSelection: () => ({ + mainProfileDraining: true, + claimMainProfile: () => { + claimCalls += 1; + return false; + }, + release: () => { selectionReleases += 1; }, + }), + resolveCodexModelEntitlements: async () => { + entitlementCalls += 1; + throw new Error("must not read stored-main entitlements while draining"); + }, + }, + )).rejects.toBeInstanceOf(CodexMainProfileDrainingError); + expect(entitlementCalls).toBe(0); + expect(claimCalls).toBe(0); + expect(selectionReleases).toBe(1); + }); + + test("Direct admission-bearer substitution fails closed when the atomic claim loses", async () => { + let entitlementCalls = 0; + let claimCalls = 0; + let selectionReleases = 0; + await expect(resolveCodexAuthContext( + new Headers({ authorization: "Bearer ocx-admission" }), + config(), + "direct", + { + modelId: "gpt-daybreak-blue-latest", + substituteMainCredentialForDirect: true, + beginCodexAccountSelection: () => ({ + mainProfileDraining: false, + claimMainProfile: () => { + claimCalls += 1; + return false; + }, + release: () => { selectionReleases += 1; }, + }), + resolveCodexModelEntitlements: async () => { + entitlementCalls += 1; + throw new Error("must not read stored-main entitlements after a lost claim"); + }, + }, + )).rejects.toBeInstanceOf(CodexMainProfileDrainingError); + expect(entitlementCalls).toBe(0); + expect(claimCalls).toBe(1); + expect(selectionReleases).toBe(1); + }); + + test("Direct admission-bearer substitution requires a turn-owned native-main claim", async () => { + let entitlementCalls = 0; + await expect(resolveCodexAuthContext( + new Headers({ authorization: "Bearer ocx-admission" }), + config(), + "direct", + { + modelId: "gpt-daybreak-blue-latest", + substituteMainCredentialForDirect: true, + resolveCodexModelEntitlements: async () => { + entitlementCalls += 1; + throw new Error("must not read stored-main entitlements without admission"); + }, + }, + )).rejects.toBeInstanceOf(CodexMainProfileDrainingError); + expect(entitlementCalls).toBe(0); }); test("account-gated native routing skips an active account without the model grant", async () => { @@ -527,6 +627,72 @@ describe("Codex auth context", () => { }); }); + test("auth resolution preserves per-model detours without replacing ordinary affinity", async () => { + const cfg = config(); + cfg.accountPoolStrategy = "round-robin"; + cfg.accountPoolStickyLimit = 1; + cfg.activeCodexAccountId = "pool-b"; + cfg.activeCodexAccountPinned = "pool-b"; + cfg.codexAccounts = [ + { id: "pool-a", email: "a@example.test", isMain: false, chatgptAccountId: "pool_acc_a" }, + { id: "pool-b", email: "b@example.test", isMain: false, chatgptAccountId: "pool_acc_b" }, + { id: "pool-c", email: "c@example.test", isMain: false, chatgptAccountId: "pool_acc_c" }, + ]; + for (const id of ["pool-a", "pool-b", "pool-c"]) { + saveCodexAccountCredential(id, { + accessToken: `${id}-token`, + refreshToken: `${id}-refresh`, + expiresAt: Date.now() + 5 * 60_000, + chatgptAccountId: `${id}-chatgpt`, + }); + } + resetCodexRoutingForManualSelection("pool-b"); + const headers = new Headers({ "x-codex-parent-thread-id": "auth-model-detour-thread" }); + const primeCodexPoolQuotas = async () => {}; + await expect(resolveCodexAuthContext(headers, cfg, "pool", { + modelId: "gpt-5.5", + primeCodexPoolQuotas, + })).resolves.toMatchObject({ kind: "pool", accountId: "pool-b" }); + + const entitlementSnapshot: CodexModelEntitlementSnapshot = { + modelsByAccount: new Map([ + ["pool-a", new Set(["gpt-daybreak-blue-latest", "gpt-5.6-sol"])], + ["pool-c", new Set(["gpt-daybreak-blue-latest", "gpt-5.6-sol"])], + ]), + confirmedAccountIds: new Set(["pool-a", "pool-b", "pool-c"]), + credentialIdentities: new Map(), + }; + const gatedOptions = { + primeCodexPoolQuotas, + resolveCodexModelEntitlements: async () => entitlementSnapshot, + }; + const first = await resolveCodexAuthContext(headers, cfg, "pool", { + ...gatedOptions, + modelId: "gpt-daybreak-blue-latest", + }); + expect(first).toMatchObject({ kind: "pool" }); + await expect(resolveCodexAuthContext(headers, cfg, "pool", { + ...gatedOptions, + modelId: "gpt-daybreak-blue-latest", + })).resolves.toMatchObject({ kind: "pool", accountId: first.accountId }); + + const secondModel = await resolveCodexAuthContext(headers, cfg, "pool", { + ...gatedOptions, + modelId: "gpt-5.6-sol", + }); + expect(secondModel).toMatchObject({ kind: "pool" }); + expect(secondModel.accountId).not.toBe(first.accountId); + await expect(resolveCodexAuthContext(headers, cfg, "pool", { + ...gatedOptions, + modelId: "gpt-5.6-sol", + })).resolves.toMatchObject({ kind: "pool", accountId: secondModel.accountId }); + + await expect(resolveCodexAuthContext(headers, cfg, "pool", { + modelId: "gpt-5.5", + primeCodexPoolQuotas, + })).resolves.toMatchObject({ kind: "pool", accountId: "pool-b" }); + }); + test("exact account-gated routing fails closed for an unentitled account", async () => { const cfg = config(); saveCodexAccountCredential("pool-a", { diff --git a/tests/codex-routing.test.ts b/tests/codex-routing.test.ts index f5bfcc31d9..0c651daf2d 100644 --- a/tests/codex-routing.test.ts +++ b/tests/codex-routing.test.ts @@ -1822,6 +1822,266 @@ describe("codex account selection order", () => { expect(resolveCodexAccountForThread("model-gated-task", config, now + 2, "shared")).toBe("b"); }); + test("repeated model-gated round-robin requests reuse a separate detour affinity", () => { + const now = 1_800_000_000_000; + const threadId = "model-detour-affinity"; + const modelId = "gpt-daybreak-blue-latest"; + const config = makeConfig({ + accountPoolStrategy: "round-robin", + accountPoolStickyLimit: 1, + activeCodexAccountId: "b", + activeCodexAccountPinned: "b", + codexAccounts: [ + { id: "a", email: "a@test", isMain: false }, + { id: "b", email: "b@test", isMain: false }, + { id: "c", email: "c@test", isMain: false }, + ], + }); + saveTestCredential("c"); + updateAccountQuota("a", 10); + updateAccountQuota("b", 10); + updateAccountQuota("c", 10); + resetCodexRoutingForManualSelection("b"); + expect(resolveCodexAccountForThread(threadId, config, now, "shared")).toBe("b"); + + const eligible = { modelEligibleAccountIds: new Set(["a", "c"]) }; + const firstPreview = previewCodexAccountForRequest( + threadId, + config, + now + 1, + "shared", + eligible, + modelId, + ); + const first = resolveCodexAccountForThreadDetailed( + threadId, + config, + now + 1, + "shared", + eligible, + modelId, + ); + expect(first).toEqual({ status: "selected", accountId: firstPreview }); + expect(["a", "c"]).toContain(firstPreview); + + expect(previewCodexAccountForRequest( + threadId, + config, + now + 2, + "shared", + eligible, + modelId, + )).toBe(firstPreview); + expect(resolveCodexAccountForThreadDetailed( + threadId, + config, + now + 2, + "shared", + eligible, + modelId, + )).toEqual(first); + + expect(config.activeCodexAccountId).toBe("b"); + expect(config.activeCodexAccountPinned).toBe("b"); + expect(getEffectiveActiveCodexAccountId(config)).toBe("b"); + expect(resolveCodexAccountForThread(threadId, config, now + 3, "shared")).toBe("b"); + + const other = firstPreview === "a" ? "c" : "a"; + expect(resolveCodexAccountForThreadDetailed( + threadId, + config, + now + 4, + "shared", + { modelEligibleAccountIds: new Set([other]) }, + modelId, + )).toEqual({ status: "selected", accountId: other }); + expect(resolveCodexAccountForThreadDetailed( + threadId, + config, + now + 5, + "shared", + { modelEligibleAccountIds: new Set(["a", "b", "c"]) }, + modelId, + )).toEqual({ status: "selected", accountId: other }); + expect(resolveCodexAccountForThread(threadId, config, now + 6, "shared")).toBe("b"); + }); + + test("model preview and final keep a live detour after ordinary affinity cleanup", () => { + const now = 1_800_000_000_000; + const threadId = "detour-after-ordinary-cleanup"; + const modelId = "gpt-daybreak-blue-latest"; + const config = makeConfig({ + accountPoolStrategy: "round-robin", + accountPoolStickyLimit: 1, + activeCodexAccountId: "b", + activeCodexAccountPinned: "b", + codexAccounts: [ + { id: "a", email: "a@test", isMain: false }, + { id: "b", email: "b@test", isMain: false }, + { id: "c", email: "c@test", isMain: false }, + ], + }); + saveTestCredential("c"); + resetCodexRoutingForManualSelection("b"); + expect(resolveCodexAccountForThread(threadId, config, now, "shared")).toBe("b"); + const eligible = { modelEligibleAccountIds: new Set(["a", "c"]) }; + const first = resolveCodexAccountForThreadDetailed( + threadId, + config, + now + 1, + "shared", + eligible, + modelId, + ); + expect(first.status).toBe("selected"); + + clearThreadAccountMapForAccount("b"); + if (first.status === "selected") { + expect(previewCodexAccountForRequest( + threadId, + config, + now + 2, + "shared", + eligible, + modelId, + )).toBe(first.accountId); + expect(resolveCodexAccountForThreadDetailed( + threadId, + config, + now + 2, + "shared", + eligible, + modelId, + )).toEqual(first); + } + expect(config.activeCodexAccountPinned).toBe("b"); + }); + + test("model detour affinities are independent within one quota scope", () => { + const now = 1_800_000_000_000; + const threadId = "independent-model-detours"; + const config = makeConfig({ + accountPoolStrategy: "round-robin", + accountPoolStickyLimit: 1, + activeCodexAccountId: "b", + activeCodexAccountPinned: "b", + codexAccounts: [ + { id: "a", email: "a@test", isMain: false }, + { id: "b", email: "b@test", isMain: false }, + { id: "c", email: "c@test", isMain: false }, + ], + }); + saveTestCredential("c"); + resetCodexRoutingForManualSelection("b"); + expect(resolveCodexAccountForThread(threadId, config, now, "shared")).toBe("b"); + const eligible = { modelEligibleAccountIds: new Set(["a", "c"]) }; + + const firstModel = resolveCodexAccountForThreadDetailed( + threadId, + config, + now + 1, + "shared", + eligible, + "gpt-daybreak-blue-latest", + ); + const secondModel = resolveCodexAccountForThreadDetailed( + threadId, + config, + now + 2, + "shared", + eligible, + "gpt-other-account-gated", + ); + expect(firstModel.status).toBe("selected"); + expect(secondModel.status).toBe("selected"); + if (firstModel.status === "selected" && secondModel.status === "selected") { + expect(secondModel.accountId).not.toBe(firstModel.accountId); + expect(resolveCodexAccountForThreadDetailed( + threadId, + config, + now + 3, + "shared", + eligible, + "gpt-daybreak-blue-latest", + )).toEqual(firstModel); + expect(resolveCodexAccountForThreadDetailed( + threadId, + config, + now + 4, + "shared", + eligible, + "gpt-other-account-gated", + )).toEqual(secondModel); + } + expect(resolveCodexAccountForThread(threadId, config, now + 5, "shared")).toBe("b"); + }); + + test("model detour LRU stays bounded without evicting ordinary affinity", () => { + const now = 1_800_000_000_000; + const threadId = "bounded-model-detours"; + const config = makeConfig({ + accountPoolStrategy: "round-robin", + accountPoolStickyLimit: 1, + activeCodexAccountId: "b", + activeCodexAccountPinned: "b", + autoSwitchThreshold: 0, + codexAccounts: [ + { id: "a", email: "a@test", isMain: false }, + { id: "b", email: "b@test", isMain: false }, + { id: "c", email: "c@test", isMain: false }, + ], + }); + saveTestCredential("c"); + resetCodexRoutingForManualSelection("b"); + expect(resolveCodexAccountForThread(threadId, config, now, "shared")).toBe("b"); + const eligible = { modelEligibleAccountIds: new Set(["a", "c"]) }; + + expect(resolveCodexAccountForThreadDetailed( + threadId, + config, + now + 1, + "shared", + eligible, + "gated-model-0", + )).toEqual({ status: "selected", accountId: "a" }); + for (let index = 1; index <= CODEX_THREAD_AFFINITY_MAX_ENTRIES; index += 1) { + expect(resolveCodexAccountForThreadDetailed( + threadId, + config, + now + index + 1, + "shared", + eligible, + `gated-model-${index}`, + ).status).toBe("selected"); + } + + // Detours are the preferred LRU victims, so model churn cannot displace the + // task's ordinary account. The oldest detour was evicted; recreating it takes + // the next RR account and then becomes sticky again. + expect(resolveCodexAccountForThread( + threadId, + config, + now + CODEX_THREAD_AFFINITY_MAX_ENTRIES + 3, + "shared", + )).toBe("b"); + expect(resolveCodexAccountForThreadDetailed( + threadId, + config, + now + CODEX_THREAD_AFFINITY_MAX_ENTRIES + 4, + "shared", + eligible, + "gated-model-0", + )).toEqual({ status: "selected", accountId: "c" }); + expect(resolveCodexAccountForThreadDetailed( + threadId, + config, + now + CODEX_THREAD_AFFINITY_MAX_ENTRIES + 5, + "shared", + eligible, + "gated-model-0", + )).toEqual({ status: "selected", accountId: "c" }); + }, STORE_BUDGET_MS); + test("a gated first request binds its actual account without replacing global active", () => { const config = makeConfig({ activeCodexAccountId: "b" }); const now = Date.now(); @@ -1922,12 +2182,22 @@ describe("codex account selection order", () => { }); } + const selectionOptions = { modelEligibleAccountIds: new Set(["a", "b"]) }; + expect(previewCodexAccountForRequest( + null, + config, + now + CODEX_TRANSIENT_SOFT_AVOID_MS + 4, + "shared", + selectionOptions, + "gpt-daybreak-blue-latest", + )).toBe("b"); + expect(resolveCodexAccountForThreadDetailed( null, config, now + CODEX_TRANSIENT_SOFT_AVOID_MS + 4, "shared", - { modelEligibleAccountIds: new Set(["a", "b"]) }, + selectionOptions, )).toEqual({ status: "selected", accountId: "b" }); expect(config.activeCodexAccountId).toBe("c"); expect(config.activeCodexAccountPinned).toBe("c"); @@ -2020,12 +2290,22 @@ describe("codex account selection order", () => { updateAccountQuota("b", 90); resetCodexRoutingForManualSelection("b"); + const selectionOptions = { modelEligibleAccountIds: new Set(["a", "b"]) }; + expect(previewCodexAccountForRequest( + null, + config, + now, + "shared", + selectionOptions, + "gpt-daybreak-blue-latest", + )).toBe("a"); + expect(resolveCodexAccountForThreadDetailed( null, config, now, "shared", - { modelEligibleAccountIds: new Set(["a", "b"]) }, + selectionOptions, )).toEqual({ status: "selected", accountId: "a" }); expect(config.activeCodexAccountId).toBe("b"); expect(config.activeCodexAccountPinned).toBeUndefined(); @@ -2053,12 +2333,22 @@ describe("codex account selection order", () => { } const resolveAt = now + CODEX_TRANSIENT_SOFT_AVOID_MS + 4; + const selectionOptions = { modelEligibleAccountIds: new Set(["a", "b"]) }; + expect(previewCodexAccountForRequest( + null, + config, + resolveAt, + "shared", + selectionOptions, + "gpt-daybreak-blue-latest", + )).toBe("a"); + expect(resolveCodexAccountForThreadDetailed( null, config, resolveAt, "shared", - { modelEligibleAccountIds: new Set(["a", "b"]) }, + selectionOptions, )).toEqual({ status: "selected", accountId: "a" }); expect(config.activeCodexAccountId).toBe("b"); expect(config.activeCodexAccountPinned).toBeUndefined(); diff --git a/tests/subagent-fallback-handle-responses.test.ts b/tests/subagent-fallback-handle-responses.test.ts index 3997bc8670..9212cbafe3 100644 --- a/tests/subagent-fallback-handle-responses.test.ts +++ b/tests/subagent-fallback-handle-responses.test.ts @@ -1551,7 +1551,9 @@ describe("native fallback account preview", () => { } // And both must actually forward it into the preview call, not merely accept it. - const forwarded = source.match(/\{ \.\.\.(previewSelectionOptions|recoverySelectionOptions), modelEligibleAccountIds \}/g) ?? []; + const forwarded = source.match( + /\{ \.\.\.(previewSelectionOptions|recoverySelectionOptions), modelEligibleAccountIds \},\s*modelId,\s*\)/g, + ) ?? []; expect(forwarded).toHaveLength(2); }); From 80d679463130473e0b82835bfc116ff36fffda06 Mon Sep 17 00:00:00 2001 From: luvs01 Date: Fri, 28 Aug 2026 21:20:41 +0900 Subject: [PATCH 5/6] fix(codex): address routing review follow-ups --- src/codex/routing.ts | 106 +++++++++--------- .../bearer-admission-routed-provider.test.ts | 38 ++++--- 2 files changed, 75 insertions(+), 69 deletions(-) diff --git a/src/codex/routing.ts b/src/codex/routing.ts index 100a7e0a9e..95e6ed06b7 100644 --- a/src/codex/routing.ts +++ b/src/codex/routing.ts @@ -1631,6 +1631,45 @@ function previewReusableAffinityAccount( return entry.accountId; } +/** + * Re-evaluate an affined account under the quota strategy. Returns a strictly + * cooler replacement, or null when the current binding should remain. + */ +function reevaluateAffinityQuota( + entry: ThreadAffinityEntry, + config: OcxConfig, + now: number, + quotaScope?: CodexQuotaScope, + selectionOptions?: CodexAccountUsabilityOptions, +): string | null { + if (normalizeAccountPoolStrategy(config.accountPoolStrategy) !== "quota") return null; + const threshold = config.autoSwitchThreshold ?? 80; + const usage = threshold > 0 + ? computeCodexUsageScore( + getAccountQuota(entry.accountId), + getPoolAccountPlanForSelection(config, entry.accountId, selectionOptions), + ) + : 0; + const overThreshold = threshold > 0 && !isUnknownUsage(usage) && usage >= threshold; + if ( + !overThreshold + && now - entry.lastReevalAt < CODEX_THREAD_AFFINITY_REEVAL_INTERVAL_MS + ) { + return null; + } + entry.lastReevalAt = now; + if (!overThreshold) return null; + const best = pickLowerUsageAccount( + config, + entry.accountId, + usage, + now, + quotaScope, + selectionOptions, + ); + return best === entry.accountId ? null : best; +} + /** * Side-effect-free preview of the Codex pool account native routing would prefer. * Used for subagent fallback quota decisions before final auth. @@ -1766,35 +1805,16 @@ export function resolveCodexAccountForThreadDetailed( // Model detours follow the same affinity policy as ordinary bindings: // RR/fill-first stay sticky, while quota strategy may re-evaluate an // over-threshold account without changing the ordinary lane. - if (normalizeAccountPoolStrategy(config.accountPoolStrategy) === "quota") { - const threshold = config.autoSwitchThreshold ?? 80; - const usage = threshold > 0 - ? computeCodexUsageScore( - getAccountQuota(detourEntry.accountId), - getPoolAccountPlanForSelection(config, detourEntry.accountId, selectionOptions), - ) - : 0; - const overThreshold = threshold > 0 && !isUnknownUsage(usage) && usage >= threshold; - if ( - overThreshold - || now - detourEntry.lastReevalAt >= CODEX_THREAD_AFFINITY_REEVAL_INTERVAL_MS - ) { - detourEntry.lastReevalAt = now; - if (overThreshold) { - const best = pickLowerUsageAccount( - config, - detourEntry.accountId, - usage, - now, - quotaScope, - selectionOptions, - ); - if (best !== detourEntry.accountId) { - bindModelDetourAffinity(threadId, best, now, modelId, quotaScope); - return { status: "selected", accountId: best }; - } - } - } + const cooler = reevaluateAffinityQuota( + detourEntry, + config, + now, + quotaScope, + selectionOptions, + ); + if (cooler) { + bindModelDetourAffinity(threadId, cooler, now, modelId, quotaScope); + return { status: "selected", accountId: cooler }; } return { status: "selected", accountId: detourEntry.accountId }; } @@ -1834,29 +1854,13 @@ export function resolveCodexAccountForThreadDetailed( // serving for up to 60s after a secondary with quota is available (#584). // Non-quota strategies (RR / fill-first) keep affinity for ongoing threads — // rotation is new-session-only (affinity policy A). - const strategy = normalizeAccountPoolStrategy(config.accountPoolStrategy); - if (strategy === "quota") { - const threshold = config.autoSwitchThreshold ?? 80; - const usage = threshold > 0 - ? computeCodexUsageScore( - getAccountQuota(entry.accountId), - getPoolAccountPlanForSelection(config, entry.accountId, selectionOptions), - ) - : 0; - const overThreshold = threshold > 0 && !isUnknownUsage(usage) && usage >= threshold; - if (overThreshold || now - entry.lastReevalAt >= CODEX_THREAD_AFFINITY_REEVAL_INTERVAL_MS) { - entry.lastReevalAt = now; - if (overThreshold) { - const best = pickLowerUsageAccount(config, entry.accountId, usage, now, quotaScope, selectionOptions); - if (best !== entry.accountId) { - if (!isIndependentCodexQuotaScope(quotaScope)) { - setActiveCodexAccount(config, best); - } - bindThreadAffinity(threadId, best, now, quotaScope); // rebinds + resets clocks - return { status: "selected", accountId: best }; - } - } + const cooler = reevaluateAffinityQuota(entry, config, now, quotaScope, selectionOptions); + if (cooler) { + if (!isIndependentCodexQuotaScope(quotaScope)) { + setActiveCodexAccount(config, cooler); } + bindThreadAffinity(threadId, cooler, now, quotaScope); // rebinds + resets clocks + return { status: "selected", accountId: cooler }; } return { status: "selected", accountId: entry.accountId }; } diff --git a/tests/bearer-admission-routed-provider.test.ts b/tests/bearer-admission-routed-provider.test.ts index 82a8c75400..6b6b2af943 100644 --- a/tests/bearer-admission-routed-provider.test.ts +++ b/tests/bearer-admission-routed-provider.test.ts @@ -278,10 +278,11 @@ describe("an admission bearer never reaches a canonical ChatGPT transport, whate ); const server = startServer(0, { inspectNativeCodexOwnership }); - await waitForNativeMainStartupGate(); - const drain = acquireNativeMainProfileDrain("custom-forward-substitution"); - expect(drain).not.toBeNull(); + let drain: ReturnType = null; try { + await waitForNativeMainStartupGate(); + drain = acquireNativeMainProfileDrain("custom-forward-substitution"); + expect(drain).not.toBeNull(); const response = await postResponses(server.url, "mirror/gpt-5.5"); expect(response.status).toBe(503); @@ -318,21 +319,22 @@ describe("an admission bearer never reaches a canonical ChatGPT transport, whate }) as typeof fetch; const server = startServer(0, { inspectNativeCodexOwnership }); - await waitForNativeMainStartupGate(); - const pending = postResponses(server.url, "mirror/gpt-5.5"); - const switchUrl = new URL("http://localhost/api/native-main-profiles/switch"); - const switchRequest = () => new Request(switchUrl, { - method: "POST", - body: JSON.stringify({ target: "target", confirmedStopped: true }), - }); - let switches = 0; - const manager = { - switch: async () => { - switches += 1; - return { ok: true }; - }, - } as unknown as NativeProfileManager; + let pending: Promise | null = null; try { + await waitForNativeMainStartupGate(); + pending = postResponses(server.url, "mirror/gpt-5.5"); + const switchUrl = new URL("http://localhost/api/native-main-profiles/switch"); + const switchRequest = () => new Request(switchUrl, { + method: "POST", + body: JSON.stringify({ target: "target", confirmedStopped: true }), + }); + let switches = 0; + const manager = { + switch: async () => { + switches += 1; + return { ok: true }; + }, + } as unknown as NativeProfileManager; await upstreamStarted; expect(getNativeMainProfileRequestCount()).toBe(1); const blocked = await handleNativeProfileAPI( @@ -359,7 +361,7 @@ describe("an admission bearer never reaches a canonical ChatGPT transport, whate expect(switches).toBe(1); } finally { releaseUpstream(); - await pending.catch(() => {}); + await pending?.catch(() => {}); await server.stop(true); } }); From 375e6f8fb849db0d60f61220cfe8716d6ab15445 Mon Sep 17 00:00:00 2001 From: luvs01 Date: Fri, 28 Aug 2026 21:51:48 +0900 Subject: [PATCH 6/6] fix(codex): skip failing quota candidates --- src/codex/routing.ts | 17 +++++- tests/codex-routing.test.ts | 101 ++++++++++++++++++++++++++++++++++++ 2 files changed, 116 insertions(+), 2 deletions(-) diff --git a/src/codex/routing.ts b/src/codex/routing.ts index 95e6ed06b7..b1d5a26b01 100644 --- a/src/codex/routing.ts +++ b/src/codex/routing.ts @@ -1006,12 +1006,14 @@ function getEligiblePoolAccounts( now = Date.now(), quotaScope?: CodexQuotaScope, selectionOptions?: CodexAccountUsabilityOptions, + skipFailoverReadyCandidates = false, ): readonly string[] { const ids = (config.codexAccounts ?? []) .filter(account => isSelectableCodexPoolAccount(account) && account.id !== excludeId && !isCodexAccountPaused(config, account.id) - && !isAccountNeedsReauth(account.id)) + && !isAccountNeedsReauth(account.id) + && (!skipFailoverReadyCandidates || !shouldFailover(config, account.id, now))) .filter(account => getCodexQuotaHealthSnapshot(account.id, quotaScope, now) === null) .filter(account => !isCodexAccountSoftAvoided(account.id, now)) .filter(account => isCodexAccountUsable(config, account.id, selectionOptions)) @@ -1024,6 +1026,7 @@ function getEligiblePoolAccounts( && !isAccountNeedsReauth(MAIN_CODEX_ACCOUNT_ID) && getCodexQuotaHealthSnapshot(MAIN_CODEX_ACCOUNT_ID, quotaScope, now) === null && !isCodexAccountSoftAvoided(MAIN_CODEX_ACCOUNT_ID, now) + && (!skipFailoverReadyCandidates || !shouldFailover(config, MAIN_CODEX_ACCOUNT_ID, now)) && isCodexAccountUsable(config, MAIN_CODEX_ACCOUNT_ID, selectionOptions) ) { ids.unshift(MAIN_CODEX_ACCOUNT_ID); @@ -1244,10 +1247,18 @@ function pickLowerUsageAccount( now: number, quotaScope?: CodexQuotaScope, selectionOptions?: CodexAccountUsabilityOptions, + skipFailoverReadyCandidates = false, ): string { let best = active; let bestUsage = activeUsage; - for (const id of getEligiblePoolAccounts(config, active, now, quotaScope, selectionOptions)) { + for (const id of getEligiblePoolAccounts( + config, + active, + now, + quotaScope, + selectionOptions, + skipFailoverReadyCandidates, + )) { const usage = computeCodexUsageScore( getAccountQuota(id), getPoolAccountPlanForSelection(config, id, selectionOptions), @@ -1623,6 +1634,7 @@ function previewReusableAffinityAccount( now, quotaScope, selectionOptions, + true, ); if (best !== entry.accountId) return best; } @@ -1666,6 +1678,7 @@ function reevaluateAffinityQuota( now, quotaScope, selectionOptions, + true, ); return best === entry.accountId ? null : best; } diff --git a/tests/codex-routing.test.ts b/tests/codex-routing.test.ts index 0c651daf2d..eb93bc14c2 100644 --- a/tests/codex-routing.test.ts +++ b/tests/codex-routing.test.ts @@ -1906,6 +1906,107 @@ describe("codex account selection order", () => { expect(resolveCodexAccountForThread(threadId, config, now + 6, "shared")).toBe("b"); }); + test("quota detour re-evaluation skips failover-ready cooler candidates", () => { + const now = 1_800_000_000_000; + const threadId = "quota-detour-failover-candidate"; + const modelId = "gpt-daybreak-blue-latest"; + const config = makeConfig({ + accountPoolStrategy: "quota", + activeCodexAccountId: "c", + codexAccounts: [ + { id: "a", email: "a@test", isMain: false }, + { id: "b", email: "b@test", isMain: false }, + { id: "c", email: "c@test", isMain: false }, + ], + }); + saveTestCredential("c"); + updateAccountQuota("a", 10); + updateAccountQuota("b", 5); + updateAccountQuota("c", 10); + resetCodexRoutingForManualSelection("c"); + expect(resolveCodexAccountForThread(threadId, config, now, "shared")).toBe("c"); + expect(resolveCodexAccountForThreadDetailed( + threadId, + config, + now + 1, + "shared", + { modelEligibleAccountIds: new Set(["a"]) }, + modelId, + )).toEqual({ status: "selected", accountId: "a" }); + // B is the highest tier after the detour exists. Filtering only after tier + // selection would drop B without ever exposing healthy C to the picker. + config.codexAccountPriorities = { b: 2, c: 1 }; + + for (let attempt = 0; attempt < 3; attempt += 1) { + recordCodexUpstreamOutcome(config, "b", 503, { + fixedAccount: true, + now: now + attempt + 2, + }); + } + updateAccountQuota("a", 90); + const resolveAt = now + CODEX_TRANSIENT_SOFT_AVOID_MS + 5; + const eligible = { modelEligibleAccountIds: new Set(["a", "b", "c"]) }; + + expect(previewCodexAccountForRequest( + threadId, + config, + resolveAt, + "shared", + eligible, + modelId, + )).toBe("c"); + expect(resolveCodexAccountForThreadDetailed( + threadId, + config, + resolveAt, + "shared", + eligible, + modelId, + )).toEqual({ status: "selected", accountId: "c" }); + expect(config.activeCodexAccountId).toBe("c"); + expect(config.activeCodexAccountPinned).toBeUndefined(); + expect(getEffectiveActiveCodexAccountId(config)).toBe("c"); + }); + + test("ordinary quota affinity re-evaluation skips a failover-ready higher tier", () => { + const now = 1_800_000_000_000; + const threadId = "ordinary-quota-failover-candidate"; + const config = makeConfig({ + accountPoolStrategy: "quota", + activeCodexAccountId: "a", + codexAccounts: [ + { id: "a", email: "a@test", isMain: false }, + { id: "b", email: "b@test", isMain: false }, + { id: "c", email: "c@test", isMain: false }, + ], + }); + saveTestCredential("c"); + updateAccountQuota("a", 10); + updateAccountQuota("b", 5); + updateAccountQuota("c", 10); + expect(resolveCodexAccountForThread(threadId, config, now, "shared")).toBe("a"); + config.codexAccountPriorities = { b: 2, c: 1 }; + + for (let attempt = 0; attempt < 3; attempt += 1) { + recordCodexUpstreamOutcome(config, "b", 503, { + fixedAccount: true, + now: now + attempt + 1, + }); + } + updateAccountQuota("a", 90); + const resolveAt = now + CODEX_TRANSIENT_SOFT_AVOID_MS + 4; + + expect(previewCodexAccountForRequest(threadId, config, resolveAt, "shared")).toBe("c"); + expect(resolveCodexAccountForThreadDetailed( + threadId, + config, + resolveAt, + "shared", + )).toEqual({ status: "selected", accountId: "c" }); + expect(config.activeCodexAccountId).toBe("c"); + expect(getEffectiveActiveCodexAccountId(config)).toBe("c"); + }); + test("model preview and final keep a live detour after ordinary affinity cleanup", () => { const now = 1_800_000_000_000; const threadId = "detour-after-ordinary-cleanup";