diff --git a/src/codex/auth-context.ts b/src/codex/auth-context.ts index f3c357d4c3..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 @@ -383,6 +412,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 +432,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, }; @@ -425,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) { @@ -436,12 +473,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/routing.ts b/src/codex/routing.ts index 10160fe913..b1d5a26b01 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,22 +975,45 @@ 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, 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)) @@ -954,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); @@ -964,7 +1037,7 @@ function getEligiblePoolAccounts( return selectPriorityTier( ids, codexAccountPriorityLookup(config), - id => hasCodexQuotaHeadroom(config, id), + id => hasCodexQuotaHeadroom(config, id, selectionOptions), pinnedCodexAccountId(config), ); } @@ -992,10 +1065,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 +1094,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 +1114,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 +1129,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 +1140,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 +1149,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 +1167,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 +1183,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 +1194,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 +1210,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, @@ -1133,11 +1247,22 @@ 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)) { - const usage = computeCodexUsageScore(getAccountQuota(id), getPoolAccountPlan(config, id)); + for (const id of getEligiblePoolAccounts( + config, + active, + now, + quotaScope, + selectionOptions, + skipFailoverReadyCandidates, + )) { + const usage = computeCodexUsageScore( + getAccountQuota(id), + getPoolAccountPlanForSelection(config, id, selectionOptions), + ); if (usage < bestUsage) { best = id; bestUsage = usage; @@ -1147,11 +1272,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 +1299,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 +1443,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 +1466,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 +1498,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 +1530,49 @@ 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 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, 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 +1583,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; @@ -1398,6 +1601,88 @@ 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, + true, + ); + if (best !== entry.accountId) return best; + } + } + } + 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, + true, + ); + 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. @@ -1412,42 +1697,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), - getPoolAccountPlan(config, entry.accountId), - ); - 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, @@ -1455,7 +1729,7 @@ export function previewCodexAccountForRequest( now, false, quotaScope, - selectionOptions, + strategySelectionOptionsForModelDetour(config, now, quotaScope, selectionOptions), ); if (strategyPick) return strategyPick; @@ -1476,7 +1750,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); } @@ -1501,12 +1778,64 @@ 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. + 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, + ) + ); + + 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. + 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 }; + } + // 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) { @@ -1514,12 +1843,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 @@ -1530,48 +1867,107 @@ 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), - getPoolAccountPlan(config, entry.accountId), - ) - : 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 }; } - 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 = strategySelectionOptionsForModelDetour( + config, + now, + quotaScope, + selectionOptions, + ); + const strategyPick = pickUnboundStrategyAccount( + config, + threadId, + now, + true, + quotaScope, + strategySelectionOptions, + !modelScopedSelection, + !preserveExistingModelScopedAffinity, + ); + if (strategyPick) { + if (threadId && preserveExistingModelScopedAffinity) { + bindModelDetourAffinity(threadId, strategyPick, now, modelId, quotaScope); + } + 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 +1985,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 +2020,13 @@ export function resolveCodexAccountForThreadDetailed( ? { status: "selected", accountId: active } : { status: "none" }; } - if (threadId) 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/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..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"; @@ -893,6 +899,7 @@ interface CodexPoolAccountRetryArgs { codexWsRuntimeIdentity?: BunRuntimeGateInput; translatorBudget: TranslatorBudget; turnAdmissionLease?: AdmissionLease; + resolveCodexModelEntitlements?: typeof resolveCodexModelEntitlements; }; firstAuthCtx: Extract; firstResponse: Response; @@ -923,6 +930,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: @@ -1011,10 +1041,22 @@ 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 resolveCodexRetryModelEntitlements( + config, + entitlementResolver, + options.turnAdmissionLease, + ); + } catch (error) { + await firstResponse.body?.cancel().catch(() => undefined); + releaseCodexAuthContextProbeLease(firstAuthCtx); + 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 +1076,20 @@ async function retryCodexPoolOnAlternateAccount( excludeAccountId: firstAuthCtx.accountId, modelId: route.modelId, beginCodexAccountSelection: codexAccountSelectionForTurn(options.turnAdmissionLease), + resolveCodexModelEntitlements: entitlementResolver, }, ); } 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); + throw error; + } } if (retryAuthCtx?.kind !== "pool" && retryAuthCtx?.kind !== "main-pool") { return { kind: "no-alternate" }; @@ -1131,24 +1178,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 +1211,23 @@ async function retryCodexPoolOnAlternateAccount( options.abortSignal, )) break; invalidateCodexModelEntitlementsForAccount(retryAuthCtx.accountId); - const refreshed = await resolveCodexModelEntitlements(config); + let refreshed: Awaited>; + try { + refreshed = await resolveCodexRetryModelEntitlements( + config, + entitlementResolver, + options.turnAdmissionLease, + ); + } catch (error) { + await upstreamResponse.body?.cancel().catch(() => undefined); + await firstResponse.body?.cancel().catch(() => undefined); + releaseCodexAuthContextProbeLease(firstAuthCtx); + releaseCodexAuthContextProbeLease(retryAuthCtx); + 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?.(); } @@ -1606,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); } @@ -2530,6 +2607,7 @@ async function handleResponsesInner( previewNow, codexQuotaScopeForModel(modelId), { ...previewSelectionOptions, modelEligibleAccountIds }, + modelId, ); const previewAccountId = route.codexAccountId ?? subagentFallbackAccountPreview( route.modelId, @@ -2651,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..6b6b2af943 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,101 @@ 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 }); + 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); + 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 }); + 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( + 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 da4d6cb5d0..6135066cdb 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, @@ -51,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"; @@ -86,6 +91,7 @@ beforeEach(() => { clearThreadAccountMap(); clearCodexUpstreamHealth(); clearAccountQuota(); + setMainAccountPlan(null); __resetGuardianState(); clearAccountNeedsReauth("pool-a"); clearAccountNeedsReauth("pool-b"); @@ -97,6 +103,7 @@ afterEach(() => { clearThreadAccountMap(); clearCodexUpstreamHealth(); clearAccountQuota(); + setMainAccountPlan(null); __resetGuardianState(); clearAccountNeedsReauth("pool-a"); clearAccountNeedsReauth("pool-b"); @@ -122,6 +129,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,15 +338,99 @@ 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")) + 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 () => { @@ -371,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(), @@ -378,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; @@ -386,10 +498,91 @@ 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 () => { 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 +611,86 @@ 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("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 () => { diff --git a/tests/codex-routing.test.ts b/tests/codex-routing.test.ts index 4bb678a95f..eb93bc14c2 100644 --- a/tests/codex-routing.test.ts +++ b/tests/codex-routing.test.ts @@ -1801,6 +1801,802 @@ 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("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("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"; + 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(); + 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, + }); + } + + 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", + selectionOptions, + )).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"); + + 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", + selectionOptions, + )).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; + + 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", + selectionOptions, + )).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..9212cbafe3 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, @@ -28,9 +29,8 @@ import { resetSubagentModelFallbackStateForTests, setSubagentQuotaPrimeForTests, } from "../src/codex/subagent-model-fallback"; -import { - 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"; import { resetAgentTaskRecoveryState } from "../src/server/responses/agent-task-recovery"; @@ -64,6 +64,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 +78,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 +97,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 +259,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; @@ -788,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({ @@ -802,6 +849,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() { @@ -826,23 +875,108 @@ describe("native fallback account preview", () => { turnAdmissionLease, resolveCodexModelEntitlements: async () => { resolverCalls += 1; + signalResolverEntered(); return discovery; }, }, ); - for (let i = 0; i < 20 && resolverCalls === 0; i += 1) await Promise.resolve(); + await resolverEntered; 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"); + 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 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); + 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 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; + globalThis.fetch = (async () => { + fetchCalls += 1; + throw new Error("must not dispatch"); + }) as typeof fetch; + + 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(entitlementCalls).toBe(2); + expect(beginCount).toBe(2); + expect(releaseCount).toBe(2); + expect(fetchCalls).toBe(0); + }); + + 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); + 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; + + await expect(postSpawn( + cfg, + { model: "team/gpt-5.6-sol", input: readableAgentInput(), stream: false }, + { + resolveCodexModelEntitlements: async () => { + throw new TypeError("programmer sentinel"); + }, + }, + )).rejects.toThrow("programmer sentinel"); + expect(fetchCalls).toBe(0); + }); + test("unentitled fixed gated primary falls through to a routed fallback", async () => { const now = 1_800_000_000_000; Date.now = () => now; @@ -917,12 +1051,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 +1082,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 +1539,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) ?? []; @@ -1303,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); }); @@ -1457,6 +1707,218 @@ 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("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; + 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 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("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({