From 5587ccfa0c25b197645546e01323ab215b09f3d1 Mon Sep 17 00:00:00 2001 From: luvs01 <27862058+luvs01@users.noreply.github.com> Date: Wed, 26 Aug 2026 05:35:53 +0900 Subject: [PATCH] fix(codex): gate subagent fallback by account entitlement --- src/codex/auth-context.ts | 39 +-- src/codex/model-entitlements.ts | 11 +- src/codex/subagent-model-fallback.ts | 77 ++++-- src/server/responses/core.ts | 59 ++++- tests/codex-model-entitlements.test.ts | 57 +++-- ...subagent-fallback-handle-responses.test.ts | 228 ++++++++++++++++++ tests/subagent-model-fallback.test.ts | 65 ++++- 7 files changed, 473 insertions(+), 63 deletions(-) diff --git a/src/codex/auth-context.ts b/src/codex/auth-context.ts index 7dd58c6b91..f3c357d4c3 100644 --- a/src/codex/auth-context.ts +++ b/src/codex/auth-context.ts @@ -29,7 +29,6 @@ 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"; @@ -334,9 +333,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 +378,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. diff --git a/src/codex/model-entitlements.ts b/src/codex/model-entitlements.ts index d06100138b..5a649d335a 100644 --- a/src/codex/model-entitlements.ts +++ b/src/codex/model-entitlements.ts @@ -40,6 +40,10 @@ 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; } const accountModelsCache = new Map(); @@ -252,9 +256,12 @@ 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, diff --git a/src/codex/subagent-model-fallback.ts b/src/codex/subagent-model-fallback.ts index 40c472bb3c..1f56be2634 100644 --- a/src/codex/subagent-model-fallback.ts +++ b/src/codex/subagent-model-fallback.ts @@ -38,6 +38,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, @@ -52,7 +53,11 @@ type SubagentQuotaPrimeFn = (config: OcxConfig, reason: string) => Promise export type SubagentPoolAccountPreview = ( modelId: string | undefined, now: number, + modelEligibleAccountIds?: ReadonlySet, ) => string | null; +export type SubagentModelEligibleAccountIds = ( + modelId: string | undefined, +) => ReadonlySet | undefined; let subagentQuotaPrimeForTests: SubagentQuotaPrimeFn | null = null; let quotaPrimeInFlight: Promise | null = null; @@ -184,10 +189,11 @@ function resolveRouteFallbackAccountId( accountId?: string | null, now = Date.now(), poolAccountPreview?: SubagentPoolAccountPreview, + modelEligibleAccountIds?: ReadonlySet, ): string | null { if (route?.codexAccountId !== undefined) return route.codexAccountId; if (route && isPoolCodexRoute(route) && poolAccountPreview) { - return poolAccountPreview(route.modelId, now); + return poolAccountPreview(route.modelId, now, modelEligibleAccountIds); } return resolvePoolFallbackAccountId(config, accountId); } @@ -249,17 +255,26 @@ export function isSubagentModelUnavailable( now = Date.now(), accountUsabilityOptions?: CodexAccountUsabilityOptions, poolAccountPreview?: SubagentPoolAccountPreview, + modelEligibleAccountIdsForModel?: SubagentModelEligibleAccountIds, ): 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; + const modelEligibleAccountIds = modelEligibleAccountIdsForModel?.(route.modelId); + const candidateAccountUsabilityOptions = modelEligibleAccountIds !== undefined + ? { + ...accountUsabilityOptions, + modelEligibleAccountIds, + } + : accountUsabilityOptions; const resolvedAccountId = resolveRouteFallbackAccountId( route, config, accountId, now, poolAccountPreview, + candidateAccountUsabilityOptions?.modelEligibleAccountIds, ); if (isModelHealthBlocked(model, config, resolvedAccountId, now)) return true; if (!isPoolCodexRoute(route)) return false; @@ -268,7 +283,7 @@ export function isSubagentModelUnavailable( // route (canonical openai defaults to pool even when codexAccountMode is omitted). if (!resolvedAccountId) return true; if (isCodexAccountPaused(config, resolvedAccountId)) return true; - if (!isCodexAccountUsable(config, resolvedAccountId, 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 @@ -298,8 +313,10 @@ export function selectAvailableSubagentModel( accountUsabilityOptions?: CodexAccountUsabilityOptions, trailingFallback: readonly string[] = [], poolAccountPreview?: SubagentPoolAccountPreview, + modelEligibleAccountIdsForModel?: SubagentModelEligibleAccountIds, + resolvedChain?: readonly string[], ): { model: string; rewritten: boolean; skipped: string[] } { - const chain = normalizedChain(primary, config, extraFallback, trailingFallback); + const chain = resolvedChain ?? normalizedChain(primary, config, extraFallback, trailingFallback); const skipped: string[] = []; for (const candidate of chain) { if (nativeFallbackOnly) { @@ -316,6 +333,7 @@ export function selectAvailableSubagentModel( now, accountUsabilityOptions, poolAccountPreview, + modelEligibleAccountIdsForModel, )) { skipped.push(candidate); continue; @@ -546,28 +564,26 @@ export function applySubagentModelFallback( nativeFallbackOnly = false, accountUsabilityOptions?: CodexAccountUsabilityOptions, poolAccountPreview?: SubagentPoolAccountPreview, + modelEligibleAccountIdsForModel?: SubagentModelEligibleAccountIds, + resolvedFallbackChain?: readonly string[] | null, ): { from?: string; to?: string; skipped?: string[] } | null { if (!isThreadSpawnRequest(headers)) return 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; + const fallbackChain = resolvedFallbackChain === undefined + ? resolveSubagentFallbackChain(parsed, config) + : resolvedFallbackChain; + if (!fallbackChain) return null; const selection = selectAvailableSubagentModel( parsed.modelId, config, - configuredFallback, + [], accountId, now, nativeFallbackOnly, accountUsabilityOptions, - tomlRoleFallback, + [], poolAccountPreview, + modelEligibleAccountIdsForModel, + fallbackChain, ); if (!selection.rewritten) return selection.skipped.length > 0 ? { from: parsed.modelId, to: parsed.modelId, skipped: selection.skipped } @@ -577,6 +593,37 @@ export function applySubagentModelFallback( return { from, to: selection.model, skipped: selection.skipped }; } +/** 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 the effective fallback chain crosses an account-gated native Pool model. */ +export function subagentFallbackNeedsModelEntitlements( + fallbackChain: readonly string[] | null, + config: OcxConfig, +): boolean { + return fallbackChain?.some((candidate) => { + const route = tryRouteFallbackModel(config, candidate); + return !!route + && isPoolCodexRoute(route) + && ACCOUNT_GATED_NATIVE_OPENAI_MODELS.has(route.modelId); + }) === true; +} + export function subagentFallbackGuidanceText(config: OcxConfig): string { const chain = config.subagentModelFallback ?? []; if (chain.length === 0) return ""; diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index 627a427563..f9a5c204b9 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -149,6 +149,7 @@ import { resolveCodexModelEntitlements, } from "../../codex/model-entitlements"; import { ACCOUNT_GATED_NATIVE_OPENAI_MODELS } from "../../codex/catalog/native-models"; +import { MAIN_CODEX_ACCOUNT_ID } from "../../codex/main-account"; import { captureCodexAffinityDiagnostic } from "../../codex/affinity-debug"; import { computeQuotaCooldown, @@ -217,6 +218,9 @@ import { applySubagentModelFallback, maybePrimeSubagentQuota, recordSubagentQuotaFailureForThreadSpawn, + resolveSubagentFallbackChain, + subagentFallbackNeedsModelEntitlements, + type SubagentModelEligibleAccountIds, type SubagentPoolAccountPreview, } from "../../codex/subagent-model-fallback"; import { isNativeMainTrafficBlocked } from "../../codex/native-profile-startup"; @@ -1291,6 +1295,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; @@ -1590,6 +1596,7 @@ async function resolveResponsesCodexAuth( modelId: route.modelId, substituteMainCredentialForDirect: substituteMainCredential, beginCodexAccountSelection: codexAccountSelectionForTurn(options.turnAdmissionLease), + resolveCodexModelEntitlements: options.resolveCodexModelEntitlements, }); options.onCodexAuthContextResolved?.(authCtx); } else { @@ -1628,6 +1635,25 @@ async function resolveResponsesCodexAuth( } } +async function resolveSubagentFallbackModelEligibility(args: { + config: OcxConfig; + fallbackChain: readonly string[] | null; + nativeMainReadsForbidden: boolean; + resolver: typeof resolveCodexModelEntitlements; +}): Promise { + if (!subagentFallbackNeedsModelEntitlements(args.fallbackChain, args.config)) return undefined; + const excludeAccountIds = args.nativeMainReadsForbidden + ? new Set([MAIN_CODEX_ACCOUNT_ID]) + : undefined; + const snapshot = await args.resolver(args.config, { excludeAccountIds }); + return (modelId) => { + const entitledAccountIds = entitledCodexAccountIdsForModel(snapshot, modelId); + return entitledAccountIds + ? new Set([...entitledAccountIds].filter(accountId => !excludeAccountIds?.has(accountId))) + : undefined; + }; +} + /** * Apply every route-dependent request mutation against the final selected route. * Must run only after subagent fallback has settled the model/provider. @@ -2435,7 +2461,12 @@ 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 initialSubagentFallbackChain = threadSpawn && !options.comboAttempt + ? resolveSubagentFallbackChain(parsed, config) + : null; + const previewSelectionAdmission = threadSpawn + && !options.comboAttempt + && (route.codexAccountId === undefined || initialSubagentFallbackChain !== null) ? codexAccountSelectionForTurn(options.turnAdmissionLease)?.() : undefined; const nativeMainRecoveryBlocked = isNativeMainTrafficBlocked(); @@ -2448,6 +2479,7 @@ async function handleResponsesInner( let selectedForwardHeaders = req.headers; let subagentFallbackAccountId = config.activeCodexAccountId ?? null; let subagentFallbackAccountPreview: SubagentPoolAccountPreview | undefined; + let subagentFallbackModelEligibleAccountIdsForModel: SubagentModelEligibleAccountIds | undefined; let subagentQuotaFailureModel = parsed.modelId; const parentThreadId = req.headers.get("x-codex-parent-thread-id")?.trim() ?? null; const poolAffinityKey = codexPoolAffinityKey(req.headers) ?? null; @@ -2465,20 +2497,35 @@ async function handleResponsesInner( // 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 + && (route.codexAccountId === undefined || initialSubagentFallbackChain !== null) + ) { // 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 fallbackChain = initialSubagentFallbackChain; + subagentFallbackModelEligibleAccountIdsForModel = await resolveSubagentFallbackModelEligibility({ + config, + fallbackChain, + nativeMainReadsForbidden, + resolver: options.resolveCodexModelEntitlements ?? resolveCodexModelEntitlements, + }); const fallbackNow = Date.now(); - subagentFallbackAccountPreview = (modelId, previewNow) => previewCodexAccountForRequest( + subagentFallbackAccountPreview = (modelId, previewNow, modelEligibleAccountIds) => previewCodexAccountForRequest( poolAffinityKey, config, previewNow, codexQuotaScopeForModel(modelId), - previewSelectionOptions, + { ...previewSelectionOptions, modelEligibleAccountIds }, + ); + const previewAccountId = route.codexAccountId ?? subagentFallbackAccountPreview( + route.modelId, + fallbackNow, + subagentFallbackModelEligibleAccountIdsForModel?.(route.modelId), ); - const previewAccountId = subagentFallbackAccountPreview(route.modelId, fallbackNow); subagentFallbackAccountId = previewAccountId ?? config.activeCodexAccountId ?? null; const fallback = applySubagentModelFallback( parsed, @@ -2489,6 +2536,8 @@ async function handleResponsesInner( unreadableEncryptedAgentTask, previewSelectionOptions, subagentFallbackAccountPreview, + subagentFallbackModelEligibleAccountIdsForModel, + fallbackChain, ); if (fallback) { (logCtx as unknown as Record).subagentModelFallbackFrom = fallback.from; diff --git a/tests/codex-model-entitlements.test.ts b/tests/codex-model-entitlements.test.ts index 017922f029..ca4e631da4 100644 --- a/tests/codex-model-entitlements.test.ts +++ b/tests/codex-model-entitlements.test.ts @@ -6,11 +6,10 @@ import { isDirectCallerEntitledToCodexModel, resetCodexModelEntitlementCacheForTests, resolveCodexModelEntitlements, - cachedAvailableAccountGatedNativeModels, - seedCodexModelEntitlementsForTests, seedCodexModelEntitlementsForTests, type CodexModelEntitlementCredentialSnapshot, } from "../src/codex/model-entitlements"; +import { MAIN_CODEX_ACCOUNT_ID } from "../src/codex/main-account"; const DAYBREAK = "gpt-daybreak-blue-latest"; const SOL = "gpt-5.6-sol"; @@ -80,6 +79,40 @@ describe("Codex account model entitlements", () => { expect(entitledCodexAccountIdsForModel(snapshot, DAYBREAK)?.size).toBe(0); }); + test("filters excluded accounts before credential and roster access", async () => { + const credentialReads: string[] = []; + const fetchedAccounts: string[] = []; + const snapshot = await resolveCodexModelEntitlements({ + codexAccounts: [ + { id: "pool-b", email: "pool-b@example.test", isMain: false }, + ], + }, { + excludeAccountIds: new Set([MAIN_CODEX_ACCOUNT_ID]), + credentialSnapshot: async (accountId) => { + credentialReads.push(accountId); + return credential(accountId); + }, + fetcher: (async (_input, init) => { + fetchedAccounts.push(new Headers(init?.headers).get("chatgpt-account-id") ?? ""); + return roster(DAYBREAK); + }) as typeof fetch, + now: 1_000, + }); + + expect(credentialReads).toEqual(["pool-b"]); + expect(fetchedAccounts).toEqual(["chatgpt-pool-b"]); + expect([...snapshot.modelsByAccount.keys()]).toEqual(["pool-b"]); + expect(snapshot.confirmedAccountIds.has(MAIN_CODEX_ACCOUNT_ID)).toBe(false); + + const supplied = await resolveCodexModelEntitlements({ codexAccounts: [] }, { + credentials: [credential(MAIN_CODEX_ACCOUNT_ID), credential("pool-c")], + excludeAccountIds: new Set([MAIN_CODEX_ACCOUNT_ID]), + fetcher: (async () => roster(DAYBREAK)) as typeof fetch, + now: 2_000, + }); + expect([...supplied.modelsByAccount.keys()]).toEqual(["pool-c"]); + }); + test("checks a Direct caller's own bearer instead of a local Pool account", async () => { let seenAuthorization = ""; let seenAccount = ""; @@ -137,24 +170,4 @@ describe("Codex account model entitlements", () => { expect([...cachedAvailableAccountGatedNativeModels(1_000)]).toContain(DAYBREAK); }); - test("Direct-caller rosters do not evict main/Pool entitlement evidence", async () => { - // The catalog projects ONLY from main/Pool keys. Under a single shared LRU, a burst of - // distinct Direct callers pushed those out and the gated row vanished from the catalog - // until rediscovery — fail-closed flapping whose cause an operator cannot see. - seedCodexModelEntitlementsForTests("main", [DAYBREAK], 1_000); - expect([...cachedAvailableAccountGatedNativeModels(1_000)]).toContain(DAYBREAK); - - // Far more distinct Direct callers than the per-class cache bound of 64. - for (let i = 0; i < 80; i += 1) { - await isDirectCallerEntitledToCodexModel( - new Headers({ authorization: `Bearer caller-${i}` }), - DAYBREAK, - { fetcher: (async () => roster(DAYBREAK)) as typeof fetch, now: 1_000 }, - ); - } - - // With one shared 64-entry LRU this read came back empty. The main grant is a different - // eviction class and is still inside its TTL, so it must survive. - expect([...cachedAvailableAccountGatedNativeModels(1_000)]).toContain(DAYBREAK); - }); }); diff --git a/tests/subagent-fallback-handle-responses.test.ts b/tests/subagent-fallback-handle-responses.test.ts index 001d182a19..d9c2e34964 100644 --- a/tests/subagent-fallback-handle-responses.test.ts +++ b/tests/subagent-fallback-handle-responses.test.ts @@ -723,6 +723,234 @@ describe("subagent fallback final-route normalization", () => { }); describe("native fallback account preview", () => { + test("fallback preview and final auth use their own entitlement snapshots", 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" }, + ], + }); + const { noteSubagentModelFailure } = await import("../src/codex/subagent-model-fallback"); + noteSubagentModelFailure("team/gpt-5.6-sol", "429", cfg, "pool-a", now); + // Make a stale account choice observable at the fallback boundary: preview must + // apply the first snapshot and move to pool-b before checking model health. + noteSubagentModelFailure("gpt-daybreak-blue-latest", "429", cfg, "pool-a", now); + const previewSnapshot = { + 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(), + }; + const finalSnapshot = { + 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; + 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 entitlementCalls === 1 ? previewSnapshot : finalSnapshot; + }, + }, + logCtx, + ); + + expect(response.status).toBe(200); + expect(entitlementCalls).toBe(2); + // Final auth is authoritative and sees the second snapshot, not the preview snapshot. + expect(finalAuth).toMatchObject({ kind: "pool", accountId: "pool-a" }); + expect((logCtx as unknown as Record).subagentModelFallbackTo) + .toBe("gpt-daybreak-blue-latest"); + expect(capture.auths[0]).toContain("pool-a_token"); + }); + + test("entitlement discovery holds and releases preview admission on rejection", async () => { + const now = 1_800_000_000_000; + Date.now = () => now; + const cfg = poolNativePlusRoutedConfig({ + activeCodexAccountId: "pool-a", + autoSwitchThreshold: 0, + codexAccountNamespaces: { team: "pool-a" }, + subagentModelFallback: ["gpt-daybreak-blue-latest"], + }); + let beginCount = 0; + let releaseCount = 0; + let resolverCalls = 0; + let rejectDiscovery!: (reason: Error) => void; + const discovery = new Promise((_resolve, reject) => { rejectDiscovery = reject; }); + const turnAdmissionLease = { + release() {}, + beginCodexAccountSelection() { + beginCount += 1; + return { + mainProfileDraining: false, + claimMainProfile: () => true, + release: () => { releaseCount += 1; }, + }; + }, + } satisfies Pick; + let fetchCalls = 0; + globalThis.fetch = (async () => { + fetchCalls += 1; + throw new Error("must not dispatch"); + }) as typeof fetch; + + const pending = postSpawn( + cfg, + { model: "team/gpt-5.6-sol", input: readableAgentInput(), stream: false }, + { + turnAdmissionLease, + resolveCodexModelEntitlements: async () => { + resolverCalls += 1; + return discovery; + }, + }, + ); + for (let i = 0; i < 20 && resolverCalls === 0; i += 1) await Promise.resolve(); + + expect(resolverCalls).toBe(1); + expect(beginCount).toBe(1); + expect(releaseCount).toBe(0); + expect(fetchCalls).toBe(0); + + rejectDiscovery(new Error("entitlement discovery unavailable")); + await expect(pending).rejects.toThrow("entitlement discovery unavailable"); + expect(releaseCount).toBe(1); + expect(fetchCalls).toBe(0); + }); + + test("unentitled fixed gated primary falls through to a routed fallback", async () => { + const now = 1_800_000_000_000; + Date.now = () => now; + installPoolCredential("pool-a", "pool_acc_a", now); + installPoolCredential("pool-b", "pool_acc_b", now); + const cfg = poolNativePlusRoutedConfig({ + activeCodexAccountId: "pool-a", + autoSwitchThreshold: 0, + codexAccountNamespaces: { restricted: "pool-b" }, + subagentModelFallback: ["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 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(), + }; + const logCtx: RequestLogContext = { model: "", provider: "" }; + const capture = { urls: [] as string[], bodies: [] as string[], auths: [] as Array }; + mockUpstream(capture); + + const response = await postSpawn( + cfg, + { model: "restricted/gpt-daybreak-blue-latest", input: readableAgentInput(), stream: false }, + { resolveCodexModelEntitlements: async () => entitlementSnapshot }, + logCtx, + ); + + expect(response.status).toBe(200); + expect((logCtx as unknown as Record).subagentModelFallbackTo).toBe("xai/grok-4.5"); + expect(capture.urls).toHaveLength(1); + expect(capture.urls[0]).toContain("api.x.ai"); + expect(capture.bodies[0]).not.toContain("gpt-daybreak-blue-latest"); + }); + + 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("Desktop fallback affinity drives the subagent preview and final native account", async () => { const now = 1_800_000_000_000; Date.now = () => now; diff --git a/tests/subagent-model-fallback.test.ts b/tests/subagent-model-fallback.test.ts index 98d83a6ad5..f37a07ef4a 100644 --- a/tests/subagent-model-fallback.test.ts +++ b/tests/subagent-model-fallback.test.ts @@ -216,7 +216,7 @@ describe("subagent model fallback chain", () => { }); }); - test("fixed account candidates do not call the Pool preview", () => { + test("fixed account candidates preserve selectors and enforce per-model entitlements", () => { updateAccountQuota("account-a", 10, undefined, 20); const config = cfg({ codexAccountNamespaces: { team: "account-a" } }); const throwingPreview = () => { @@ -228,11 +228,74 @@ describe("subagent model fallback chain", () => { config, "pool-a", Date.now(), + { modelEligibleAccountIds: new Set(["account-a"]) }, + throwingPreview, + () => undefined, + )).toBe(false); + expect(isSubagentModelUnavailable( + "team/gpt-daybreak-blue-latest", + config, + "pool-a", + Date.now(), + undefined, + throwingPreview, + () => new Set(["account-b"]), + )).toBe(true); + expect(isSubagentModelUnavailable( + "team/gpt-daybreak-blue-latest", + config, + "pool-a", + Date.now(), undefined, throwingPreview, + () => new Set(["account-a"]), )).toBe(false); }); + test("unqualified gated candidates pass their entitlement set into Pool preview", () => { + const now = 1_800_000_000_000; + const config = cfg({ + autoSwitchThreshold: 0, + codexAccountNamespaces: { team: "account-a" }, + subagentModelFallback: ["gpt-daybreak-blue-latest", "xai/grok-4.5"], + }); + updateAccountQuota("account-a", 10, undefined, 20); + updateAccountQuota("account-b", 10, undefined, 20); + noteSubagentModelFailure("team/gpt-5.6-sol", "429", config, "account-a", now); + + const previews: Array<{ modelId: string | undefined; eligible: string[] | undefined }> = []; + const selected = selectAvailableSubagentModel( + "team/gpt-5.6-sol", + config, + [], + "account-a", + now, + false, + undefined, + [], + (modelId, _previewNow, eligibleAccountIds) => { + previews.push({ + modelId, + eligible: eligibleAccountIds ? [...eligibleAccountIds] : undefined, + }); + return eligibleAccountIds?.has("account-b") ? "account-b" : "account-a"; + }, + modelId => modelId === "gpt-daybreak-blue-latest" + ? new Set(["account-b"]) + : undefined, + ); + + expect(selected).toEqual({ + model: "gpt-daybreak-blue-latest", + rewritten: true, + skipped: ["team/gpt-5.6-sol"], + }); + expect(previews).toEqual([{ + modelId: "gpt-daybreak-blue-latest", + eligible: ["account-b"], + }]); + }); + 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"] });