From 25fbdd10dfb0f4f237ab6a5da4405dc0ce4e0ffe 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/routing.ts | 9 + src/codex/subagent-model-fallback.ts | 102 ++++++++- src/server/responses/core.ts | 49 +++- ...subagent-fallback-handle-responses.test.ts | 210 ++++++++++++++++++ tests/subagent-model-fallback.test.ts | 137 ++++++++++++ 5 files changed, 489 insertions(+), 18 deletions(-) diff --git a/src/codex/routing.ts b/src/codex/routing.ts index fa6fa63d83..10160fe913 100644 --- a/src/codex/routing.ts +++ b/src/codex/routing.ts @@ -612,6 +612,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. diff --git a/src/codex/subagent-model-fallback.ts b/src/codex/subagent-model-fallback.ts index da40ebffff..db2bdf7971 100644 --- a/src/codex/subagent-model-fallback.ts +++ b/src/codex/subagent-model-fallback.ts @@ -15,12 +15,12 @@ import { CODEX_HOME, getCodexHome } from "./paths"; import { CODEX_UNKNOWN_USAGE_SCORE, getAccountQuota } from "./quota"; import { canAcquireCodexQuotaProbeLease, + canAcquireCodexQuotaScopeProbeLease, codexQuotaScopeForModel, computeCodexUsageScore, getCodexQuotaHealthSnapshot, getEffectiveActiveCodexAccountId, getPoolAccountPlan, - isCodexAccountInCooldown, } from "./routing"; import { isCodexAccountUsable, @@ -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, @@ -48,6 +49,15 @@ 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; let subagentQuotaPrimeForTests: SubagentQuotaPrimeFn | null = null; let quotaPrimeInFlight: Promise | null = null; @@ -177,8 +187,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 +248,62 @@ 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, ): 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 { + const quotaScope = codexQuotaScopeForModel(route.modelId); + const cooldown = getCodexQuotaHealthSnapshot(resolvedAccountId, quotaScope, now); + if (cooldown !== null) { + const probeAvailable = cooldown.quotaScope + ? canAcquireCodexQuotaScopeProbeLease(resolvedAccountId, cooldown.quotaScope, now) + : canAcquireCodexQuotaProbeLease(resolvedAccountId, now); + if (!probeAvailable) return true; + } } - return isNativeModelQuotaExhausted(model, config, accountId, now); + return isNativeModelQuotaExhausted(model, config, resolvedAccountId, now); } export function selectAvailableSubagentModel( @@ -275,6 +315,8 @@ export function selectAvailableSubagentModel( nativeFallbackOnly = false, accountUsabilityOptions?: CodexAccountUsabilityOptions, trailingFallback: readonly string[] = [], + poolAccountPreview?: SubagentPoolAccountPreview, + modelEligibleAccountIdsForModel?: SubagentModelEligibleAccountIds, ): { model: string; rewritten: boolean; skipped: string[] } { const chain = normalizedChain(primary, config, extraFallback, trailingFallback); const skipped: string[] = []; @@ -286,7 +328,15 @@ export function selectAvailableSubagentModel( continue; } } - if (isSubagentModelUnavailable(candidate, config, accountId, now, accountUsabilityOptions)) { + if (isSubagentModelUnavailable( + candidate, + config, + accountId, + now, + accountUsabilityOptions, + poolAccountPreview, + modelEligibleAccountIdsForModel, + )) { skipped.push(candidate); continue; } @@ -515,6 +565,8 @@ export function applySubagentModelFallback( now = Date.now(), nativeFallbackOnly = false, accountUsabilityOptions?: CodexAccountUsabilityOptions, + poolAccountPreview?: SubagentPoolAccountPreview, + modelEligibleAccountIdsForModel?: SubagentModelEligibleAccountIds, ): { from?: string; to?: string; skipped?: string[] } | null { if (!isThreadSpawnRequest(headers)) return null; const tomlRoleFallback = resolveAgentModelFallbackForPrimary( @@ -536,6 +588,8 @@ export function applySubagentModelFallback( nativeFallbackOnly, accountUsabilityOptions, tomlRoleFallback, + poolAccountPreview, + modelEligibleAccountIdsForModel, ); if (!selection.rewritten) return selection.skipped.length > 0 ? { from: parsed.modelId, to: parsed.modelId, skipped: selection.skipped } @@ -545,6 +599,30 @@ export function applySubagentModelFallback( return { from, to: selection.model, skipped: selection.skipped }; } +/** Whether this request's configured fallback chain contains an account-gated native model. */ +export function subagentFallbackNeedsModelEntitlements( + parsed: OcxParsedRequest, + config: OcxConfig, +): boolean { + const tomlRoleFallback = resolveAgentModelFallbackForPrimary( + parsed.modelId, + getCodexHome(), + config.codexAccountNamespaces, + ); + const configuredFallback = resolveConfiguredModelFallbackForPrimary(parsed.modelId, config); + const globalFallback = config.subagentModelFallback ?? []; + if (globalFallback.length === 0 && configuredFallback.length === 0 && tomlRoleFallback.length === 0) { + return false; + } + return normalizedChain(parsed.modelId, config, configuredFallback, tomlRoleFallback) + .some((candidate) => { + const route = tryRouteFallbackModel(config, candidate); + return !!route + && isPoolCodexRoute(route) + && ACCOUNT_GATED_NATIVE_OPENAI_MODELS.has(route.modelId); + }); +} + 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 8e6afce616..fb57268059 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -207,6 +207,9 @@ import { applySubagentModelFallback, maybePrimeSubagentQuota, recordSubagentQuotaFailureForThreadSpawn, + subagentFallbackNeedsModelEntitlements, + type SubagentModelEligibleAccountIds, + type SubagentPoolAccountPreview, } from "../../codex/subagent-model-fallback"; import { isNativeMainTrafficBlocked } from "../../codex/native-profile-startup"; import { @@ -1224,6 +1227,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; @@ -1515,6 +1520,7 @@ async function resolveResponsesCodexAuth( modelId: route.modelId, substituteMainCredentialForDirect: substituteMainCredential, beginCodexAccountSelection: codexAccountSelectionForTurn(options.turnAdmissionLease), + resolveCodexModelEntitlements: options.resolveCodexModelEntitlements, }); options.onCodexAuthContextResolved?.(authCtx); } else { @@ -2361,6 +2367,8 @@ async function handleResponsesInner( let selectedForwardHeaders = req.headers; let subagentFallbackAccountId = config.activeCodexAccountId ?? null; let subagentFallbackPreviewAccountId: string | null | undefined; + let subagentFallbackAccountPreview: SubagentPoolAccountPreview | undefined; + let subagentFallbackModelEligibleAccountIds: SubagentModelEligibleAccountIds | undefined; let subagentQuotaFailureModel = parsed.modelId; const parentThreadId = req.headers.get("x-codex-parent-thread-id")?.trim() ?? null; const poolAffinityKey = codexPoolAffinityKey(req.headers) ?? null; @@ -2374,6 +2382,20 @@ async function handleResponsesInner( await maybePrimeSubagentQuota(config, Date.now(), { nativeMainReadsForbidden }); } + if ( + threadSpawn + && !options.comboAttempt + && route.codexAccountId === undefined + && subagentFallbackNeedsModelEntitlements(parsed, config) + ) { + const entitlementSnapshot = await (options.resolveCodexModelEntitlements + ?? resolveCodexModelEntitlements)(config); + subagentFallbackModelEligibleAccountIds = (modelId) => entitledCodexAccountIdsForModel( + entitlementSnapshot, + modelId, + ); + } + // 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 @@ -2383,12 +2405,22 @@ async function handleResponsesInner( // 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( + const fallbackNow = Date.now(); + subagentFallbackAccountPreview = ( + modelId, + previewNow, + modelEligibleAccountIds, + ) => previewCodexAccountForRequest( poolAffinityKey, config, - Date.now(), - codexQuotaScopeForModel(route.modelId), - previewSelectionOptions, + previewNow, + codexQuotaScopeForModel(modelId), + { ...previewSelectionOptions, modelEligibleAccountIds }, + ); + const previewAccountId = subagentFallbackAccountPreview( + route.modelId, + fallbackNow, + subagentFallbackModelEligibleAccountIds?.(route.modelId), ); subagentFallbackPreviewAccountId = previewAccountId; subagentFallbackAccountId = previewAccountId ?? config.activeCodexAccountId ?? null; @@ -2397,9 +2429,11 @@ async function handleResponsesInner( req.headers, config, previewAccountId, - Date.now(), + fallbackNow, unreadableEncryptedAgentTask, previewSelectionOptions, + subagentFallbackAccountPreview, + subagentFallbackModelEligibleAccountIds, ); if (fallback) { (logCtx as unknown as Record).subagentModelFallbackFrom = fallback.from; @@ -2482,14 +2516,17 @@ 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 fallbackNow = Date.now(); const fallback = applySubagentModelFallback( parsed, req.headers, config, subagentFallbackPreviewAccountId, - Date.now(), + fallbackNow, false, previewSelectionOptions, + subagentFallbackAccountPreview, + subagentFallbackModelEligibleAccountIds, ); if (fallback) { (logCtx as unknown as Record).subagentModelFallbackFrom = fallback.from; diff --git a/tests/subagent-fallback-handle-responses.test.ts b/tests/subagent-fallback-handle-responses.test.ts index 937b0aadeb..6a29531422 100644 --- a/tests/subagent-fallback-handle-responses.test.ts +++ b/tests/subagent-fallback-handle-responses.test.ts @@ -29,10 +29,16 @@ import { } from "../src/codex/subagent-model-fallback"; import { resolveCodexAuthContext, type CodexAuthContext } from "../src/codex/auth-context"; import { handleResponses } from "../src/server/responses"; +import { resetAgentTaskRecoveryState } from "../src/server/responses/agent-task-recovery"; import { isEagerRelaySseResponse } from "../src/server/relay"; 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); @@ -52,6 +58,7 @@ beforeEach(() => { clearCodexUpstreamHealth(); clearAccountQuota(); resetSubagentModelFallbackStateForTests(); + resetAgentTaskRecoveryState(); }); afterEach(() => { @@ -61,6 +68,7 @@ afterEach(() => { clearCodexUpstreamHealth(); clearAccountQuota(); resetSubagentModelFallbackStateForTests(); + resetAgentTaskRecoveryState(); rmSync(testDir, { recursive: true, force: true }); if (previousOpencodexHome === undefined) delete process.env.OPENCODEX_HOME; else process.env.OPENCODEX_HOME = previousOpencodexHome; @@ -775,6 +783,208 @@ describe("native fallback account preview", () => { expect(finalAuth).toMatchObject({ kind: "pool", accountId: "pool-a" }); }); + 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({ + activeCodexAccountId: "pool-a", + 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" }, + ], + }); + 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("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"], + 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); + + 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 + 60_001; + 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 }, + { onCodexAuthContextResolved: (ctx) => { finalAuth = ctx; } }, + { 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(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; diff --git a/tests/subagent-model-fallback.test.ts b/tests/subagent-model-fallback.test.ts index 3d433aef12..35943f6b62 100644 --- a/tests/subagent-model-fallback.test.ts +++ b/tests/subagent-model-fallback.test.ts @@ -17,6 +17,7 @@ import { scanCodexAgentRolesWithTomlModelFallback, selectAvailableSubagentModel, setSubagentQuotaPrimeForTests, + subagentFallbackNeedsModelEntitlements, subagentFallbackGuidanceText, } from "../src/codex/subagent-model-fallback"; import { saveCodexAccountCredential } from "../src/codex/account-store"; @@ -24,6 +25,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 +149,24 @@ 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; + expect(subagentFallbackNeedsModelEntitlements(parsed, cfg({ + subagentModelFallback: ["gpt-daybreak-blue-latest"], + }))).toBe(true); + expect(subagentFallbackNeedsModelEntitlements(parsed, cfg({ + providers: { + ...cfg().providers, + openai: { + adapter: "openai-responses", + baseUrl: "https://chatgpt.com/backend-api/codex", + codexAccountMode: "direct", + }, + }, + subagentModelFallback: ["gpt-daybreak-blue-latest"], + }))).toBe(false); + }); + test("selectAvailableSubagentModel skips quota-exhausted native models", () => { resetSubagentModelFallbackStateForTests(); updateAccountQuota("pool-a", 95, undefined, 20); @@ -215,6 +235,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 +365,75 @@ 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: [], + }); + }); + + 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);