fix(codex): honor scoped cooldowns in subagent fallback

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,16 +331,16 @@ 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. */
   isMainAccountTokenLive?: () => boolean;
   getMainAccountToken?: typeof getMainAccountToken;
   primeCodexPoolQuotas?: (config: OcxConfig, reason: string) => Promise<void>;
   /** Test seam for account-gated native model discovery. */
-  resolveCodexModelEntitlements?: (
-    config: Pick<OcxConfig, "codexAccounts">,
-  ) => Promise<CodexModelEntitlementSnapshot>;
+  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);
     }

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<string>;
+  /** Last-mile circuit fence checked synchronously before a new roster fetch starts. */
+  readonly allowNetworkFetch?: () => boolean;
 }
 
 const accountModelsCache = new Map<string, CachedAccountModels>();
@@ -195,6 +201,7 @@ async function modelsForCredential(
   credential: CodexModelEntitlementCredentialSnapshot,
   fetcher: typeof fetch,
   now: number,
+  allowNetworkFetch?: () => boolean,
 ): Promise<CachedAccountModels> {
   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<CodexModelEntitlementSnapshot> {
   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])),

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