From 76af70ca916e234755b891564924786def1ce835 Mon Sep 17 00:00:00 2001 From: luvs01 <27862058+luvs01@users.noreply.github.com> Date: Sun, 23 Aug 2026 05:58:14 +0900 Subject: [PATCH] fix(codex): honor scoped cooldowns in subagent fallback --- src/codex/auth-context.ts | 59 +- src/codex/model-entitlements.ts | 26 +- src/codex/routing.ts | 44 + src/codex/subagent-model-fallback.ts | 330 ++++- src/server/responses/core.ts | 328 ++++- tests/codex-model-entitlements.test.ts | 51 + ...subagent-fallback-handle-responses.test.ts | 1168 ++++++++++++++++- tests/subagent-model-fallback.test.ts | 149 +++ 8 files changed, 2044 insertions(+), 111 deletions(-) diff --git a/src/codex/auth-context.ts b/src/codex/auth-context.ts index 7dd58c6b91..58891b8cb0 100644 --- a/src/codex/auth-context.ts +++ b/src/codex/auth-context.ts @@ -18,6 +18,7 @@ import type { NativeMainStartupBlockReason } from "./native-profile-startup"; import { codexQuotaScopeForModel, getCodexQuotaHealthSnapshot, + isCodexQuotaProbeReservationActive, releaseCodexQuotaProbeLease, releaseCodexQuotaScopeProbeLease, tryAcquireCodexQuotaProbeLease, @@ -29,10 +30,13 @@ import { entitledCodexAccountIdsForModel, isDirectCallerEntitledToCodexModel, resolveCodexModelEntitlements, - type CodexModelEntitlementSnapshot, } from "./model-entitlements"; import { ACCOUNT_GATED_NATIVE_OPENAI_MODELS } from "./catalog/native-models"; -import type { CodexCooldownSource, CodexQuotaScope } from "./routing"; +import type { + CodexCooldownSource, + CodexQuotaProbeReservation, + CodexQuotaScope, +} from "./routing"; import { maskAccountId } from "../lib/privacy"; import { formatErrorResponse } from "../bridge"; import { getAccountQuota } from "./quota"; @@ -327,6 +331,8 @@ export interface ResolveCodexAuthContextOptions { accountId?: string; /** Final native model selected for this request, used to select its quota group. */ modelId?: string; + /** Atomic quota-probe ownership reserved by subagent fallback selection. */ + preAcquiredQuotaProbeReservation?: CodexQuotaProbeReservation; /** Short reservation converted to turn ownership before native `__main__` token materialization. */ beginCodexAccountSelection?: () => CodexAccountSelectionAdmission | undefined; /** Test-only native credential read seams. */ @@ -334,9 +340,7 @@ export interface ResolveCodexAuthContextOptions { getMainAccountToken?: typeof getMainAccountToken; primeCodexPoolQuotas?: (config: OcxConfig, reason: string) => Promise; /** Test seam for account-gated native model discovery. */ - resolveCodexModelEntitlements?: ( - config: Pick, - ) => Promise; + resolveCodexModelEntitlements?: typeof resolveCodexModelEntitlements; /** Direct requests admitted with a proxy bearer substitute the stored native-main credential. */ substituteMainCredentialForDirect?: boolean; /** Test seam for a Direct request's own forwarded ChatGPT credential. */ @@ -381,28 +385,34 @@ export async function resolveCodexAuthContext( return { kind: "main", accountId: null }; } const affinityKey = fixedAccountId === undefined ? codexPoolAffinityKey(headers) : undefined; - const entitlementSnapshot = options.modelId && ACCOUNT_GATED_NATIVE_OPENAI_MODELS.has(options.modelId) - ? await (options.resolveCodexModelEntitlements ?? resolveCodexModelEntitlements)(config) - : undefined; - const modelEligibleAccountIds = entitlementSnapshot - ? entitledCodexAccountIdsForModel(entitlementSnapshot, options.modelId) - : undefined; // Retained startup recovery makes the physical main identity ineligible. Routing // can still preserve service by selecting a healthy configured pool account. const nativeMainTrafficBlocked = isNativeMainTrafficBlocked(); const selectionAdmission = options.beginCodexAccountSelection?.(); const nativeMainReadsForbidden = nativeMainTrafficBlocked || selectionAdmission?.mainProfileDraining === true; - 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, - isMainAccountTokenLive: options.isMainAccountTokenLive, - modelEligibleAccountIds, - }; let accountId: string; const quotaScope = codexQuotaScopeForModel(options.modelId); try { + const excludeAccountIds = nativeMainReadsForbidden + ? new Set([MAIN_CODEX_ACCOUNT_ID]) + : undefined; + const entitlementSnapshot = options.modelId && ACCOUNT_GATED_NATIVE_OPENAI_MODELS.has(options.modelId) + ? await (options.resolveCodexModelEntitlements ?? resolveCodexModelEntitlements)(config, { excludeAccountIds }) + : undefined; + const entitledAccountIds = entitlementSnapshot + ? entitledCodexAccountIdsForModel(entitlementSnapshot, options.modelId) + : undefined; + const modelEligibleAccountIds = entitledAccountIds + ? new Set([...entitledAccountIds].filter(candidate => !excludeAccountIds?.has(candidate))) + : undefined; + 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, + isMainAccountTokenLive: options.isMainAccountTokenLive, + modelEligibleAccountIds, + }; // A pre-drain selector reserves the native identity while reconciliation and // routing inspect it. Selectors arriving after the fence skip reconciliation // and may still route to non-main pool accounts without touching switch state. @@ -511,9 +521,16 @@ export async function resolveCodexAuthContext( throw new CodexAccountCooldownError(accountId, cooldownUntil, cooldown?.cooldownSource, cooldown?.quotaScope); } probeQuotaScope = cooldown?.quotaScope; - probeLeaseId = probeQuotaScope + const preAcquiredProbe = options.preAcquiredQuotaProbeReservation; + const matchingPreAcquiredProbe = preAcquiredProbe + && preAcquiredProbe.accountId === accountId + && preAcquiredProbe.quotaScope === probeQuotaScope + && isCodexQuotaProbeReservationActive(preAcquiredProbe) + ? preAcquiredProbe + : undefined; + probeLeaseId = matchingPreAcquiredProbe?.leaseId ?? (probeQuotaScope ? tryAcquireCodexQuotaScopeProbeLease(accountId, probeQuotaScope) ?? undefined - : tryAcquireCodexQuotaProbeLease(accountId) ?? undefined; + : tryAcquireCodexQuotaProbeLease(accountId) ?? undefined); if (!probeLeaseId) { throw new CodexAccountCooldownError(accountId, cooldownUntil, cooldown?.cooldownSource, cooldown?.quotaScope); } diff --git a/src/codex/model-entitlements.ts b/src/codex/model-entitlements.ts index d06100138b..fec88316ba 100644 --- a/src/codex/model-entitlements.ts +++ b/src/codex/model-entitlements.ts @@ -40,6 +40,12 @@ export interface CodexModelEntitlementResolveOptions { readonly now?: number; /** Test-only credential seam; production callers enumerate local main + Pool credentials. */ readonly credentials?: readonly CodexModelEntitlementCredentialSnapshot[]; + /** Test-only seam for proving lifecycle exclusions happen before credential reads. */ + readonly credentialSnapshot?: typeof accountCredentialSnapshot; + /** Accounts whose credentials must not be read while another lifecycle owns them. */ + readonly excludeAccountIds?: ReadonlySet; + /** Last-mile circuit fence checked synchronously before a new roster fetch starts. */ + readonly allowNetworkFetch?: () => boolean; } const accountModelsCache = new Map(); @@ -195,6 +201,7 @@ async function modelsForCredential( credential: CodexModelEntitlementCredentialSnapshot, fetcher: typeof fetch, now: number, + allowNetworkFetch?: () => boolean, ): Promise { const cached = accountModelsCache.get(credential.accountId); if ( @@ -204,6 +211,16 @@ async function modelsForCredential( ) return cached; const flightKey = `${credential.accountId}\u0000${credential.credentialIdentity}`; + if (allowNetworkFetch?.() === false) { + // Do not cache a circuit-fenced miss. The host circuit owns retry pacing; + // a local failure TTL would hide a grant after the circuit recovers. + return { + credentialIdentity: credential.credentialIdentity, + expiresAt: now, + models: new Set(), + confirmed: false, + }; + } const existing = accountModelsFlights.get(flightKey); if (existing) return existing; const flight = fetchAccountModels(credential, fetcher, now) @@ -252,13 +269,16 @@ export async function resolveCodexModelEntitlements( ): Promise { const now = options.now ?? Date.now(); const fetcher = options.fetcher ?? fetch; + const allowedAccountIds = candidateAccountIds(config) + .filter(accountId => !options.excludeAccountIds?.has(accountId)); + const credentialSnapshot = options.credentialSnapshot ?? accountCredentialSnapshot; const credentials = options.credentials - ? [...options.credentials] - : (await Promise.all(candidateAccountIds(config).map(accountCredentialSnapshot))) + ? [...options.credentials].filter(credential => !options.excludeAccountIds?.has(credential.accountId)) + : (await Promise.all(allowedAccountIds.map(credentialSnapshot))) .filter((value): value is CodexModelEntitlementCredentialSnapshot => value !== null); const results = await Promise.all(credentials.map(async credential => ({ credential, - result: await modelsForCredential(credential, fetcher, now), + result: await modelsForCredential(credential, fetcher, now, options.allowNetworkFetch), }))); return { modelsByAccount: new Map(results.map(({ credential, result }) => [credential.accountId, result.models])), diff --git a/src/codex/routing.ts b/src/codex/routing.ts index fa6fa63d83..555dc7e154 100644 --- a/src/codex/routing.ts +++ b/src/codex/routing.ts @@ -138,6 +138,13 @@ export type CodexCooldownSource = "retry-after" | "reset-derived" | "default"; */ export type CodexQuotaScope = "shared" | "spark"; +/** Probe ownership reserved before final Codex authentication. */ +export type CodexQuotaProbeReservation = Readonly<{ + accountId: string; + leaseId: string; + quotaScope?: CodexQuotaScope; +}>; + export type CodexQuotaRecoveryProbeClaim = { accountId: string; scope?: CodexQuotaScope; @@ -612,6 +619,15 @@ export function tryAcquireCodexQuotaScopeProbeLease( return probeLeaseId; } +/** Side-effect-free check for a confirmed model-specific quota probe. */ +export function canAcquireCodexQuotaScopeProbeLease( + accountId: string, + scope: CodexQuotaScope, + now = Date.now(), +): boolean { + return canAcquireQuotaProbeLease(scopedHealthFor(accountId, scope), now); +} + /** * Hand a probe lease back without recording an upstream outcome. Used by paths * that take a lease and then fail before any request reaches upstream. @@ -634,6 +650,34 @@ export function releaseCodexQuotaScopeProbeLease( setScopedHealth(accountId, scope, withProbeLeaseReleased(health, now)); } +/** Verify that a pre-auth reservation still owns the current cooldown generation. */ +export function isCodexQuotaProbeReservationActive( + reservation: CodexQuotaProbeReservation, +): boolean { + const health = reservation.quotaScope + ? scopedHealthFor(reservation.accountId, reservation.quotaScope) + : upstreamHealth.get(reservation.accountId); + return health?.probeLeaseId === reservation.leaseId + && (health.probeLeaseGeneration ?? 0) === (health.cooldownGeneration ?? 0); +} + +/** Return a pre-auth reservation that never reached upstream. */ +export function releaseCodexQuotaProbeReservation( + reservation: CodexQuotaProbeReservation, + now = Date.now(), +): void { + if (reservation.quotaScope) { + releaseCodexQuotaScopeProbeLease( + reservation.accountId, + reservation.quotaScope, + reservation.leaseId, + now, + ); + } else { + releaseCodexQuotaProbeLease(reservation.accountId, reservation.leaseId, now); + } +} + /** * True when this outcome belongs to the account's in-flight probe. The * undefined-id guard matters: without it an outcome carrying no lease would match diff --git a/src/codex/subagent-model-fallback.ts b/src/codex/subagent-model-fallback.ts index da40ebffff..afdb0f5acd 100644 --- a/src/codex/subagent-model-fallback.ts +++ b/src/codex/subagent-model-fallback.ts @@ -15,12 +15,14 @@ import { CODEX_HOME, getCodexHome } from "./paths"; import { CODEX_UNKNOWN_USAGE_SCORE, getAccountQuota } from "./quota"; import { canAcquireCodexQuotaProbeLease, + canAcquireCodexQuotaScopeProbeLease, codexQuotaScopeForModel, computeCodexUsageScore, getCodexQuotaHealthSnapshot, getEffectiveActiveCodexAccountId, getPoolAccountPlan, - isCodexAccountInCooldown, + type CodexQuotaProbeReservation, + type CodexQuotaScope, } from "./routing"; import { isCodexAccountUsable, @@ -38,6 +40,7 @@ import { 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 { getUpstreamHostHealth, normalizeUpstreamHostCircuitThreshold, @@ -48,6 +51,29 @@ export const DEFAULT_SUBAGENT_MODEL_FALLBACK_POLL_MS = 60_000; const CODEX_FORWARD_ORIGIN = new URL(CODEX_FORWARD_BASE_URL).origin.toLowerCase(); type SubagentQuotaPrimeFn = (config: OcxConfig, reason: string) => Promise; +/** Side-effect-free pool account preview for one resolved candidate model. */ +export type SubagentPoolAccountPreview = ( + modelId: string | undefined, + now: number, + modelEligibleAccountIds?: ReadonlySet, +) => string | null; +export type SubagentModelEligibleAccountIds = ( + modelId: string | undefined, +) => ReadonlySet | undefined; +export type SubagentQuotaProbeReserver = ( + accountId: string, + quotaScope: CodexQuotaScope | undefined, + now: number, +) => CodexQuotaProbeReservation | null; +export type AppliedSubagentModelFallback = { + from?: string; + to?: string; + skipped?: string[]; + quotaProbeReservation?: CodexQuotaProbeReservation; +} | null; +export type SubagentFallbackEntitlementBoundaryResult = + | { kind: "complete"; fallback: AppliedSubagentModelFallback } + | { kind: "entitlements-required" }; let subagentQuotaPrimeForTests: SubagentQuotaPrimeFn | null = null; let quotaPrimeInFlight: Promise | null = null; @@ -177,8 +203,15 @@ function resolveRouteFallbackAccountId( route: RouteResult | null, config: OcxConfig, accountId?: string | null, + now = Date.now(), + poolAccountPreview?: SubagentPoolAccountPreview, + modelEligibleAccountIds?: ReadonlySet, ): string | null { - return route?.codexAccountId ?? resolvePoolFallbackAccountId(config, accountId); + if (route?.codexAccountId !== undefined) return route.codexAccountId; + if (route && isPoolCodexRoute(route) && poolAccountPreview) { + return poolAccountPreview(route.modelId, now, modelEligibleAccountIds); + } + return resolvePoolFallbackAccountId(config, accountId); } function isRoutableFallbackModel(model: string, config: OcxConfig): boolean { @@ -231,39 +264,76 @@ export function isModelHealthBlocked( return !!health && health.unavailableUntil > now; } +/** + * Check one fallback candidate against the pool account its resolved model scope would use. + * The optional preview must not bind affinity, move an account cursor, or acquire a probe lease. + */ export function isSubagentModelUnavailable( model: string, config: OcxConfig, accountId?: string | null, now = Date.now(), accountUsabilityOptions?: CodexAccountUsabilityOptions, + poolAccountPreview?: SubagentPoolAccountPreview, + modelEligibleAccountIdsForModel?: SubagentModelEligibleAccountIds, + quotaProbeReserver?: SubagentQuotaProbeReserver, + onQuotaProbeReserved?: (reservation: CodexQuotaProbeReservation) => void, ): boolean { if (isDisabledFallbackModel(model, config)) return true; if (!isRoutableFallbackModel(model, config)) return true; const route = tryRouteFallbackModel(config, model); if (!route || route.provider.disabled === true) return true; - if (isModelHealthBlocked(model, config, accountId, now)) return true; + const candidateAccountUsabilityOptions = modelEligibleAccountIdsForModel + ? { + ...accountUsabilityOptions, + modelEligibleAccountIds: modelEligibleAccountIdsForModel(route.modelId), + } + : accountUsabilityOptions; + const resolvedAccountId = resolveRouteFallbackAccountId( + route, + config, + accountId, + now, + 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). - const resolvedAccountId = resolveRouteFallbackAccountId(route, config, accountId); if (!resolvedAccountId) return true; if (isCodexAccountPaused(config, resolvedAccountId)) return true; - if (!isCodexAccountUsable(config, resolvedAccountId, accountUsabilityOptions)) 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 if ( - isCodexAccountInCooldown(resolvedAccountId, now) - && !canAcquireCodexQuotaProbeLease(resolvedAccountId, now) - ) { - return true; + } else { + if (isNativeModelQuotaExhausted(model, config, resolvedAccountId, now)) return true; + const quotaScope = codexQuotaScopeForModel(route.modelId); + const cooldown = getCodexQuotaHealthSnapshot(resolvedAccountId, quotaScope, now); + if (cooldown !== null) { + if (quotaProbeReserver) { + const reservation = quotaProbeReserver( + resolvedAccountId, + cooldown.quotaScope, + now, + ); + if (!reservation) return true; + onQuotaProbeReserved?.(reservation); + } else { + const probeAvailable = cooldown.quotaScope + ? canAcquireCodexQuotaScopeProbeLease(resolvedAccountId, cooldown.quotaScope, now) + : canAcquireCodexQuotaProbeLease(resolvedAccountId, now); + if (!probeAvailable) return true; + } + } + return false; } - return isNativeModelQuotaExhausted(model, config, accountId, now); + return isNativeModelQuotaExhausted(model, config, resolvedAccountId, now); } export function selectAvailableSubagentModel( @@ -275,10 +345,20 @@ export function selectAvailableSubagentModel( nativeFallbackOnly = false, accountUsabilityOptions?: CodexAccountUsabilityOptions, trailingFallback: readonly string[] = [], -): { model: string; rewritten: boolean; skipped: string[] } { - const chain = normalizedChain(primary, config, extraFallback, trailingFallback); + poolAccountPreview?: SubagentPoolAccountPreview, + modelEligibleAccountIdsForModel?: SubagentModelEligibleAccountIds, + resolvedChain?: readonly string[], + quotaProbeReserver?: SubagentQuotaProbeReserver, +): { + model: string; + rewritten: boolean; + skipped: string[]; + quotaProbeReservation?: CodexQuotaProbeReservation; +} { + const chain = resolvedChain ?? normalizedChain(primary, config, extraFallback, trailingFallback); const skipped: string[] = []; for (const candidate of chain) { + let quotaProbeReservation: CodexQuotaProbeReservation | undefined; if (nativeFallbackOnly) { const route = tryRouteFallbackModel(config, candidate); if (!route || !isCanonicalOpenAiForwardProvider(route.provider)) { @@ -286,15 +366,56 @@ export function selectAvailableSubagentModel( continue; } } - if (isSubagentModelUnavailable(candidate, config, accountId, now, accountUsabilityOptions)) { + if (isSubagentModelUnavailable( + candidate, + config, + accountId, + now, + accountUsabilityOptions, + poolAccountPreview, + modelEligibleAccountIdsForModel, + quotaProbeReserver, + (reservation) => { quotaProbeReservation = reservation; }, + )) { skipped.push(candidate); continue; } - return { model: candidate, rewritten: !slugsEquivalent(candidate, primary), skipped }; + return { + model: candidate, + rewritten: !slugsEquivalent(candidate, primary), + skipped, + ...(quotaProbeReservation ? { quotaProbeReservation } : {}), + }; } return { model: primary, rewritten: false, skipped }; } +function applySelectedSubagentModelFallback( + parsed: OcxParsedRequest, + selection: ReturnType, +): AppliedSubagentModelFallback { + if (!selection.rewritten) return selection.skipped.length > 0 || selection.quotaProbeReservation + ? { + from: parsed.modelId, + to: parsed.modelId, + skipped: selection.skipped, + ...(selection.quotaProbeReservation + ? { quotaProbeReservation: selection.quotaProbeReservation } + : {}), + } + : null; + const from = parsed.modelId; + rewriteParsedModel(parsed, selection.model); + return { + from, + to: selection.model, + skipped: selection.skipped, + ...(selection.quotaProbeReservation + ? { quotaProbeReservation: selection.quotaProbeReservation } + : {}), + }; +} + export function noteSubagentModelFailure( model: string, message: string, @@ -418,10 +539,27 @@ export function resolveAgentModelFallbackForPrimary( return merged; } -function subagentQuotaPrimeBlockedByHostCircuit(config: OcxConfig): boolean { - if (normalizeUpstreamHostCircuitThreshold(config.upstreamHostCircuitThreshold) === 0) return false; +/** + * Whether subagent-side ChatGPT reads must defer to the logical request that + * owns the host circuit admission. An expired cooldown still counts here: the + * next request must acquire the single half-open probe before any credential- + * bearing quota or entitlement discovery reaches the same origin. + */ +export type SubagentChatGptHostCircuitState = "closed" | "cooldown" | "probe-due"; + +export function subagentChatGptHostCircuitState( + config: OcxConfig, + now = Date.now(), +): SubagentChatGptHostCircuitState { + if (normalizeUpstreamHostCircuitThreshold(config.upstreamHostCircuitThreshold) === 0) return "closed"; const key = upstreamHostHealthKey(OPENAI_CODEX_PROVIDER_ID, CODEX_FORWARD_ORIGIN); - return getUpstreamHostHealth(key)?.cooldownUntil !== undefined; + const cooldownUntil = getUpstreamHostHealth(key)?.cooldownUntil; + if (cooldownUntil === undefined) return "closed"; + return cooldownUntil > now ? "cooldown" : "probe-due"; +} + +export function isSubagentChatGptHostCircuitBlocked(config: OcxConfig): boolean { + return subagentChatGptHostCircuitState(config) !== "closed"; } /** @@ -469,7 +607,7 @@ export function maybePrimeSubagentQuota( options: { nativeMainReadsForbidden?: boolean } = {}, ): Promise { if (options.nativeMainReadsForbidden) return Promise.resolve(); - if (subagentQuotaPrimeBlockedByHostCircuit(config)) return Promise.resolve(); + if (isSubagentChatGptHostCircuitBlocked(config)) return Promise.resolve(); if (quotaPrimeInFlight) return quotaPrimeInFlight; if (!shouldPrimeSubagentQuota(config, now)) return Promise.resolve(); @@ -477,7 +615,7 @@ export function maybePrimeSubagentQuota( try { // Re-check after claiming single-flight ownership so a circuit opened by // a concurrent request cannot race us into a fresh usage-probe pass. - if (subagentQuotaPrimeBlockedByHostCircuit(config)) return; + if (isSubagentChatGptHostCircuitBlocked(config)) return; if (subagentQuotaPrimeForTests) { await subagentQuotaPrimeForTests(config, "subagent-spawn"); } else { @@ -515,34 +653,152 @@ export function applySubagentModelFallback( now = Date.now(), nativeFallbackOnly = false, accountUsabilityOptions?: CodexAccountUsabilityOptions, -): { from?: string; to?: string; skipped?: string[] } | null { + poolAccountPreview?: SubagentPoolAccountPreview, + modelEligibleAccountIdsForModel?: SubagentModelEligibleAccountIds, + resolvedFallbackChain?: readonly string[] | null, + quotaProbeReserver?: SubagentQuotaProbeReserver, +): AppliedSubagentModelFallback { if (!isThreadSpawnRequest(headers)) return null; - const tomlRoleFallback = resolveAgentModelFallbackForPrimary( + const fallbackChain = resolvedFallbackChain === undefined + ? resolveSubagentFallbackChain(parsed, config) + : resolvedFallbackChain; + if (!fallbackChain) return null; + const selection = selectAvailableSubagentModel( parsed.modelId, - getCodexHome(), - config.codexAccountNamespaces, + config, + [], + accountId, + now, + nativeFallbackOnly, + accountUsabilityOptions, + [], + poolAccountPreview, + modelEligibleAccountIdsForModel, + fallbackChain, + quotaProbeReserver, ); - // Config-keyed chains are the supported per-role home (#1190); TOML `model_fallback` - // stays readable for backwards compatibility with homes written before Codex 0.146. - const configuredFallback = resolveConfiguredModelFallbackForPrimary(parsed.modelId, config); - const globalFallback = config.subagentModelFallback ?? []; - if (globalFallback.length === 0 && configuredFallback.length === 0 && tomlRoleFallback.length === 0) return null; + return applySelectedSubagentModelFallback(parsed, selection); +} + +function subagentFallbackEntitlementBoundaryIndex( + fallbackChain: readonly string[] | null, + config: OcxConfig, +): number { + if (!fallbackChain) return -1; + return fallbackChain.findIndex((candidate) => { + const route = tryRouteFallbackModel(config, candidate); + return !!route + && isPoolCodexRoute(route) + && ACCOUNT_GATED_NATIVE_OPENAI_MODELS.has(route.modelId); + }); +} + +/** + * Apply fallback only through the candidates that precede the first + * account-gated Pool model. Reaching that boundary asks the caller to resolve + * entitlement evidence before restarting selection from the full chain. + */ +export function applySubagentModelFallbackUntilEntitlementBoundary( + parsed: OcxParsedRequest, + headers: Headers, + config: OcxConfig, + accountId: string | null | undefined, + now: number, + nativeFallbackOnly: boolean, + accountUsabilityOptions: CodexAccountUsabilityOptions | undefined, + poolAccountPreview: SubagentPoolAccountPreview, + resolvedFallbackChain: readonly string[] | null, + quotaProbeReserver?: SubagentQuotaProbeReserver, +): SubagentFallbackEntitlementBoundaryResult { + if (!isThreadSpawnRequest(headers) || !resolvedFallbackChain) { + return { + kind: "complete", + fallback: applySubagentModelFallback( + parsed, + headers, + config, + accountId, + now, + nativeFallbackOnly, + accountUsabilityOptions, + poolAccountPreview, + undefined, + resolvedFallbackChain, + quotaProbeReserver, + ), + }; + } + const entitlementBoundary = subagentFallbackEntitlementBoundaryIndex( + resolvedFallbackChain, + config, + ); + if (entitlementBoundary < 0) { + return { + kind: "complete", + fallback: applySubagentModelFallback( + parsed, + headers, + config, + accountId, + now, + nativeFallbackOnly, + accountUsabilityOptions, + poolAccountPreview, + undefined, + resolvedFallbackChain, + quotaProbeReserver, + ), + }; + } + + const prefix = resolvedFallbackChain.slice(0, entitlementBoundary); const selection = selectAvailableSubagentModel( parsed.modelId, config, - configuredFallback, + [], accountId, now, nativeFallbackOnly, accountUsabilityOptions, - tomlRoleFallback, + [], + poolAccountPreview, + undefined, + prefix, + quotaProbeReserver, ); - if (!selection.rewritten) return selection.skipped.length > 0 - ? { from: parsed.modelId, to: parsed.modelId, skipped: selection.skipped } - : null; - const from = parsed.modelId; - rewriteParsedModel(parsed, selection.model); - return { from, to: selection.model, skipped: selection.skipped }; + if (selection.skipped.length === prefix.length) { + return { kind: "entitlements-required" }; + } + return { + kind: "complete", + fallback: applySelectedSubagentModelFallback(parsed, selection), + }; +} + +/** Resolve the effective fallback chain once for one logical spawn request. */ +export function resolveSubagentFallbackChain( + parsed: OcxParsedRequest, + config: OcxConfig, +): readonly string[] | null { + const tomlRoleFallback = resolveAgentModelFallbackForPrimary( + parsed.modelId, + getCodexHome(), + config.codexAccountNamespaces, + ); + // Config-keyed chains are the supported per-role home (#1190); TOML `model_fallback` + // stays readable for backwards compatibility with homes written before Codex 0.146. + const configuredFallback = resolveConfiguredModelFallbackForPrimary(parsed.modelId, config); + const globalFallback = config.subagentModelFallback ?? []; + if (globalFallback.length === 0 && configuredFallback.length === 0 && tomlRoleFallback.length === 0) return null; + return normalizedChain(parsed.modelId, config, configuredFallback, tomlRoleFallback); +} + +/** Whether this request's configured fallback chain contains an account-gated native model. */ +export function subagentFallbackNeedsModelEntitlements( + fallbackChain: readonly string[] | null, + config: OcxConfig, +): boolean { + return subagentFallbackEntitlementBoundaryIndex(fallbackChain, config) >= 0; } export function subagentFallbackGuidanceText(config: OcxConfig): string { diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index 67005be579..fd45d94b2b 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -134,6 +134,7 @@ import { invalidateCodexModelEntitlementsForAccount, resolveCodexModelEntitlements, } from "../../codex/model-entitlements"; +import { MAIN_CODEX_ACCOUNT_ID } from "../../codex/main-account"; import { ACCOUNT_GATED_NATIVE_OPENAI_MODELS } from "../../codex/catalog/native-models"; import { captureCodexAffinityDiagnostic } from "../../codex/affinity-debug"; import { @@ -142,6 +143,10 @@ import { formatCodexProviderForLog, previewCodexAccountForRequest, recordCodexUpstreamOutcome, + releaseCodexQuotaProbeReservation, + tryAcquireCodexQuotaProbeLease, + tryAcquireCodexQuotaScopeProbeLease, + type CodexQuotaProbeReservation, type CodexUpstreamOutcome, } from "../../codex/routing"; import { codexAuthContextLogLabel } from "../../codex/account-label"; @@ -201,8 +206,15 @@ import { supportedLadderFor } from "../effort-policy"; import { isThreadSpawnRequest } from "../effort-policy"; import { applySubagentModelFallback, + applySubagentModelFallbackUntilEntitlementBoundary, + isSubagentChatGptHostCircuitBlocked, maybePrimeSubagentQuota, recordSubagentQuotaFailureForThreadSpawn, + resolveSubagentFallbackChain, + subagentChatGptHostCircuitState, + type SubagentModelEligibleAccountIds, + type SubagentPoolAccountPreview, + type SubagentQuotaProbeReserver, } from "../../codex/subagent-model-fallback"; import { isNativeMainTrafficBlocked } from "../../codex/native-profile-startup"; import { @@ -753,6 +765,188 @@ export function preAuthUpstreamHostCircuitKey( return upstreamHostHealthKey(route.providerName, safeOriginLabel(route.provider.baseUrl ?? "")); } +const NO_ELIGIBLE_CODEX_ACCOUNTS: ReadonlySet = new Set(); + +type SubagentFallbackModelEligibilityResolution = Readonly<{ + modelEligibleAccountIds: SubagentModelEligibleAccountIds; + preservePrimaryForHostProbe: boolean; + preAuthHostAdmissionLease?: UpstreamHostAdmissionLease; +}>; + +/** + * Resolve account-gated fallback eligibility without bypassing ChatGPT host + * admission. While the circuit is open, fail closed only for gated native + * models so an unrelated routed primary (for example xAI) can still proceed. + */ +async function resolveSubagentFallbackModelEligibleAccountIds( + config: OcxConfig, + resolver: typeof resolveCodexModelEntitlements, + primaryRoute: RouteResult, + excludeAccountIds?: ReadonlySet, +): Promise { + const circuitFencedResolution = ( + hostCircuitState: Exclude, "closed">, + ): SubagentFallbackModelEligibilityResolution => { + const preAuthHostKey = hostCircuitState === "probe-due" + && ACCOUNT_GATED_NATIVE_OPENAI_MODELS.has(primaryRoute.modelId) + ? preAuthUpstreamHostCircuitKey(primaryRoute, config) + : null; + if (preAuthHostKey) { + // Reserve the single half-open probe atomically with the preservation + // decision. A read-only health snapshot cannot distinguish another + // request acquiring the lease between selection and final admission. + const admission = acquireUpstreamHostAdmission( + preAuthHostKey, + config.upstreamHostCircuitThreshold, + ); + if (admission.kind === "admitted" && admission.lease) { + return { + preservePrimaryForHostProbe: true, + preAuthHostAdmissionLease: admission.lease, + modelEligibleAccountIds: (modelId) => modelId + && ACCOUNT_GATED_NATIVE_OPENAI_MODELS.has(modelId) + ? NO_ELIGIBLE_CODEX_ACCOUNTS + : undefined, + }; + } + } + return { + preservePrimaryForHostProbe: false, + modelEligibleAccountIds: (modelId) => modelId + && ACCOUNT_GATED_NATIVE_OPENAI_MODELS.has(modelId) + ? NO_ELIGIBLE_CODEX_ACCOUNTS + : undefined, + }; + }; + const hostCircuitState = subagentChatGptHostCircuitState(config); + if (hostCircuitState !== "closed") { + return circuitFencedResolution(hostCircuitState); + } + const entitlementSnapshot = await resolver(config, { + excludeAccountIds, + // Credential snapshots may await disk/token work. Re-check synchronously + // inside the resolver immediately before every new roster fetch so a + // circuit opened during that gap still fences network I/O. + allowNetworkFetch: () => !isSubagentChatGptHostCircuitBlocked(config), + }); + const settledHostCircuitState = subagentChatGptHostCircuitState(config); + if (settledHostCircuitState !== "closed") { + return circuitFencedResolution(settledHostCircuitState); + } + return { + preservePrimaryForHostProbe: false, + modelEligibleAccountIds: (modelId) => { + const entitledAccountIds = entitledCodexAccountIdsForModel(entitlementSnapshot, modelId); + return entitledAccountIds + ? new Set([...entitledAccountIds].filter(accountId => !excludeAccountIds?.has(accountId))) + : undefined; + }, + }; +} + +type LazySubagentFallbackApplication = Readonly<{ + fallback: ReturnType; + previewAccountId: string | null; + preAuthHostAdmissionLease?: UpstreamHostAdmissionLease; + preAuthQuotaProbeReservation?: CodexQuotaProbeReservation; +}>; + +async function applySubagentModelFallbackLazily(options: { + parsed: OcxParsedRequest; + headers: Headers; + config: OcxConfig; + route: RouteResult; + poolAffinityKey: string | null; + selectionOptions: { nativeMainSelectionOnly: boolean }; + fallbackChain: readonly string[] | null; + nativeMainReadsForbidden: boolean; + nativeFallbackOnly: boolean; + resolver: typeof resolveCodexModelEntitlements; +}): Promise { + const accountPreview: SubagentPoolAccountPreview = (modelId, previewNow, eligibleAccountIds) => + previewCodexAccountForRequest( + options.poolAffinityKey, + options.config, + previewNow, + codexQuotaScopeForModel(modelId), + { ...options.selectionOptions, modelEligibleAccountIds: eligibleAccountIds }, + ); + const reserveQuotaProbe: SubagentQuotaProbeReserver = (accountId, quotaScope, reserveNow) => { + const leaseId = quotaScope + ? tryAcquireCodexQuotaScopeProbeLease(accountId, quotaScope, reserveNow) + : tryAcquireCodexQuotaProbeLease(accountId, reserveNow); + return leaseId + ? { accountId, leaseId, ...(quotaScope ? { quotaScope } : {}) } + : null; + }; + const preselectionNow = Date.now(); + const previewAccountId = options.route.codexAccountId ?? accountPreview( + options.route.modelId, + preselectionNow, + ); + const preselection = applySubagentModelFallbackUntilEntitlementBoundary( + options.parsed, + options.headers, + options.config, + previewAccountId, + preselectionNow, + options.nativeFallbackOnly, + options.selectionOptions, + accountPreview, + options.fallbackChain, + reserveQuotaProbe, + ); + if (preselection.kind === "complete") { + return { + fallback: preselection.fallback, + previewAccountId, + preAuthQuotaProbeReservation: preselection.fallback?.quotaProbeReservation, + }; + } + + const excludeAccountIds = options.nativeMainReadsForbidden + ? new Set([MAIN_CODEX_ACCOUNT_ID]) + : undefined; + const entitlementResolution = await resolveSubagentFallbackModelEligibleAccountIds( + options.config, + options.resolver, + options.route, + excludeAccountIds, + ); + // Entitlement discovery is an async boundary. Restart the full selection with + // a fresh clock so cooldown/health changes during discovery cannot make the + // preselection authoritative. + const settledNow = Date.now(); + const entitledPreviewAccountId = options.route.codexAccountId ?? accountPreview( + options.route.modelId, + settledNow, + entitlementResolution.preservePrimaryForHostProbe + ? undefined + : entitlementResolution.modelEligibleAccountIds(options.route.modelId), + ); + const fallback = entitlementResolution.preservePrimaryForHostProbe + ? null + : applySubagentModelFallback( + options.parsed, + options.headers, + options.config, + entitledPreviewAccountId, + settledNow, + options.nativeFallbackOnly, + options.selectionOptions, + accountPreview, + entitlementResolution.modelEligibleAccountIds, + options.fallbackChain, + reserveQuotaProbe, + ); + return { + previewAccountId: entitledPreviewAccountId, + preAuthHostAdmissionLease: entitlementResolution.preAuthHostAdmissionLease, + fallback, + preAuthQuotaProbeReservation: fallback?.quotaProbeReservation, + }; +} + export function upstreamHostCircuitOpenResponse(retryAfterSeconds: number): Response { return formatErrorResponse( 503, @@ -1227,6 +1421,8 @@ export interface HandleResponsesOptions { /** One-shot TTFT callback: first non-empty model output observed (WP4). */ onFirstOutput?: () => void; onCodexAuthContextResolved?: (context: CodexAuthContext | undefined) => void; + /** Internal deterministic seam for account-gated native fallback tests. */ + resolveCodexModelEntitlements?: typeof resolveCodexModelEntitlements; recordTerminalOutcomes?: boolean; setTerminalOutcomeRecorder?: (recorder: ((status: ResponsesTerminalStatus, httpStatusOverride?: number) => void) | undefined) => void; onNativePassthroughTerminal?: (status: ResponsesTerminalStatus) => void; @@ -1483,6 +1679,7 @@ async function resolveResponsesCodexAuth( config: OcxConfig, route: RouteResult, options: HandleResponsesOptions, + preAcquiredQuotaProbeReservation?: CodexQuotaProbeReservation, ): Promise { try { // #1686: a caller that proved admission with a BEARER presented one of our own secrets. @@ -1516,8 +1713,10 @@ async function resolveResponsesCodexAuth( authCtx = await resolveCodexAuthContext(req.headers, config, route.codexAccountMode, { accountId: route.codexAccountId, modelId: route.modelId, + preAcquiredQuotaProbeReservation, substituteMainCredentialForDirect: substituteMainCredential, beginCodexAccountSelection: codexAccountSelectionForTurn(options.turnAdmissionLease), + resolveCodexModelEntitlements: options.resolveCodexModelEntitlements, }); options.onCodexAuthContextResolved?.(authCtx); } else { @@ -2152,6 +2351,7 @@ async function handleResponsesInner( options: HandleResponsesOptions & { translatorBudget: TranslatorBudget }, ): Promise { let pendingHostAdmissionLease: UpstreamHostAdmissionLease | null = null; + let pendingQuotaProbeReservation: CodexQuotaProbeReservation | null = null; let authCtx: CodexAuthContext = { kind: "main", accountId: null }; try { // The Chat and Anthropic surfaces replay through here with a Responses-shaped body, @@ -2352,7 +2552,7 @@ async function handleResponsesInner( // also fail closed without polling quota upstream. Cached fallback state can still select a // provider with native continuation support below. const threadSpawn = isThreadSpawnRequest(req.headers); - const previewSelectionAdmission = threadSpawn && route.codexAccountId === undefined + const previewSelectionAdmission = threadSpawn && !options.comboAttempt ? codexAccountSelectionForTurn(options.turnAdmissionLease)?.() : undefined; const nativeMainRecoveryBlocked = isNativeMainTrafficBlocked(); @@ -2364,12 +2564,16 @@ async function handleResponsesInner( }; let selectedForwardHeaders = req.headers; let subagentFallbackAccountId = config.activeCodexAccountId ?? null; - let subagentFallbackPreviewAccountId: string | null | undefined; + let subagentFallbackChain: readonly string[] | null = null; + let subagentFallbackApplication: LazySubagentFallbackApplication | undefined; let subagentQuotaFailureModel = parsed.modelId; const parentThreadId = req.headers.get("x-codex-parent-thread-id")?.trim() ?? null; const poolAffinityKey = codexPoolAffinityKey(req.headers) ?? null; try { + subagentFallbackChain = threadSpawn && !options.comboAttempt + ? resolveSubagentFallbackChain(parsed, config) + : null; if ( threadSpawn && route.codexAccountId === undefined @@ -2378,33 +2582,36 @@ async function handleResponsesInner( await maybePrimeSubagentQuota(config, Date.now(), { nativeMainReadsForbidden }); } + if (threadSpawn && !options.comboAttempt) { + subagentFallbackApplication = await applySubagentModelFallbackLazily({ + parsed, + headers: req.headers, + config, + route, + poolAffinityKey, + selectionOptions: previewSelectionOptions, + fallbackChain: subagentFallbackChain, + nativeMainReadsForbidden, + nativeFallbackOnly: unreadableEncryptedAgentTask, + resolver: options.resolveCodexModelEntitlements ?? resolveCodexModelEntitlements, + }); + pendingHostAdmissionLease = subagentFallbackApplication.preAuthHostAdmissionLease ?? null; + pendingQuotaProbeReservation = subagentFallbackApplication.preAuthQuotaProbeReservation ?? null; + } + // Subagent fallback must settle the final model/provider BEFORE route-dependent // normalization (virtual models, effort caps, service tier, wire protocol). // Preview the preferred Codex account without acquiring a probe lease or refreshing // tokens — auth is resolved only after the final route is selected. - if (threadSpawn && !options.comboAttempt && route.codexAccountId === undefined) { + if (threadSpawn && !options.comboAttempt) { + // An account-qualified primary remains pinned to its exact account. Later pooled + // fallback candidates still need a candidate-scoped preview, so construct the + // side-effect-free preview independently of the primary route ownership. // The final resolveCodexAuthContext binds under codexQuotaScopeForModel(route.modelId), - // so the preview must read the same scope slot — an undefined scope would map to the - // "legacy" affinity bucket and never find a binding made under "shared" or a native - // model scope, making the preview diverge from the account that actually authenticates. - const previewAccountId = previewCodexAccountForRequest( - poolAffinityKey, - config, - Date.now(), - codexQuotaScopeForModel(route.modelId), - previewSelectionOptions, - ); - subagentFallbackPreviewAccountId = previewAccountId; - subagentFallbackAccountId = previewAccountId ?? config.activeCodexAccountId ?? null; - const fallback = applySubagentModelFallback( - parsed, - req.headers, - config, - previewAccountId, - Date.now(), - unreadableEncryptedAgentTask, - previewSelectionOptions, - ); + // so each pooled preview must read the same scope slot. + const selection = subagentFallbackApplication!; + subagentFallbackAccountId = selection.previewAccountId ?? config.activeCodexAccountId ?? null; + const fallback = selection.fallback; if (fallback) { (logCtx as unknown as Record).subagentModelFallbackFrom = fallback.from; (logCtx as unknown as Record).subagentModelFallbackTo = fallback.to; @@ -2414,7 +2621,7 @@ async function handleResponsesInner( } subagentQuotaFailureModel = fallback?.to ?? parsed.modelId; - if (fallback?.to && !slugsEquivalent(fallback.to, route.modelId)) { + if (fallback?.to && fallback.from && !slugsEquivalent(fallback.to, fallback.from)) { try { route = routeModel(config, fallback.to, evidenceFromBody(parsed._rawBody)); logCtx.routeDecision = route.routeDecision; @@ -2486,15 +2693,42 @@ async function handleResponsesInner( // The ciphertext-only pass intentionally excludes routed candidates. Once recovery // makes the assignment readable, run selection again with the full configured chain // and keep the route in sync with any newly selected fallback. - const fallback = applySubagentModelFallback( - parsed, - req.headers, - config, - subagentFallbackPreviewAccountId, - Date.now(), - false, - previewSelectionOptions, - ); + const recoverySelectionAdmission = codexAccountSelectionForTurn(options.turnAdmissionLease)?.(); + let fallback: ReturnType; + try { + const recoveryNativeMainBlocked = isNativeMainTrafficBlocked(); + const recoveryNativeMainReadsForbidden = recoveryNativeMainBlocked + || recoverySelectionAdmission?.mainProfileDraining === true; + const recoverySelectionOptions = { + nativeMainSelectionOnly: !recoveryNativeMainBlocked + && recoverySelectionAdmission?.mainProfileDraining === true, + }; + const recoverySelection = await applySubagentModelFallbackLazily({ + parsed, + headers: req.headers, + config, + route, + poolAffinityKey, + selectionOptions: recoverySelectionOptions, + fallbackChain: subagentFallbackChain, + nativeMainReadsForbidden: recoveryNativeMainReadsForbidden, + nativeFallbackOnly: false, + resolver: options.resolveCodexModelEntitlements ?? resolveCodexModelEntitlements, + }); + pendingHostAdmissionLease = recoverySelection.preAuthHostAdmissionLease + ?? pendingHostAdmissionLease; + const recoveredProbeReservation = recoverySelection.preAuthQuotaProbeReservation ?? null; + if ( + pendingQuotaProbeReservation + && pendingQuotaProbeReservation.leaseId !== recoveredProbeReservation?.leaseId + ) { + releaseCodexQuotaProbeReservation(pendingQuotaProbeReservation); + } + pendingQuotaProbeReservation = recoveredProbeReservation; + fallback = recoverySelection.fallback; + } finally { + recoverySelectionAdmission?.release(); + } if (fallback) { (logCtx as unknown as Record).subagentModelFallbackFrom = fallback.from; (logCtx as unknown as Record).subagentModelFallbackTo = fallback.to; @@ -2504,7 +2738,7 @@ async function handleResponsesInner( } subagentQuotaFailureModel = fallback?.to ?? parsed.modelId; - if (fallback?.to && !slugsEquivalent(fallback.to, route.modelId)) { + if (fallback?.to && fallback.from && !slugsEquivalent(fallback.to, fallback.from)) { try { route = routeModel(config, fallback.to, evidenceFromBody(parsed._rawBody)); logCtx.routeDecision = route.routeDecision; @@ -2596,7 +2830,11 @@ async function handleResponsesInner( } } const preAuthHostKey = preAuthUpstreamHostCircuitKey(route, config); - if (preAuthHostKey) { + if (pendingHostAdmissionLease && pendingHostAdmissionLease.key !== preAuthHostKey) { + releaseUpstreamHostAdmission(pendingHostAdmissionLease); + pendingHostAdmissionLease = null; + } + if (preAuthHostKey && !pendingHostAdmissionLease) { const admission = acquireUpstreamHostAdmission( preAuthHostKey, config.upstreamHostCircuitThreshold, @@ -2609,9 +2847,22 @@ async function handleResponsesInner( let substituteMainCredential = false; { - const finalAuth = await resolveResponsesCodexAuth(req, config, route, options); + const finalAuth = await resolveResponsesCodexAuth( + req, + config, + route, + options, + pendingQuotaProbeReservation ?? undefined, + ); if (!finalAuth.ok) return finalAuth.response; authCtx = finalAuth.authCtx; + if ( + pendingQuotaProbeReservation + && pendingQuotaProbeReservation.leaseId !== codexProbeLeaseId(authCtx) + ) { + releaseCodexQuotaProbeReservation(pendingQuotaProbeReservation); + } + pendingQuotaProbeReservation = null; selectedForwardHeaders = finalAuth.headers; substituteMainCredential = finalAuth.substituteMainCredential; } @@ -5409,6 +5660,9 @@ async function handleResponsesInner( return formatErrorResponse(400, "invalid_request_error", "Non-streaming not supported by this adapter"); } finally { + if (pendingQuotaProbeReservation) { + releaseCodexQuotaProbeReservation(pendingQuotaProbeReservation); + } if (pendingHostAdmissionLease) { releaseUpstreamHostAdmission(pendingHostAdmissionLease); releaseCodexAuthContextProbeLease(authCtx); diff --git a/tests/codex-model-entitlements.test.ts b/tests/codex-model-entitlements.test.ts index 85a55c703e..5df625c737 100644 --- a/tests/codex-model-entitlements.test.ts +++ b/tests/codex-model-entitlements.test.ts @@ -61,6 +61,57 @@ describe("Codex account model entitlements", () => { expect(availableAccountGatedNativeModels(snapshot).size).toBe(0); }); + test("checks the network fence after credential work without caching the fenced miss", async () => { + let networkAllowed = false; + let fetchCalls = 0; + const options = { + credentials: [credential("main")], + allowNetworkFetch: () => networkAllowed, + fetcher: (async () => { + fetchCalls += 1; + return roster(DAYBREAK); + }) as typeof fetch, + now: 1_000, + }; + + const fenced = await resolveCodexModelEntitlements({ codexAccounts: [] }, options); + expect(fetchCalls).toBe(0); + expect(fenced.confirmedAccountIds.size).toBe(0); + + networkAllowed = true; + const recovered = await resolveCodexModelEntitlements({ codexAccounts: [] }, { + ...options, + now: 1_001, + }); + expect(fetchCalls).toBe(1); + expect(recovered.confirmedAccountIds.has("main")).toBe(true); + expect(entitledCodexAccountIdsForModel(recovered, DAYBREAK)?.has("main")).toBe(true); + }); + + test("does not read credentials excluded by a lifecycle owner", async () => { + const snapshotAccounts: string[] = []; + const seenAccounts: string[] = []; + const snapshot = await resolveCodexModelEntitlements({ + codexAccounts: [{ id: "secondary", email: "secondary@example.test", isMain: false }], + }, { + excludeAccountIds: new Set(["__main__"]), + credentialSnapshot: async accountId => { + snapshotAccounts.push(accountId); + return credential(accountId); + }, + fetcher: (async (_input, init) => { + seenAccounts.push(new Headers(init?.headers).get("chatgpt-account-id") ?? ""); + return roster(DAYBREAK); + }) as typeof fetch, + now: 1_000, + }); + + expect(snapshotAccounts).toEqual(["secondary"]); + expect(seenAccounts).toEqual(["chatgpt-secondary"]); + expect(snapshot.modelsByAccount.has("__main__")).toBe(false); + expect(snapshot.modelsByAccount.has("secondary")).toBe(true); + }); + test("ignores hidden or API-disabled rows", async () => { const snapshot = await resolveCodexModelEntitlements({ codexAccounts: [] }, { credentials: [credential("main")], diff --git a/tests/subagent-fallback-handle-responses.test.ts b/tests/subagent-fallback-handle-responses.test.ts index 937b0aadeb..63d26ac4b6 100644 --- a/tests/subagent-fallback-handle-responses.test.ts +++ b/tests/subagent-fallback-handle-responses.test.ts @@ -5,7 +5,7 @@ * 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 { mkdtempSync, rmSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { saveCodexAccountCredential } from "../src/codex/account-store"; @@ -14,25 +14,46 @@ import { updateAccountQuota, } from "../src/codex/quota"; import { + canAcquireCodexQuotaScopeProbeLease, CODEX_QUOTA_PROBE_INTERVAL_MS, clearCodexUpstreamHealth, clearThreadAccountMap, + getCodexQuotaHealthSnapshot, getCodexUpstreamHealth, previewCodexAccountForRequest, recordCodexUpstreamOutcome, + releaseCodexQuotaScopeProbeLease, resolveCodexAccountForThreadDetailed, + tryAcquireCodexQuotaScopeProbeLease, } from "../src/codex/routing"; import { + DEFAULT_SUBAGENT_MODEL_FALLBACK_POLL_MS, isModelHealthBlocked, resetSubagentModelFallbackStateForTests, setSubagentQuotaPrimeForTests, } from "../src/codex/subagent-model-fallback"; import { resolveCodexAuthContext, type CodexAuthContext } from "../src/codex/auth-context"; +import { resolveCodexModelEntitlements } from "../src/codex/model-entitlements"; +import { + acquireUpstreamHostAdmission, + clearUpstreamHostHealth, + getUpstreamHostHealth, + recordUpstreamHostFailure, + releaseUpstreamHostAdmission, + upstreamHostHealthKey, +} from "../src/codex/upstream-host-health"; import { handleResponses } from "../src/server/responses"; +import { resetAgentTaskRecoveryState } from "../src/server/responses/agent-task-recovery"; import { isEagerRelaySseResponse } from "../src/server/relay"; +import type { ActiveTurnLease } from "../src/server/lifecycle"; import type { OcxConfig } from "../src/types"; import type { RequestLogContext } from "../src/server/request-log"; import type { ResponsesTerminalStatus } from "../src/bridge"; +import { + codexHeaders, + encryptedInput as recoverableEncryptedInput, + recoverySse, +} from "./helpers/agent-task-recovery"; setDefaultTimeout(30_000); @@ -50,8 +71,10 @@ beforeEach(() => { process.env.CODEX_HOME = testDir; clearThreadAccountMap(); clearCodexUpstreamHealth(); + clearUpstreamHostHealth(); clearAccountQuota(); resetSubagentModelFallbackStateForTests(); + resetAgentTaskRecoveryState(); }); afterEach(() => { @@ -59,8 +82,10 @@ afterEach(() => { Date.now = originalNow; clearThreadAccountMap(); clearCodexUpstreamHealth(); + clearUpstreamHostHealth(); clearAccountQuota(); resetSubagentModelFallbackStateForTests(); + resetAgentTaskRecoveryState(); rmSync(testDir, { recursive: true, force: true }); if (previousOpencodexHome === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = previousOpencodexHome; @@ -199,6 +224,405 @@ async function postSpawn( } describe("subagent fallback without primary auth cooldown failure", () => { + test("fallback-chain resolution failure releases the preview selection admission", async () => { + writeFileSync(join(testDir, "agents"), "not a directory", "utf8"); + let selectionStarts = 0; + let selectionReleases = 0; + const turnAdmissionLease = { + release() {}, + beginCodexAccountSelection() { + selectionStarts += 1; + return { + mainProfileDraining: false, + claimMainProfile: () => false, + release: () => { selectionReleases += 1; }, + }; + }, + } satisfies Pick; + const cfg = poolNativePlusRoutedConfig({ + subagentModelFallback: ["xai/grok-4.5"], + }); + + await expect(postSpawn( + cfg, + { model: "gpt-5.6-sol", input: readableAgentInput(), stream: false }, + { turnAdmissionLease }, + )).rejects.toThrow(); + + expect(selectionStarts).toBe(1); + expect(selectionReleases).toBe(1); + }); + + test("open ChatGPT circuit skips gated entitlement discovery for a healthy routed primary", async () => { + const now = 1_800_000_000_000; + Date.now = () => now; + const cfg = poolNativePlusRoutedConfig({ + upstreamHostCircuitThreshold: 1, + subagentModelFallback: ["gpt-daybreak-blue-latest"], + }); + const hostKey = upstreamHostHealthKey("openai", "https://chatgpt.com"); + const admission = acquireUpstreamHostAdmission(hostKey, 1, now); + expect(admission.kind).toBe("admitted"); + if (admission.kind !== "admitted" || !admission.lease) { + throw new Error("expected ChatGPT host circuit admission lease"); + } + recordUpstreamHostFailure(hostKey, { + code: "ECONNREFUSED", + now, + threshold: 1, + lease: admission.lease, + }); + + let entitlementCalls = 0; + const capture = { urls: [] as string[], bodies: [] as string[], auths: [] as Array }; + mockUpstream(capture); + const response = await postSpawn( + cfg, + { model: "xai/grok-4.5", input: readableAgentInput(), stream: false }, + { + resolveCodexModelEntitlements: async () => { + entitlementCalls += 1; + throw new Error("must not bypass ChatGPT host admission"); + }, + }, + ); + + expect(response.status).toBe(200); + expect(entitlementCalls).toBe(0); + expect(capture.urls).toHaveLength(1); + expect(capture.urls[0]).toContain("api.x.ai"); + expect(capture.urls.some(url => url.includes("chatgpt.com"))).toBe(false); + }); + + test("healthy routed primary does not discover entitlements for a later gated fallback", async () => { + const now = 1_800_000_000_000; + Date.now = () => now; + const cfg = poolNativePlusRoutedConfig({ + subagentModelFallback: ["gpt-daybreak-blue-latest"], + }); + setSubagentQuotaPrimeForTests(async () => {}); + + let entitlementCalls = 0; + const capture = { urls: [] as string[], bodies: [] as string[], auths: [] as Array }; + mockUpstream(capture); + const response = await postSpawn( + cfg, + { model: "xai/grok-4.5", input: readableAgentInput(), stream: false }, + { + resolveCodexModelEntitlements: async () => { + entitlementCalls += 1; + throw new Error("healthy routed primary must finish selection before gated discovery"); + }, + }, + ); + + expect(response.status).toBe(200); + expect(entitlementCalls).toBe(0); + expect(capture.urls).toHaveLength(1); + expect(capture.urls[0]).toContain("api.x.ai"); + expect(capture.urls.some(url => url.includes("chatgpt.com"))).toBe(false); + }); + + test("unavailable routed primary discovers entitlements only after reaching gated fallback", async () => { + const now = 1_800_000_000_000; + Date.now = () => now; + installPoolCredential("pool-a", "pool_acc", now); + const cfg = poolNativePlusRoutedConfig({ + autoSwitchThreshold: 0, + subagentModelFallback: ["gpt-daybreak-blue-latest"], + }); + setSubagentQuotaPrimeForTests(async () => {}); + const { noteSubagentModelFailure } = await import("../src/codex/subagent-model-fallback"); + noteSubagentModelFailure("xai/grok-4.5", "429", cfg, undefined, now); + const entitlementSnapshot = { + modelsByAccount: new Map([ + ["pool-a", new Set(["gpt-daybreak-blue-latest"])], + ]), + confirmedAccountIds: new Set(["pool-a"]), + credentialIdentities: new Map(), + }; + + let entitlementCalls = 0; + let finalAuth: CodexAuthContext | undefined; + const capture = { urls: [] as string[], bodies: [] as string[], auths: [] as Array }; + mockUpstream(capture); + const response = await postSpawn( + cfg, + { model: "xai/grok-4.5", input: readableAgentInput(), stream: false }, + { + onCodexAuthContextResolved: (ctx) => { finalAuth = ctx; }, + resolveCodexModelEntitlements: async () => { + entitlementCalls += 1; + return entitlementSnapshot; + }, + }, + ); + + expect(response.status).toBe(200); + // One lazy selection lookup plus final-auth revalidation. + expect(entitlementCalls).toBe(2); + expect(finalAuth).toMatchObject({ kind: "pool", accountId: "pool-a" }); + expect(capture.urls.some(url => url.includes("chatgpt.com/backend-api/codex"))).toBe(true); + expect(capture.urls.some(url => url.includes("api.x.ai"))).toBe(false); + expect(capture.bodies.some(body => body.includes('"model":"gpt-5.6-sol"'))).toBe(true); + }); + + test("entitlement discovery restarts selection from a newly healthy routed primary", async () => { + const now = 1_800_000_000_000; + let currentNow = now; + Date.now = () => currentNow; + installPoolCredential("pool-a", "pool_acc", now); + const cfg = poolNativePlusRoutedConfig({ + autoSwitchThreshold: 0, + subagentModelFallback: ["gpt-daybreak-blue-latest"], + }); + setSubagentQuotaPrimeForTests(async () => {}); + const { noteSubagentModelFailure } = await import("../src/codex/subagent-model-fallback"); + noteSubagentModelFailure("xai/grok-4.5", "429", cfg, undefined, now, 1_000); + const entitlementSnapshot = { + modelsByAccount: new Map([ + ["pool-a", new Set(["gpt-daybreak-blue-latest"])], + ]), + confirmedAccountIds: new Set(["pool-a"]), + credentialIdentities: new Map(), + }; + let entitlementCalls = 0; + let resolveDiscoveryStarted!: () => void; + const discoveryStarted = new Promise((resolve) => { resolveDiscoveryStarted = resolve; }); + let releaseDiscovery!: () => void; + const discoveryGate = new Promise((resolve) => { releaseDiscovery = resolve; }); + const capture = { urls: [] as string[], bodies: [] as string[], auths: [] as Array }; + mockUpstream(capture); + + const responsePromise = postSpawn( + cfg, + { model: "xai/grok-4.5", input: readableAgentInput(), stream: false }, + { + resolveCodexModelEntitlements: async () => { + entitlementCalls += 1; + resolveDiscoveryStarted(); + await discoveryGate; + return entitlementSnapshot; + }, + }, + ); + await discoveryStarted; + currentNow = now + 1_001; + releaseDiscovery(); + const response = await responsePromise; + + expect(response.status).toBe(200); + expect(entitlementCalls).toBe(1); + expect(capture.urls).toHaveLength(1); + expect(capture.urls[0]).toContain("api.x.ai"); + expect(capture.urls.some(url => url.includes("chatgpt.com"))).toBe(false); + }); + + test("circuit opening during credential work fences every gated roster fetch", async () => { + const now = 1_800_000_000_000; + Date.now = () => now; + const cfg = poolNativePlusRoutedConfig({ + upstreamHostCircuitThreshold: 1, + subagentModelFallback: ["gpt-daybreak-blue-latest", "xai/grok-3"], + }); + setSubagentQuotaPrimeForTests(async () => {}); + const { noteSubagentModelFailure } = await import("../src/codex/subagent-model-fallback"); + noteSubagentModelFailure("xai/grok-4.5", "429", cfg, undefined, now); + const hostKey = upstreamHostHealthKey("openai", "https://chatgpt.com"); + let rosterFetches = 0; + let circuitOpened = false; + let releaseRoster!: () => void; + const rosterGate = new Promise((resolve) => { releaseRoster = resolve; }); + const entitlementCredential = { + accountId: "pool-a", + accessToken: "pool-a-token", + chatgptAccountId: "pool_acc", + credentialIdentity: "test:pool-a", + }; + const deferredRosterFetcher = (async () => { + rosterFetches += 1; + await rosterGate; + return Response.json({ models: [] }); + }) as typeof fetch; + // Another request already owns a discovery flight before this request's + // credential work opens the circuit. + const staleFlight = resolveCodexModelEntitlements(cfg, { + credentials: [entitlementCredential], + fetcher: deferredRosterFetcher, + now, + }); + expect(rosterFetches).toBe(1); + const capture = { urls: [] as string[], bodies: [] as string[], auths: [] as Array }; + mockUpstream(capture); + + const responsePromise = postSpawn( + cfg, + { model: "xai/grok-4.5", input: readableAgentInput(), stream: false }, + { + resolveCodexModelEntitlements: async (config, resolveOptions) => { + const admission = acquireUpstreamHostAdmission(hostKey, 1, now); + expect(admission.kind).toBe("admitted"); + if (admission.kind !== "admitted" || !admission.lease) { + throw new Error("expected ChatGPT host circuit admission lease"); + } + recordUpstreamHostFailure(hostKey, { + code: "ECONNREFUSED", + now, + threshold: 1, + lease: admission.lease, + }); + circuitOpened = true; + return resolveCodexModelEntitlements(config, { + ...resolveOptions, + credentials: [entitlementCredential], + fetcher: deferredRosterFetcher, + }); + }, + }, + ); + const completedBeforeStaleFlight = await Promise.race([ + responsePromise.then(() => true), + new Promise((resolve) => setTimeout(() => resolve(false), 2_000)), + ]); + releaseRoster(); + const [response] = await Promise.all([responsePromise, staleFlight]); + + expect(response.status).toBe(200); + expect(completedBeforeStaleFlight).toBe(true); + expect(circuitOpened).toBe(true); + expect(rosterFetches).toBe(1); + expect(capture.urls).toHaveLength(1); + expect(capture.urls[0]).toContain("api.x.ai"); + expect(capture.urls.some(url => url.includes("chatgpt.com"))).toBe(false); + }); + + test("probe-due gated primary keeps half-open ownership from a same-model fixed fallback", async () => { + const now = 1_800_000_000_000; + let currentNow = now; + Date.now = () => currentNow; + installPoolCredential("pool-a", "pool_acc", now); + installPoolCredential("pool-b", "pool_b_acc", now); + const cfg = poolNativePlusRoutedConfig({ + activeCodexAccountId: "pool-a", + upstreamHostCircuitThreshold: 1, + codexAccountNamespaces: { team: "pool-b" }, + subagentModelFallback: ["team/gpt-daybreak-blue-latest", "xai/grok-4.5"], + codexAccounts: [ + { id: "main", email: "main@example.test", isMain: true }, + { id: "pool-a", email: "a@example.test", isMain: false, chatgptAccountId: "pool_acc" }, + { id: "pool-b", email: "b@example.test", isMain: false, chatgptAccountId: "pool_b_acc" }, + ], + }); + const hostKey = upstreamHostHealthKey("openai", "https://chatgpt.com"); + const admission = acquireUpstreamHostAdmission(hostKey, 1, now); + expect(admission.kind).toBe("admitted"); + if (admission.kind !== "admitted" || !admission.lease) { + throw new Error("expected ChatGPT host circuit admission lease"); + } + recordUpstreamHostFailure(hostKey, { + code: "ECONNREFUSED", + now, + threshold: 1, + lease: admission.lease, + }); + const cooldownUntil = getUpstreamHostHealth(hostKey)?.cooldownUntil; + expect(cooldownUntil).toBeNumber(); + currentNow = (cooldownUntil ?? now) + 1; + const { noteSubagentModelFailure } = await import("../src/codex/subagent-model-fallback"); + noteSubagentModelFailure("gpt-daybreak-blue-latest", "429", cfg, "pool-a", now); + + const entitlementSnapshot = { + modelsByAccount: new Map([ + ["pool-a", new Set(["gpt-daybreak-blue-latest"])], + ["pool-b", new Set(["gpt-daybreak-blue-latest"])], + ]), + confirmedAccountIds: new Set(["pool-a", "pool-b"]), + credentialIdentities: new Map(), + }; + let entitlementCalls = 0; + let finalAuth: CodexAuthContext | undefined; + const capture = { urls: [] as string[], bodies: [] as string[], auths: [] as Array }; + mockUpstream(capture); + const response = await postSpawn( + cfg, + { model: "gpt-daybreak-blue-latest", input: readableAgentInput(), stream: false }, + { + onCodexAuthContextResolved: (ctx) => { finalAuth = ctx; }, + resolveCodexModelEntitlements: async () => { + entitlementCalls += 1; + return entitlementSnapshot; + }, + }, + ); + + expect(response.status).toBe(200); + // One call belongs to final auth after half-open host admission. The former + // speculative pre-admission discovery would make this two. + expect(entitlementCalls).toBe(1); + expect(capture.urls.length).toBeGreaterThan(0); + expect(capture.urls.some(url => url.includes("chatgpt.com/backend-api/codex/responses"))).toBe(true); + expect(capture.urls.some(url => url.includes("/backend-api/codex/models"))).toBe(false); + expect(capture.urls.some(url => url.includes("api.x.ai"))).toBe(false); + expect(finalAuth).toMatchObject({ kind: "pool", accountId: "pool-a" }); + expect(getUpstreamHostHealth(hostKey)).toBeNull(); + }); + + test("occupied half-open probe lets a concurrent gated primary use a healthy routed fallback", async () => { + const now = 1_800_000_000_000; + let currentNow = now; + Date.now = () => currentNow; + const cfg = poolNativePlusRoutedConfig({ + upstreamHostCircuitThreshold: 1, + subagentModelFallback: ["xai/grok-4.5"], + }); + const hostKey = upstreamHostHealthKey("openai", "https://chatgpt.com"); + const openingAdmission = acquireUpstreamHostAdmission(hostKey, 1, now); + expect(openingAdmission.kind).toBe("admitted"); + if (openingAdmission.kind !== "admitted" || !openingAdmission.lease) { + throw new Error("expected ChatGPT host circuit admission lease"); + } + recordUpstreamHostFailure(hostKey, { + code: "ECONNREFUSED", + now, + threshold: 1, + lease: openingAdmission.lease, + }); + const cooldownUntil = getUpstreamHostHealth(hostKey)?.cooldownUntil; + expect(cooldownUntil).toBeNumber(); + currentNow = (cooldownUntil ?? now) + 1; + const probeOwner = acquireUpstreamHostAdmission(hostKey, 1, currentNow); + expect(probeOwner.kind).toBe("admitted"); + if (probeOwner.kind !== "admitted" || !probeOwner.lease) { + throw new Error("expected half-open probe owner"); + } + expect(probeOwner.lease.halfOpen).toBe(true); + + let entitlementCalls = 0; + const capture = { urls: [] as string[], bodies: [] as string[], auths: [] as Array }; + mockUpstream(capture); + const response = await postSpawn( + cfg, + { model: "gpt-daybreak-blue-latest", input: readableAgentInput(), stream: false }, + { + resolveCodexModelEntitlements: async () => { + entitlementCalls += 1; + throw new Error("entitlement discovery must remain fenced"); + }, + }, + ); + + expect(response.status).toBe(200); + expect(entitlementCalls).toBe(0); + expect(capture.urls).toHaveLength(1); + expect(capture.urls[0]).toContain("api.x.ai"); + expect(capture.urls.some(url => url.includes("chatgpt.com"))).toBe(false); + expect(acquireUpstreamHostAdmission(hostKey, 1, currentNow)).toEqual({ + kind: "blocked", + retryAfterSeconds: 1, + }); + expect(releaseUpstreamHostAdmission(probeOwner.lease, currentNow)).toBe(true); + }); + test("exact account child bypasses quota priming and fallback on an empty 503", async () => { const now = 1_800_000_000_000; Date.now = () => now; @@ -346,6 +770,223 @@ describe("subagent fallback without primary auth cooldown failure", () => { expect(response.status).not.toBe(429); }); + test("concurrent probe-due spawns reserve once and advance the loser to routed fallback", async () => { + const now = 1_800_000_000_000; + let currentNow = now; + Date.now = () => currentNow; + installPoolCredential("pool-a", "pool_acc", now); + const cfg = poolNativePlusRoutedConfig({ + autoSwitchThreshold: 0, + subagentModelFallback: ["gpt-daybreak-blue-latest", "xai/grok-3"], + }); + setSubagentQuotaPrimeForTests(async () => {}); + updateAccountQuota("pool-a", 10, undefined, 20); + recordCodexUpstreamOutcome(cfg, "pool-a", 429, { + modelId: "gpt-daybreak-blue-latest", + now, + resetAt: Math.floor((now + 60 * 60_000) / 1_000), + }); + const { noteSubagentModelFailure } = await import("../src/codex/subagent-model-fallback"); + noteSubagentModelFailure("xai/grok-4.5", "429", cfg, undefined, now, 60 * 60_000); + currentNow = now + CODEX_QUOTA_PROBE_INTERVAL_MS + 1; + + const entitlementSnapshot = { + modelsByAccount: new Map([ + ["pool-a", new Set(["gpt-daybreak-blue-latest"])], + ]), + confirmedAccountIds: new Set(["pool-a"]), + credentialIdentities: new Map(), + }; + let lazySelectionsStarted = 0; + let releaseLazySelections!: () => void; + const lazySelectionGate = new Promise((resolve) => { releaseLazySelections = resolve; }); + let resolveBothLazySelectionsStarted!: () => void; + const bothLazySelectionsStarted = new Promise((resolve) => { + resolveBothLazySelectionsStarted = resolve; + }); + const resolverCalls = [0, 0]; + let probeOccupiedBeforeFinalAuth = false; + const entitlementResolver = (requestIndex: number) => async () => { + resolverCalls[requestIndex] += 1; + if (resolverCalls[requestIndex] === 1) { + lazySelectionsStarted += 1; + if (lazySelectionsStarted === 2) { + resolveBothLazySelectionsStarted(); + releaseLazySelections(); + } + await lazySelectionGate; + } else if (resolverCalls[requestIndex] === 2) { + probeOccupiedBeforeFinalAuth = !canAcquireCodexQuotaScopeProbeLease( + "pool-a", + "shared", + currentNow, + ); + } + return entitlementSnapshot; + }; + + const authPublications: Array = []; + const capture = { urls: [] as string[], bodies: [] as string[], auths: [] as Array }; + mockUpstream(capture); + const requests = [0, 1].map((requestIndex) => postSpawn( + cfg, + { model: "xai/grok-4.5", input: readableAgentInput(), stream: false }, + { + onCodexAuthContextResolved: (ctx) => authPublications.push(ctx), + resolveCodexModelEntitlements: entitlementResolver(requestIndex), + }, + { model: "", provider: "" }, + { + "session-id": `probe-race-session-${requestIndex}`, + "thread-id": `probe-race-thread-${requestIndex}`, + }, + )); + await bothLazySelectionsStarted; + const responses = await Promise.all(requests); + + expect(responses.map(response => response.status)).toEqual([200, 200]); + expect(resolverCalls.reduce((sum, calls) => sum + calls, 0)).toBe(3); + expect(capture.urls.filter(url => url.includes("chatgpt.com/backend-api/codex"))).toHaveLength(1); + expect(capture.urls.filter(url => url.includes("api.x.ai"))).toHaveLength(1); + const nativeAuth = authPublications.filter( + (ctx): ctx is CodexAuthContext => ctx !== undefined, + ); + expect(nativeAuth).toHaveLength(1); + expect(nativeAuth[0]).toMatchObject({ + kind: "pool", + accountId: "pool-a", + probeQuotaScope: "shared", + }); + expect(probeOccupiedBeforeFinalAuth).toBe(true); + expect((nativeAuth[0] as { probeLeaseId?: string }).probeLeaseId).toBeTruthy(); + expect(authPublications.filter(ctx => ctx === undefined)).toHaveLength(1); + expect(getCodexQuotaHealthSnapshot("pool-a", "shared", currentNow)).toBeNull(); + }); + + test("pre-auth continuation refusal releases a reserved scoped probe", async () => { + const now = 1_800_000_000_000; + let currentNow = now; + Date.now = () => currentNow; + installPoolCredential("pool-a", "pool_acc", now); + const cfg = poolNativePlusRoutedConfig({ + autoSwitchThreshold: 0, + subagentModelFallback: ["xai/grok-4.5"], + }); + setSubagentQuotaPrimeForTests(async () => {}); + updateAccountQuota("pool-a", 10, undefined, 20); + recordCodexUpstreamOutcome(cfg, "pool-a", 429, { + modelId: "gpt-5.6-sol", + now, + resetAt: Math.floor((now + 60 * 60_000) / 1_000), + }); + currentNow = now + CODEX_QUOTA_PROBE_INTERVAL_MS + 1; + + let fetchCalls = 0; + globalThis.fetch = (async () => { + fetchCalls += 1; + throw new Error("must not dispatch"); + }) as typeof fetch; + const response = await postSpawn(cfg, { + model: "gpt-5.6-sol", + input: readableAgentInput(), + previous_response_id: "resp_missing", + stream: false, + }); + + expect(response.status).toBe(400); + expect(fetchCalls).toBe(0); + expect(getCodexQuotaHealthSnapshot("pool-a", "shared", currentNow)).toMatchObject({ + quotaScope: "shared", + }); + currentNow += CODEX_QUOTA_PROBE_INTERVAL_MS + 1; + const reacquiredLease = tryAcquireCodexQuotaScopeProbeLease( + "pool-a", + "shared", + currentNow, + ); + expect(reacquiredLease).toBeTruthy(); + if (reacquiredLease) { + releaseCodexQuotaScopeProbeLease("pool-a", "shared", reacquiredLease, currentNow); + } + }); + + test("newer scoped cooldown fences and releases a stale pre-auth reservation", async () => { + const now = 1_800_000_000_000; + let currentNow = now; + Date.now = () => currentNow; + installPoolCredential("pool-a", "pool_acc", now); + const cfg = poolNativePlusRoutedConfig({ + autoSwitchThreshold: 0, + subagentModelFallback: ["gpt-daybreak-blue-latest", "xai/grok-3"], + }); + setSubagentQuotaPrimeForTests(async () => {}); + updateAccountQuota("pool-a", 10, undefined, 20); + recordCodexUpstreamOutcome(cfg, "pool-a", 429, { + modelId: "gpt-daybreak-blue-latest", + now, + resetAt: Math.floor((now + 60 * 60_000) / 1_000), + }); + const { noteSubagentModelFailure } = await import("../src/codex/subagent-model-fallback"); + noteSubagentModelFailure("xai/grok-4.5", "429", cfg, undefined, now, 60 * 60_000); + currentNow = now + CODEX_QUOTA_PROBE_INTERVAL_MS + 1; + + const entitlementSnapshot = { + modelsByAccount: new Map([ + ["pool-a", new Set(["gpt-daybreak-blue-latest"])], + ]), + confirmedAccountIds: new Set(["pool-a"]), + credentialIdentities: new Map(), + }; + let entitlementCalls = 0; + let sawReservedProbe = false; + 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 }, + { + resolveCodexModelEntitlements: async () => { + entitlementCalls += 1; + if (entitlementCalls === 2) { + sawReservedProbe = !canAcquireCodexQuotaScopeProbeLease( + "pool-a", + "shared", + currentNow, + ); + recordCodexUpstreamOutcome(cfg, "pool-a", 429, { + modelId: "gpt-daybreak-blue-latest", + now: currentNow, + resetAt: Math.floor((currentNow + 60 * 60_000) / 1_000), + }); + } + return entitlementSnapshot; + }, + }, + ); + + expect(response.status).toBe(429); + expect(entitlementCalls).toBe(2); + expect(sawReservedProbe).toBe(true); + expect(fetchCalls).toBe(0); + expect(getCodexQuotaHealthSnapshot("pool-a", "shared", currentNow)).toMatchObject({ + cooldownSource: "reset-derived", + quotaScope: "shared", + }); + currentNow += CODEX_QUOTA_PROBE_INTERVAL_MS + 1; + const reacquiredLease = tryAcquireCodexQuotaScopeProbeLease( + "pool-a", + "shared", + currentNow, + ); + expect(reacquiredLease).toBeTruthy(); + if (reacquiredLease) { + releaseCodexQuotaScopeProbeLease("pool-a", "shared", reacquiredLease, currentNow); + } + }); + test("final-route direct auth failure does not acquire a pool probe lease", async () => { const now = 1_800_000_000_000; Date.now = () => now; @@ -775,29 +1416,492 @@ describe("native fallback account preview", () => { expect(finalAuth).toMatchObject({ kind: "pool", accountId: "pool-a" }); }); - test("uses healthier pool account B when active A is above threshold", async () => { + test("fallback previews the pool account separately for each candidate quota scope", async () => { const now = 1_800_000_000_000; Date.now = () => now; installPoolCredential("pool-a", "pool_acc_a", now); installPoolCredential("pool-b", "pool_acc_b", now); const cfg = poolNativePlusRoutedConfig({ - defaultProvider: "xai", activeCodexAccountId: "pool-a", - subagentModelFallback: ["gpt-5.6-terra"], + autoSwitchThreshold: 0, + subagentModelFallback: ["gpt-5.3-codex-spark", "xai/grok-4.5"], codexAccounts: [ { id: "main", email: "main@example.test", isMain: true }, { 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" }, ], - providers: { - xai: { - adapter: "openai-chat", - baseUrl: "https://api.x.ai/v1", - authMode: "key", - apiKey: "xai-test", - }, - openai: { - adapter: "openai-responses", + }); + const desktopHeaders = { + "session-id": "candidate-scope-session-private", + "thread-id": "candidate-scope-thread-private", + }; + const bound = await resolveCodexAuthContext(new Headers(desktopHeaders), cfg, "pool", { + modelId: "gpt-5.6-sol", + }); + expect(bound).toMatchObject({ kind: "pool", accountId: "pool-a" }); + if (bound.kind !== "pool") throw new Error("expected pool context"); + + recordCodexUpstreamOutcome(cfg, "pool-a", 429, { + modelId: "gpt-5.3-codex-spark", + now, + resetAt: Math.floor((now + 60 * 60_000) / 1_000), + }); + const { noteSubagentModelFailure } = await import("../src/codex/subagent-model-fallback"); + noteSubagentModelFailure("gpt-5.3-codex-spark", "429", cfg, "pool-a", now); + noteSubagentModelFailure("gpt-5.6-sol", "429", cfg, "pool-a", now); + + expect(previewCodexAccountForRequest(bound.affinityKey ?? null, cfg, now, "shared")).toBe("pool-a"); + expect(previewCodexAccountForRequest(bound.affinityKey ?? null, cfg, now, "spark")).toBe("pool-b"); + + let finalAuth: CodexAuthContext | undefined; + const capture = { urls: [] as string[], bodies: [] as string[], auths: [] as Array }; + mockUpstream(capture); + + const response = await postSpawn( + cfg, + { model: "gpt-5.6-sol", input: readableAgentInput(), stream: false }, + { onCodexAuthContextResolved: (ctx) => { finalAuth = ctx; } }, + { model: "", provider: "" }, + desktopHeaders, + ); + + expect(response.status).toBe(200); + expect(finalAuth).toMatchObject({ kind: "pool", accountId: "pool-b" }); + expect(capture.urls.some((url) => url.includes("chatgpt.com/backend-api/codex"))).toBe(true); + expect(capture.auths.some((auth) => auth?.includes("pool-b_token"))).toBe(true); + expect(capture.bodies.some((body) => body.includes('"model":"gpt-5.3-codex-spark"'))).toBe(true); + }); + + test("account-gated fallback previews and authenticates an entitled pool account", async () => { + const now = 1_800_000_000_000; + Date.now = () => now; + installPoolCredential("pool-a", "pool_acc_a", now); + installPoolCredential("pool-b", "pool_acc_b", now); + const cfg = poolNativePlusRoutedConfig({ + activeCodexAccountId: "pool-a", + autoSwitchThreshold: 0, + subagentModelFallback: ["gpt-daybreak-blue-latest", "xai/grok-4.5"], + codexAccounts: [ + { id: "main", email: "main@example.test", isMain: true }, + { 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" }, + ], + }); + const desktopHeaders = { + "session-id": "candidate-entitlement-session-private", + "thread-id": "candidate-entitlement-thread-private", + }; + const bound = await resolveCodexAuthContext(new Headers(desktopHeaders), cfg, "pool", { + modelId: "gpt-5.6-sol", + }); + expect(bound).toMatchObject({ kind: "pool", accountId: "pool-a" }); + const { noteSubagentModelFailure } = await import("../src/codex/subagent-model-fallback"); + noteSubagentModelFailure("gpt-5.6-sol", "429", cfg, "pool-a", now); + const entitlementSnapshot = { + modelsByAccount: new Map([ + ["pool-a", new Set(["gpt-5.6-sol"])], + ["pool-b", new Set(["gpt-5.6-sol", "gpt-daybreak-blue-latest"])], + ]), + confirmedAccountIds: new Set(["pool-a", "pool-b"]), + credentialIdentities: new Map(), + }; + let entitlementCalls = 0; + + let finalAuth: CodexAuthContext | undefined; + const capture = { urls: [] as string[], bodies: [] as string[], auths: [] as Array }; + mockUpstream(capture); + + const response = await postSpawn( + cfg, + { model: "gpt-5.6-sol", input: readableAgentInput(), stream: false }, + { + onCodexAuthContextResolved: (ctx) => { finalAuth = ctx; }, + resolveCodexModelEntitlements: async () => { + entitlementCalls += 1; + return entitlementSnapshot; + }, + }, + { model: "", provider: "" }, + desktopHeaders, + ); + + expect(response.status).toBe(200); + expect(entitlementCalls).toBe(2); + expect(finalAuth).toMatchObject({ kind: "pool", accountId: "pool-b" }); + expect(capture.urls.some((url) => url.includes("chatgpt.com/backend-api/codex"))).toBe(true); + expect(capture.auths.some((auth) => auth?.includes("pool-b_token"))).toBe(true); + expect(capture.urls.some((url) => url.includes("api.x.ai"))).toBe(false); + }); + + test("account-qualified primary previews an entitled pool account for a gated fallback", async () => { + const now = 1_800_000_000_000; + Date.now = () => now; + installPoolCredential("pool-a", "pool_acc_a", now); + installPoolCredential("pool-b", "pool_acc_b", now); + const cfg = poolNativePlusRoutedConfig({ + activeCodexAccountId: "pool-a", + autoSwitchThreshold: 0, + codexAccountNamespaces: { team: "pool-a" }, + subagentModelFallback: ["gpt-daybreak-blue-latest", "xai/grok-4.5"], + codexAccounts: [ + { id: "main", email: "main@example.test", isMain: true }, + { 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" }, + ], + }); + 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([ + ["pool-a", new Set(["gpt-5.6-sol"])], + ["pool-b", new Set(["gpt-5.6-sol", "gpt-daybreak-blue-latest"])], + ]), + confirmedAccountIds: new Set(["pool-a", "pool-b"]), + credentialIdentities: new Map(), + }; + let entitlementCalls = 0; + let finalAuth: CodexAuthContext | undefined; + const logCtx: RequestLogContext = { model: "", provider: "" }; + const capture = { urls: [] as string[], bodies: [] as string[], auths: [] as Array }; + mockUpstream(capture); + + const response = await postSpawn( + cfg, + { model: "team/gpt-5.6-sol", input: readableAgentInput(), stream: false }, + { + onCodexAuthContextResolved: (ctx) => { finalAuth = ctx; }, + resolveCodexModelEntitlements: async () => { + entitlementCalls += 1; + return entitlementSnapshot; + }, + }, + logCtx, + ); + + expect(response.status).toBe(200); + expect(entitlementCalls).toBe(2); + expect(finalAuth).toMatchObject({ kind: "pool", accountId: "pool-b" }); + expect((logCtx as unknown as Record).subagentModelFallbackTo) + .toBe("gpt-daybreak-blue-latest"); + expect(capture.bodies.filter(body => body.length > 0)).toHaveLength(1); + expect(capture.auths[0]).toContain("pool-b_token"); + expect(capture.auths.some((auth) => auth?.includes("pool-b_token"))).toBe(true); + }); + + test("account-qualified primary reroutes when fallback drops its selector from the same model", async () => { + const now = 1_800_000_000_000; + Date.now = () => now; + installPoolCredential("pool-a", "pool_acc_a", now); + installPoolCredential("pool-b", "pool_acc_b", now); + const cfg = poolNativePlusRoutedConfig({ + activeCodexAccountId: "pool-a", + autoSwitchThreshold: 0, + codexAccountNamespaces: { team: "pool-a" }, + subagentModelFallback: ["gpt-5.6-sol", "xai/grok-4.5"], + codexAccounts: [ + { id: "main", email: "main@example.test", isMain: true }, + { 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" }, + ], + }); + recordCodexUpstreamOutcome(cfg, "pool-a", 429, { + fixedAccount: true, + modelId: "gpt-5.6-sol", + now, + resetAt: Math.floor((now + 60 * 60_000) / 1_000), + }); + let finalAuth: CodexAuthContext | undefined; + const logCtx: RequestLogContext = { model: "", provider: "" }; + const capture = { urls: [] as string[], bodies: [] as string[], auths: [] as Array }; + mockUpstream(capture); + + const response = await postSpawn( + cfg, + { model: "team/gpt-5.6-sol", input: readableAgentInput(), stream: false }, + { onCodexAuthContextResolved: (ctx) => { finalAuth = ctx; } }, + logCtx, + ); + + expect(response.status).toBe(200); + expect((logCtx as unknown as Record).subagentModelFallbackFrom) + .toBe("team/gpt-5.6-sol"); + expect((logCtx as unknown as Record).subagentModelFallbackTo) + .toBe("gpt-5.6-sol"); + expect(finalAuth).toMatchObject({ kind: "pool", accountId: "pool-b" }); + const responseAuths = capture.auths.filter((_auth, index) => capture.bodies[index]?.length > 0); + expect(responseAuths).toHaveLength(1); + expect(responseAuths[0]).toContain("pool-b_token"); + expect(capture.urls.some((url) => url.includes("api.x.ai"))).toBe(false); + }); + + test("account-qualified primary skips an unentitled fixed gated fallback", async () => { + const now = 1_800_000_000_000; + Date.now = () => now; + installPoolCredential("pool-a", "pool_acc_a", now); + installPoolCredential("pool-b", "pool_acc_b", now); + const cfg = poolNativePlusRoutedConfig({ + activeCodexAccountId: "pool-a", + autoSwitchThreshold: 0, + codexAccountNamespaces: { team: "pool-a", restricted: "pool-b" }, + subagentModelFallback: ["restricted/gpt-daybreak-blue-latest", "xai/grok-4.5"], + codexAccounts: [ + { id: "main", email: "main@example.test", isMain: true }, + { 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" }, + ], + }); + 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([ + ["pool-a", new Set(["gpt-5.6-sol", "gpt-daybreak-blue-latest"])], + ["pool-b", new Set(["gpt-5.6-sol"])], + ]), + confirmedAccountIds: new Set(["pool-a", "pool-b"]), + credentialIdentities: new Map(), + }; + let entitlementCalls = 0; + const logCtx: RequestLogContext = { model: "", provider: "" }; + const capture = { urls: [] as string[], bodies: [] as string[], auths: [] as Array }; + mockUpstream(capture); + + const response = await postSpawn( + cfg, + { model: "team/gpt-5.6-sol", input: readableAgentInput(), stream: false }, + { + resolveCodexModelEntitlements: async () => { + entitlementCalls += 1; + return entitlementSnapshot; + }, + }, + logCtx, + ); + + expect(response.status).toBe(200); + expect(entitlementCalls).toBe(1); + expect((logCtx as unknown as Record).subagentModelFallbackTo).toBe("xai/grok-4.5"); + expect(capture.urls).toHaveLength(1); + expect(capture.urls.some((url) => url.includes("api.x.ai"))).toBe(true); + expect(capture.urls.some((url) => url.includes("chatgpt.com/backend-api/codex"))).toBe(false); + expect(capture.bodies.some((body) => body.includes("gpt-daybreak-blue-latest"))).toBe(false); + }); + + test("account-qualified fallback excludes native main entitlement reads during profile drain", async () => { + const now = 1_800_000_000_000; + Date.now = () => now; + installPoolCredential("pool-a", "pool_acc_a", now); + installPoolCredential("pool-b", "pool_acc_b", now); + const cfg = poolNativePlusRoutedConfig({ + activeCodexAccountId: "pool-a", + autoSwitchThreshold: 0, + codexAccountNamespaces: { team: "pool-a" }, + subagentModelFallback: ["gpt-daybreak-blue-latest", "xai/grok-4.5"], + codexAccounts: [ + { id: "main", email: "main@example.test", isMain: true }, + { 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" }, + ], + }); + 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([ + ["__main__", new Set(["gpt-daybreak-blue-latest"])], + ["pool-b", new Set(["gpt-daybreak-blue-latest"])], + ]), + confirmedAccountIds: new Set(["__main__", "pool-b"]), + credentialIdentities: new Map(), + }; + const mainExclusions: boolean[] = []; + let selectionReleases = 0; + const turnAdmissionLease = { + release() {}, + beginCodexAccountSelection() { + return { + mainProfileDraining: true, + claimMainProfile: () => false, + release: () => { selectionReleases += 1; }, + }; + }, + } satisfies Pick; + let finalAuth: CodexAuthContext | undefined; + const capture = { urls: [] as string[], bodies: [] as string[], auths: [] as Array }; + mockUpstream(capture); + + const response = await postSpawn( + cfg, + { model: "team/gpt-5.6-sol", input: readableAgentInput(), stream: false }, + { + turnAdmissionLease, + onCodexAuthContextResolved: (ctx) => { finalAuth = ctx; }, + resolveCodexModelEntitlements: async (_config, resolveOptions) => { + mainExclusions.push(resolveOptions?.excludeAccountIds?.has("__main__") === true); + return entitlementSnapshot; + }, + }, + ); + + expect(response.status).toBe(200); + expect(mainExclusions).toEqual([true, true]); + expect(selectionReleases).toBe(2); + expect(finalAuth).toMatchObject({ kind: "pool", accountId: "pool-b" }); + expect(capture.auths[0]).toContain("pool-b_token"); + }); + + test("recovery second pass re-previews the account for the newly available candidate scope", async () => { + const now = 1_800_000_000_000; + let currentNow = now; + Date.now = () => currentNow; + installPoolCredential("pool-a", "pool_acc_a", now); + installPoolCredential("pool-b", "pool_acc_b", now); + const cfg = poolNativePlusRoutedConfig({ + defaultProvider: "xai", + activeCodexAccountId: "pool-a", + autoSwitchThreshold: 0, + agentTaskRecovery: { enabled: true }, + subagentModelFallback: ["gpt-5.3-codex-spark", "gpt-daybreak-blue-latest"], + codexAccounts: [ + { id: "main", email: "main@example.test", isMain: true }, + { 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" }, + ], + }); + const requestHeaders = codexHeaders("caller-account", { + "session-id": "recovery-candidate-scope-session-private", + "thread-id": "recovery-candidate-scope-thread-private", + }); + const bound = await resolveCodexAuthContext(requestHeaders, cfg, "pool", { + modelId: "gpt-5.6-sol", + }); + expect(bound).toMatchObject({ kind: "pool", accountId: "pool-a" }); + if (bound.kind !== "pool") throw new Error("expected pool context"); + + recordCodexUpstreamOutcome(cfg, "pool-a", 429, { + modelId: "gpt-5.3-codex-spark", + now, + resetAt: Math.floor((now + 60 * 60_000) / 1_000), + }); + const { noteSubagentModelFailure } = await import("../src/codex/subagent-model-fallback"); + noteSubagentModelFailure("gpt-5.3-codex-spark", "429", cfg, "pool-b", now); + noteSubagentModelFailure("xai/grok-4.5", "429", cfg, undefined, now, 10 * 60_000); + const entitlementSnapshot = { + modelsByAccount: new Map([ + ["pool-a", new Set(["gpt-5.6-sol"])], + ["pool-b", new Set(["gpt-5.6-sol"])], + ]), + confirmedAccountIds: new Set(["pool-a", "pool-b"]), + credentialIdentities: new Map(), + }; + let entitlementCalls = 0; + + let selectionStarts = 0; + let selectionReleases = 0; + const turnAdmissionLease = { + release() {}, + beginCodexAccountSelection() { + selectionStarts += 1; + return { + mainProfileDraining: false, + claimMainProfile: () => false, + release: () => { selectionReleases += 1; }, + }; + }, + } satisfies Pick; + let finalAuth: CodexAuthContext | undefined; + const fetchedUrls: string[] = []; + const forwardedBodies: string[] = []; + const forwardedAuths: Array = []; + globalThis.fetch = (async (input, init) => { + const url = String(input); + const raw = typeof init?.body === "string" ? init.body : ""; + fetchedUrls.push(url); + forwardedBodies.push(raw); + forwardedAuths.push(new Headers(init?.headers).get("authorization")); + if (raw.includes("capture_assignment")) { + currentNow = now + DEFAULT_SUBAGENT_MODEL_FALLBACK_POLL_MS + 1; + return new Response(recoverySse("Use the recovered candidate-scope assignment."), { + status: 200, + headers: { "content-type": "text/event-stream" }, + }); + } + return Response.json({ + id: "resp_recovered_candidate_scope", + object: "response", + status: "completed", + model: "gpt-5.3-codex-spark", + output: [], + usage: { input_tokens: 1, output_tokens: 1, total_tokens: 2 }, + }); + }) as typeof fetch; + + const response = await postSpawn( + cfg, + { model: "xai/grok-4.5", input: recoverableEncryptedInput(), stream: false }, + { + turnAdmissionLease, + onCodexAuthContextResolved: (ctx) => { finalAuth = ctx; }, + resolveCodexModelEntitlements: async () => { + entitlementCalls += 1; + return entitlementSnapshot; + }, + }, + { model: "", provider: "" }, + requestHeaders, + ); + + expect(response.status).toBe(200); + const bodyRequests = forwardedBodies.map((body, index) => ({ + body, + url: fetchedUrls[index], + auth: forwardedAuths[index], + })).filter(({ body }) => body.length > 0); + expect(bodyRequests).toHaveLength(2); + expect(bodyRequests[0]?.body).toContain("capture_assignment"); + expect(bodyRequests[1]?.body).toContain("Use the recovered candidate-scope assignment."); + expect(selectionStarts).toBe(3); + expect(selectionReleases).toBe(3); + // The ciphertext pass reaches the gated boundary once. Recovery selects + // the now-available ungated Spark candidate before that boundary. + expect(entitlementCalls).toBe(1); + expect(finalAuth).toMatchObject({ kind: "pool", accountId: "pool-b" }); + expect(bodyRequests[1]?.auth).toContain("pool-b_token"); + }); + + test("uses healthier pool account B when active A is above threshold", async () => { + const now = 1_800_000_000_000; + Date.now = () => now; + installPoolCredential("pool-a", "pool_acc_a", now); + installPoolCredential("pool-b", "pool_acc_b", now); + const cfg = poolNativePlusRoutedConfig({ + defaultProvider: "xai", + activeCodexAccountId: "pool-a", + subagentModelFallback: ["gpt-5.6-terra"], + codexAccounts: [ + { id: "main", email: "main@example.test", isMain: true }, + { 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" }, + ], + providers: { + xai: { + adapter: "openai-chat", + baseUrl: "https://api.x.ai/v1", + authMode: "key", + apiKey: "xai-test", + }, + openai: { + adapter: "openai-responses", baseUrl: "https://chatgpt.com/backend-api/codex", authMode: "forward", codexAccountMode: "pool", @@ -911,6 +2015,44 @@ describe("native fallback account preview", () => { }); describe("encrypted child native-only fallback", () => { + test("selects ungated native candidate before gated discovery for encrypted input", async () => { + const now = 1_800_000_000_000; + Date.now = () => now; + installPoolCredential("pool-a", "pool_acc", now); + const cfg = poolNativePlusRoutedConfig({ + defaultProvider: "xai", + autoSwitchThreshold: 0, + subagentModelFallback: ["gpt-5.6-terra", "gpt-daybreak-blue-latest"], + }); + setSubagentQuotaPrimeForTests(async () => {}); + let entitlementCalls = 0; + const capture = { urls: [] as string[], bodies: [] as string[], auths: [] as Array }; + mockUpstream(capture); + + const response = await postSpawn( + cfg, + { model: "xai/grok-4.5", input: encryptedAgentInput(), stream: false }, + { + resolveCodexModelEntitlements: async () => { + entitlementCalls += 1; + throw new Error("ungated native candidate must win before gated discovery"); + }, + }, + ); + + expect(response.status).toBe(200); + expect(entitlementCalls).toBe(0); + expect(capture.urls.some(url => url.includes("api.x.ai"))).toBe(false); + const bodyRequests = capture.bodies.map((body, index) => ({ + body, + url: capture.urls[index], + })).filter(({ body }) => body.length > 0); + expect(bodyRequests).toHaveLength(1); + expect(bodyRequests[0]?.url).toContain("chatgpt.com/backend-api/codex"); + expect(bodyRequests[0]?.body).toContain('"model":"gpt-5.6-terra"'); + expect(bodyRequests[0]?.body).toContain(FERNET_TASK); + }); + test("rejects encrypted routed primary when only routed fallbacks exist", async () => { const cfg = poolNativePlusRoutedConfig({ defaultProvider: "xai", diff --git a/tests/subagent-model-fallback.test.ts b/tests/subagent-model-fallback.test.ts index 3d433aef12..1b52770e52 100644 --- a/tests/subagent-model-fallback.test.ts +++ b/tests/subagent-model-fallback.test.ts @@ -14,9 +14,11 @@ import { resetSubagentModelFallbackStateForTests, resolveAgentModelFallbackForPrimary, resolveConfiguredModelFallbackForPrimary, + resolveSubagentFallbackChain, scanCodexAgentRolesWithTomlModelFallback, selectAvailableSubagentModel, setSubagentQuotaPrimeForTests, + subagentFallbackNeedsModelEntitlements, subagentFallbackGuidanceText, } from "../src/codex/subagent-model-fallback"; import { saveCodexAccountCredential } from "../src/codex/account-store"; @@ -24,6 +26,7 @@ import { clearAccountNeedsReauth, markAccountNeedsReauth } from "../src/codex/ac import { clearAccountQuota, updateAccountQuota } from "../src/codex/quota"; import { canAcquireCodexQuotaProbeLease, + canAcquireCodexQuotaScopeProbeLease, clearCodexUpstreamHealthForAccount, CODEX_QUOTA_PROBE_INTERVAL_MS, recordCodexUpstreamOutcome, @@ -147,6 +150,32 @@ describe("subagent model fallback chain", () => { }))).toEqual(["kimi/k3"]); }); + test("entitlement discovery is limited to account-gated pool fallback candidates", () => { + const parsed = { modelId: "gpt-5.6-sol" } as never; + const poolConfig = cfg({ + subagentModelFallback: ["gpt-daybreak-blue-latest"], + }); + expect(subagentFallbackNeedsModelEntitlements( + resolveSubagentFallbackChain(parsed, poolConfig), + poolConfig, + )).toBe(true); + const directConfig = cfg({ + providers: { + ...cfg().providers, + openai: { + adapter: "openai-responses", + baseUrl: "https://chatgpt.com/backend-api/codex", + codexAccountMode: "direct", + }, + }, + subagentModelFallback: ["gpt-daybreak-blue-latest"], + }); + expect(subagentFallbackNeedsModelEntitlements( + resolveSubagentFallbackChain(parsed, directConfig), + directConfig, + )).toBe(false); + }); + test("selectAvailableSubagentModel skips quota-exhausted native models", () => { resetSubagentModelFallbackStateForTests(); updateAccountQuota("pool-a", 95, undefined, 20); @@ -215,6 +244,54 @@ describe("subagent model fallback chain", () => { }); }); + test("fixed account candidates ignore pool preview and enforce candidate entitlements", () => { + updateAccountQuota("account-a", 10, undefined, 20); + const config = cfg({ codexAccountNamespaces: { team: "account-a" } }); + const throwingPreview = () => { + throw new Error("fixed account must not call pool preview"); + }; + + expect(isSubagentModelUnavailable( + "team/gpt-daybreak-blue-latest", + config, + "pool-a", + Date.now(), + undefined, + throwingPreview, + () => new Set(["account-a"]), + )).toBe(false); + expect(isSubagentModelUnavailable( + "team/gpt-daybreak-blue-latest", + config, + "pool-a", + Date.now(), + undefined, + throwingPreview, + () => new Set(["account-b"]), + )).toBe(true); + }); + + test("a null candidate account preview does not fall back to the active pool account", () => { + updateAccountQuota("pool-a", 10, undefined, 20); + const config = cfg({ subagentModelFallback: ["kimi/k3"] }); + + expect(selectAvailableSubagentModel( + "gpt-5.6-sol", + config, + [], + "pool-a", + Date.now(), + false, + undefined, + [], + () => null, + )).toEqual({ + model: "kimi/k3", + rewritten: true, + skipped: ["gpt-5.6-sol"], + }); + }); + test("case-distinct account selector fallbacks remain independent", () => { updateAccountQuota("pool-a", 95, undefined, 20); const config = cfg({ @@ -297,6 +374,78 @@ describe("subagent model fallback chain", () => { }); }); + test("pool fallback skips a reset-derived cooldown in the model's quota scope", () => { + const now = 1_800_000_000_000; + updateAccountQuota("pool-a", 10, undefined, 20); + const config = cfg({ subagentModelFallback: ["kimi/k3"] }); + recordCodexUpstreamOutcome(config, "pool-a", 429, { + modelId: "gpt-5.6-sol", + now, + resetAt: Math.floor((now + 60 * 60_000) / 1_000), + }); + + expect(selectAvailableSubagentModel("gpt-5.6-sol", config, [], "pool-a", now + 1)).toEqual({ + model: "kimi/k3", + rewritten: true, + skipped: ["gpt-5.6-sol"], + }); + }); + + test("pool fallback admits a due reset-derived probe in the model's quota scope", () => { + const now = 1_800_000_000_000; + const probeAt = now + CODEX_QUOTA_PROBE_INTERVAL_MS + 1; + updateAccountQuota("pool-a", 10, undefined, 20); + const config = cfg({ subagentModelFallback: ["kimi/k3"] }); + recordCodexUpstreamOutcome(config, "pool-a", 429, { + modelId: "gpt-5.6-sol", + now, + resetAt: Math.floor((now + 60 * 60_000) / 1_000), + }); + + expect(canAcquireCodexQuotaScopeProbeLease("pool-a", "shared", probeAt)).toBe(true); + expect(selectAvailableSubagentModel("gpt-5.6-sol", config, [], "pool-a", probeAt)).toEqual({ + model: "gpt-5.6-sol", + rewritten: false, + skipped: [], + }); + // Public selection remains a pure preview. Production request handling passes + // the internal atomic reserver only after every candidate predicate succeeds. + expect(canAcquireCodexQuotaScopeProbeLease("pool-a", "shared", probeAt)).toBe(true); + }); + + test("pool fallback ignores a reset-derived cooldown for an unrelated quota scope", () => { + const now = 1_800_000_000_000; + updateAccountQuota("pool-a", 10, undefined, 20); + const config = cfg({ subagentModelFallback: ["kimi/k3"] }); + recordCodexUpstreamOutcome(config, "pool-a", 429, { + modelId: "gpt-5.3-codex-spark", + now, + resetAt: Math.floor((now + 60 * 60_000) / 1_000), + }); + + expect(selectAvailableSubagentModel("gpt-5.6-sol", config, [], "pool-a", now + 1)).toEqual({ + model: "gpt-5.6-sol", + rewritten: false, + skipped: [], + }); + }); + + test("pool fallback preserves account-wide cooldown probe pacing", () => { + const now = 1_800_000_000_000; + const probeAt = now + CODEX_QUOTA_PROBE_INTERVAL_MS + 1; + updateAccountQuota("pool-a", 10, undefined, 20); + const config = cfg({ subagentModelFallback: ["kimi/k3"] }); + recordCodexUpstreamOutcome(config, "pool-a", 429, { + now, + resetAt: Math.floor((now + 60 * 60_000) / 1_000), + }); + + expect(selectAvailableSubagentModel("gpt-5.6-sol", config, [], "pool-a", now + 1).model).toBe("kimi/k3"); + expect(canAcquireCodexQuotaProbeLease("pool-a", probeAt)).toBe(true); + expect(selectAvailableSubagentModel("gpt-5.6-sol", config, [], "pool-a", probeAt).model) + .toBe("gpt-5.6-sol"); + }); + test("account selector fallbacks still reject invalid or disabled native models", () => { resetSubagentModelFallbackStateForTests(); updateAccountQuota("pool-a", 95, undefined, 20);