From a4350d01d256f6dc7edc8db47061f3b6ac1b7f8d Mon Sep 17 00:00:00 2001 From: luvs01 Date: Sun, 30 Aug 2026 20:57:25 +0900 Subject: [PATCH 1/4] fix(codex): throttle repeated failed pool quota primes A pool account whose WHAM lookup fails stores no quota, so it stays "unknown" and every later prime trigger re-selects it as stale and repeats the same failing request. Successful lookups are already bounded by POOL_CACHE_TTL; failures had no backoff at all. Record the last prime attempt per account and give a failed lookup the same TTL window. The record is keyed by credential generation, so a re-authentication, refresh, or account removal retries immediately instead of waiting out a backoff earned by the previous credential. --- src/codex/auth-api.ts | 32 +++++++++++++- tests/codex-quota-prime.test.ts | 78 +++++++++++++++++++++++++++++++++ 2 files changed, 109 insertions(+), 1 deletion(-) diff --git a/src/codex/auth-api.ts b/src/codex/auth-api.ts index 91734a3229..18ec156016 100644 --- a/src/codex/auth-api.ts +++ b/src/codex/auth-api.ts @@ -1243,6 +1243,16 @@ async function fetchPoolAccountQuota(accountId: string, forceRefresh = false, co } let primeInFlight: Promise | null = null; +/** + * Last prime attempt per pool account. A failed WHAM lookup stores no quota, so + * without this the account stays "unknown" and every later prime trigger re-selects + * it as stale and repeats the same failing request. Successful lookups are already + * throttled by their stored updatedAt; this gives failures the same TTL backoff. + * + * Keyed by credential generation so a re-authentication, refresh, or account removal + * retries immediately instead of waiting out a backoff earned by the old credential. + */ +const poolQuotaPrimeAttemptedAt = new Map(); let cooldownRecoveryInFlight: Promise | null = null; export async function runCodexCooldownRecoveryProbes(config: OcxConfig, now = Date.now()): Promise { @@ -1331,7 +1341,15 @@ export async function primeCodexPoolQuotas( const pool = (runtimeConfig.codexAccounts ?? []).filter(isSelectableCodexPoolAccount); const stale = pool.filter(a => { const q = getAccountQuota(a.id); - return !q || Date.now() - q.updatedAt >= POOL_CACHE_TTL; + if (q) return Date.now() - q.updatedAt >= POOL_CACHE_TTL; + // No stored quota: either never primed, or the last attempt failed. Retry only + // once per TTL window so an unreachable or rejecting account cannot turn every + // prime trigger into another upstream request. + const lastAttempt = poolQuotaPrimeAttemptedAt.get(a.id); + if (!lastAttempt) return true; + // A newer credential invalidates the previous failure: retry without waiting. + if (lastAttempt.generation !== readCodexAccountRecord(a.id)?.generation) return true; + return Date.now() - lastAttempt.at >= POOL_CACHE_TTL; }); const primeMain = async () => { const mainLease = tryAcquireNativeMainPrimeLease(); @@ -1359,6 +1377,10 @@ export async function primeCodexPoolQuotas( primeMain(), mapWithConcurrency(stale, POOL_QUOTA_REFRESH_CONCURRENCY, async a => { if (!getCodexAccountCredential(a.id)) return; + poolQuotaPrimeAttemptedAt.set(a.id, { + generation: readCodexAccountRecord(a.id)?.generation ?? 0, + at: Date.now(), + }); await fetchPoolAccountQuota(a.id, false, a.plan); }), ]); @@ -1376,6 +1398,14 @@ export async function primeCodexPoolQuotas( * from another suite cannot coalesce into the next prime. */ export function clearCodexQuotaPrimeState(): void { primeInFlight = null; + poolQuotaPrimeAttemptedAt.clear(); +} + +/** Test-only: drop the shared single-flight promise while keeping the per-account + * failure backoff, so a test can trigger a second real prime pass and still observe + * the throttle a production caller would see. */ +export function clearCodexQuotaPrimeSingleFlightForTests(): void { + primeInFlight = null; } /** Test-only reset for the worker-level single-flight. */ diff --git a/tests/codex-quota-prime.test.ts b/tests/codex-quota-prime.test.ts index 01f02f19b2..4b6fdcc7d8 100644 --- a/tests/codex-quota-prime.test.ts +++ b/tests/codex-quota-prime.test.ts @@ -7,6 +7,7 @@ import { updateAccountQuota, clearAccountQuota, clearCodexQuotaPrimeState, + clearCodexQuotaPrimeSingleFlightForTests, clearMainAccountInfoCache, } from "../src/codex/auth-api"; import { saveCodexAccountCredential } from "../src/codex/account-store"; @@ -400,6 +401,83 @@ describe("primeCodexPoolQuotas", () => { } }); + test("a failed pool quota fetch is throttled for the rest of the TTL window", async () => { + const config = makeConfig(); + seedPoolAccount(config, "p1"); + const originalFetch = globalThis.fetch; + let calls = 0; + try { + // Upstream is unavailable, so no quota is ever stored for this account. The + // account therefore stays "unknown" and, without an attempt record, every + // later prime re-selects it as stale and re-issues the same failing fetch. + globalThis.fetch = async (input: RequestInfo | URL) => { + if (String(input).includes("/backend-api/wham/usage")) { + calls += 1; + return new Response("upstream unavailable", { status: 503 }); + } + return originalFetch(input); + }; + await primeCodexPoolQuotas(config, "test"); + expect(calls).toBe(1); + expect(getAccountQuota("p1")).toBeNull(); + + // Only the single-flight promise is dropped between passes; the throttle state + // must survive so a later trigger does not repeat the failing lookup. + clearCodexQuotaPrimeSingleFlightForTests(); + await primeCodexPoolQuotas(config, "test"); + clearCodexQuotaPrimeSingleFlightForTests(); + await primeCodexPoolQuotas(config, "test"); + + // A failed lookup must back off for the same POOL_CACHE_TTL window that a + // successful one gets, instead of retrying on every prime trigger. + expect(calls).toBe(1); + expect(getAccountQuota("p1")).toBeNull(); + } finally { + globalThis.fetch = originalFetch; + } + }); + + test("re-authenticating a failed account retries without waiting out the backoff", async () => { + const config = makeConfig(); + seedPoolAccount(config, "p1"); + const originalFetch = globalThis.fetch; + let calls = 0; + let upstreamHealthy = false; + try { + globalThis.fetch = async (input: RequestInfo | URL) => { + if (String(input).includes("/backend-api/wham/usage")) { + calls += 1; + return upstreamHealthy ? whamResponse(20) : new Response("down", { status: 503 }); + } + return originalFetch(input); + }; + await primeCodexPoolQuotas(config, "test"); + expect(calls).toBe(1); + expect(getAccountQuota("p1")).toBeNull(); + + // Throttled while the same credential keeps failing. + clearCodexQuotaPrimeSingleFlightForTests(); + await primeCodexPoolQuotas(config, "test"); + expect(calls).toBe(1); + + // A re-authentication bumps the credential generation, which must invalidate the + // backoff earned by the old credential instead of hiding a now-usable account. + upstreamHealthy = true; + saveCodexAccountCredential("p1", { + accessToken: "access-p1-renewed", + refreshToken: "refresh-p1-renewed", + expiresAt: Date.now() + 5 * 60_000, + chatgptAccountId: "acct-p1", + }); + clearCodexQuotaPrimeSingleFlightForTests(); + await primeCodexPoolQuotas(config, "test"); + expect(calls).toBe(2); + expect(getAccountQuota("p1")).toMatchObject({ weeklyPercent: 20 }); + } finally { + globalThis.fetch = originalFetch; + } + }); + test("one blocked account does not sink the rest", async () => { const config = makeConfig(); seedPoolAccount(config, "ok"); From 7ce7db776f433a171875c27aba586dc65fbedadd Mon Sep 17 00:00:00 2001 From: luvs01 Date: Sun, 30 Aug 2026 23:20:59 +0900 Subject: [PATCH 2/4] fix(codex): bind quota prime backoff to admitted probe --- src/codex/auth-api.ts | 34 ++++++++++++-- tests/codex-quota-prime.test.ts | 81 ++++++++++++++++++++++++++++++++- 2 files changed, 109 insertions(+), 6 deletions(-) diff --git a/src/codex/auth-api.ts b/src/codex/auth-api.ts index 18ec156016..f5444b69c9 100644 --- a/src/codex/auth-api.ts +++ b/src/codex/auth-api.ts @@ -1377,11 +1377,35 @@ export async function primeCodexPoolQuotas( primeMain(), mapWithConcurrency(stale, POOL_QUOTA_REFRESH_CONCURRENCY, async a => { if (!getCodexAccountCredential(a.id)) return; - poolQuotaPrimeAttemptedAt.set(a.id, { - generation: readCodexAccountRecord(a.id)?.generation ?? 0, - at: Date.now(), - }); - await fetchPoolAccountQuota(a.id, false, a.plan); + const attemptedAt = Date.now(); + const startGeneration = readCodexAccountRecord(a.id)?.generation ?? 0; + try { + const result = await fetchPoolAccountQuota(a.id, false, a.plan); + // Credential admission can settle without sending WHAM (for example while + // another refresh owns the grant). Do not turn that local deferral into a + // five-minute upstream backoff. + if (result.quotaProbeSkipped) return; + poolQuotaPrimeAttemptedAt.set(a.id, { + // getValidCodexToken may rotate the credential before WHAM is sent. + // Bind the backoff to the generation that actually made the request; + // otherwise the next prime sees a false generation change and retries + // the same failed WHAM call immediately. + generation: result.credentialGeneration ?? startGeneration, + at: attemptedAt, + }); + } catch (error) { + // Global quota-flight admission rejected this account before any WHAM + // request existed. Leave it immediately eligible for the next prime. + if (error instanceof PoolQuotaProbeBusyError) return; + // Unexpected failures are still bounded, but only against the credential + // whose attempt started; a concurrent replacement remains immediately + // eligible through the generation comparison above. + poolQuotaPrimeAttemptedAt.set(a.id, { + generation: startGeneration, + at: attemptedAt, + }); + throw error; + } }), ]); } catch { diff --git a/tests/codex-quota-prime.test.ts b/tests/codex-quota-prime.test.ts index 4b6fdcc7d8..a9eb921a6b 100644 --- a/tests/codex-quota-prime.test.ts +++ b/tests/codex-quota-prime.test.ts @@ -9,8 +9,9 @@ import { clearCodexQuotaPrimeState, clearCodexQuotaPrimeSingleFlightForTests, clearMainAccountInfoCache, + seedCodexAuthAdmissionForTests, } from "../src/codex/auth-api"; -import { saveCodexAccountCredential } from "../src/codex/account-store"; +import { readCodexAccountRecord, saveCodexAccountCredential } from "../src/codex/account-store"; import { resetMainCodexAccountIdentityTrackingForTests } from "../src/codex/account-lifecycle"; import { MAIN_CODEX_ACCOUNT_ID } from "../src/codex/main-account"; import { @@ -478,6 +479,84 @@ describe("primeCodexPoolQuotas", () => { } }); + test("an admission-busy prime does not back off an account it never probed", async () => { + const config = makeConfig(); + seedPoolAccount(config, "p1"); + const originalFetch = globalThis.fetch; + const releaseAdmission = seedCodexAuthAdmissionForTests({ quotaFlights: 16 }); + let whamCalls = 0; + try { + globalThis.fetch = async (input: RequestInfo | URL) => { + if (String(input).includes("/backend-api/wham/usage")) { + whamCalls += 1; + return whamResponse(20); + } + return originalFetch(input); + }; + await primeCodexPoolQuotas(config, "test"); + expect(whamCalls).toBe(0); + expect(getAccountQuota("p1")).toBeNull(); + + releaseAdmission(); + clearCodexQuotaPrimeSingleFlightForTests(); + await primeCodexPoolQuotas(config, "test"); + + expect(whamCalls).toBe(1); + expect(getAccountQuota("p1")).toMatchObject({ weeklyPercent: 20 }); + } finally { + releaseAdmission(); + globalThis.fetch = originalFetch; + } + }); + + test("a refreshed credential keeps the backoff earned by its failed WHAM request", async () => { + const config = makeConfig(); + seedPoolAccount(config, "p1"); + saveCodexAccountCredential("p1", { + accessToken: "expiring-p1", + refreshToken: "refresh-p1", + expiresAt: Date.now() + 30_000, + chatgptAccountId: "acct-p1", + }); + const startGeneration = readCodexAccountRecord("p1")?.generation; + const originalFetch = globalThis.fetch; + let oauthCalls = 0; + let whamCalls = 0; + try { + globalThis.fetch = async (input: RequestInfo | URL) => { + const url = String(input); + if (url.includes("/oauth/token")) { + oauthCalls += 1; + return Response.json({ + access_token: "fresh-p1", + refresh_token: "fresh-refresh-p1", + expires_in: 3600, + }); + } + if (url.includes("/backend-api/wham/usage")) { + whamCalls += 1; + return new Response("upstream unavailable", { status: 503 }); + } + return originalFetch(input); + }; + + await primeCodexPoolQuotas(config, "test"); + expect(oauthCalls).toBe(1); + expect(whamCalls).toBe(1); + expect(readCodexAccountRecord("p1")?.generation).toBe((startGeneration ?? 0) + 1); + expect(getAccountQuota("p1")).toBeNull(); + + clearCodexQuotaPrimeSingleFlightForTests(); + await primeCodexPoolQuotas(config, "test"); + + expect(oauthCalls).toBe(1); + expect(whamCalls).toBe(1); + expect(getAccountQuota("p1")).toBeNull(); + } finally { + globalThis.fetch = originalFetch; + } + }); + test("one blocked account does not sink the rest", async () => { const config = makeConfig(); seedPoolAccount(config, "ok"); From a78b410930a1757b03950df91f348ee2ed418d86 Mon Sep 17 00:00:00 2001 From: luvs01 Date: Tue, 1 Sep 2026 15:44:48 +0900 Subject: [PATCH 3/4] fix(codex): prove quota prime dispatch before backoff --- src/codex/auth-api.ts | 126 ++++++++++++------ tests/codex-quota-prime.test.ts | 218 +++++++++++++++++++++++++++++++- 2 files changed, 306 insertions(+), 38 deletions(-) diff --git a/src/codex/auth-api.ts b/src/codex/auth-api.ts index f5444b69c9..aabc3b0a43 100644 --- a/src/codex/auth-api.ts +++ b/src/codex/auth-api.ts @@ -843,6 +843,23 @@ interface PoolQuotaResult { /** Present only when this call's WHAM response included `rate_limit_reset_credits.available_count`. */ freshResetCredits?: number; quotaProbeSkipped?: true; + /** Positive evidence captured immediately before an upstream WHAM dispatch. */ + quotaProbeAttempted?: { at: number; credentialGeneration: number }; +} + +interface PoolQuotaProbeEvidence { + attempted?: NonNullable; +} + +function markQuotaProbeAttempted(evidence: PoolQuotaProbeEvidence, credentialGeneration: number): void { + evidence.attempted = { at: Date.now(), credentialGeneration }; +} + +function withQuotaProbeEvidence( + result: PoolQuotaResult, + evidence: PoolQuotaProbeEvidence, +): PoolQuotaResult { + return evidence.attempted ? { ...result, quotaProbeAttempted: evidence.attempted } : result; } interface PoolQuotaRefreshFlight { @@ -988,6 +1005,7 @@ async function recoverPoolQuotaFrom401(ctx: { rejectedAccessToken: string; rejectedGeneration: number; resp: Response; + quotaProbeEvidence: PoolQuotaProbeEvidence; onCredentialGeneration?: (generation: number) => void; }): Promise { const { accountId, existing, configuredPlan, rejectedAccessToken, rejectedGeneration, resp } = ctx; @@ -1061,6 +1079,7 @@ async function recoverPoolQuotaFrom401(ctx: { ctx.onCredentialGeneration?.(refreshed.generation); const writerGeneration = captureConfigGeneration(); + markQuotaProbeAttempted(ctx.quotaProbeEvidence, refreshed.generation); const replay = await fetch("https://chatgpt.com/backend-api/wham/usage", { headers: { Authorization: `Bearer ${refreshed.accessToken}`, @@ -1147,57 +1166,75 @@ async function fetchFreshPoolAccountQuota( existing: StoredAccountQuota | null, configuredPlan?: string, onCredentialGeneration?: (generation: number) => void, + getValidToken: typeof getValidCodexToken = getValidCodexToken, ): Promise { const writerGeneration = captureConfigGeneration(); let requestCredentialGeneration = readCodexAccountRecord(accountId)?.generation; + const quotaProbeEvidence: PoolQuotaProbeEvidence = {}; try { - const { accessToken, chatgptAccountId, generation } = await getValidCodexToken(accountId); + const { accessToken, chatgptAccountId, generation } = await getValidToken(accountId); requestCredentialGeneration = generation; onCredentialGeneration?.(generation); + markQuotaProbeAttempted(quotaProbeEvidence, generation); const resp = await fetch("https://chatgpt.com/backend-api/wham/usage", { headers: { Authorization: `Bearer ${accessToken}`, "ChatGPT-Account-Id": chatgptAccountId }, signal: AbortSignal.timeout(8000), }); if (!resp.ok) { if (resp.status !== 401) { - return { quota: existing ?? null, needsReauth: false, credentialGeneration: generation }; + return withQuotaProbeEvidence( + { quota: existing ?? null, needsReauth: false, credentialGeneration: generation }, + quotaProbeEvidence, + ); } // A bare 401 is what a stale-but-refreshable bearer produces after a plan change, so // quarantining on it tells the operator to re-authenticate an account that was fine // (#3019). Refresh once, replay once, and only then decide. - return await recoverPoolQuotaFrom401({ + const recovered = await recoverPoolQuotaFrom401({ accountId, existing, configuredPlan, rejectedAccessToken: accessToken, rejectedGeneration: generation, resp, + quotaProbeEvidence, onCredentialGeneration, }); + return withQuotaProbeEvidence(recovered, quotaProbeEvidence); } - return await commitPoolQuotaResponse(resp, { + const committed = await commitPoolQuotaResponse(resp, { accountId, existing, configuredPlan, generation, writerGeneration, }); + return withQuotaProbeEvidence(committed, quotaProbeEvidence); } catch (e) { if (e instanceof CodexCredentialGenerationConflictError || e instanceof CodexCredentialRefreshLockTimeoutError || e instanceof CodexCredentialRefreshBusyError || e instanceof CodexCredentialRefreshStaleError) { - return { + return withQuotaProbeEvidence({ quota: existing ?? null, needsReauth: false, credentialGeneration: requestCredentialGeneration, - ...(e instanceof CodexCredentialRefreshBusyError || e instanceof CodexCredentialRefreshStaleError - ? { quotaProbeSkipped: true as const } - : {}), - }; + quotaProbeSkipped: true, + }, quotaProbeEvidence); } if (e instanceof TokenRefreshError) { - return { quota: existing ?? null, needsReauth: true, credentialGeneration: requestCredentialGeneration }; + return withQuotaProbeEvidence( + { quota: existing ?? null, needsReauth: true, credentialGeneration: requestCredentialGeneration }, + quotaProbeEvidence, + ); } - return { quota: existing ?? null, needsReauth: false, credentialGeneration: requestCredentialGeneration }; + return withQuotaProbeEvidence( + { quota: existing ?? null, needsReauth: false, credentialGeneration: requestCredentialGeneration }, + quotaProbeEvidence, + ); } } -async function fetchPoolAccountQuota(accountId: string, forceRefresh = false, configuredPlan?: string): Promise { +async function fetchPoolAccountQuota( + accountId: string, + forceRefresh = false, + configuredPlan?: string, + getValidToken: typeof getValidCodexToken = getValidCodexToken, +): Promise { const existing = getAccountQuota(accountId); if (!forceRefresh && existing && Date.now() - existing.updatedAt < POOL_CACHE_TTL) { return { @@ -1227,6 +1264,7 @@ async function fetchPoolAccountQuota(accountId: string, forceRefresh = false, co existing, configuredPlan, generation => { state.resolvedCredentialGeneration = generation; }, + getValidToken, ); const flight: PoolQuotaRefreshFlight = { state, promise: refresh }; const activeFlights = flights ?? new Set(); @@ -1304,6 +1342,19 @@ export interface PrimeCodexPoolQuotasOptions { fetchMainInfo?: typeof fetchMainAccountInfo; } +let getValidPoolTokenForPrime = getValidCodexToken; + +/** Test-only: inject a deterministic pre-dispatch credential outcome for quota priming. */ +export function setCodexPoolQuotaTokenResolverForTests( + resolver: typeof getValidCodexToken, +): () => void { + const previous = getValidPoolTokenForPrime; + getValidPoolTokenForPrime = resolver; + return () => { + if (getValidPoolTokenForPrime === resolver) getValidPoolTokenForPrime = previous; + }; +} + function tryAcquireNativeMainPrimeLease(): AdmissionLease | null { return tryAcquireNativeMainProfileClaim(); } @@ -1335,9 +1386,13 @@ export async function primeCodexPoolQuotas( || !isCanonicalOpenAiForwardProvider(openai) || providerCodexAccountMode(OPENAI_CODEX_PROVIDER_ID, openai) !== "pool" ) return; + const runtimeConfig = getRuntimeConfig(config); + const configuredPoolIds = new Set((runtimeConfig.codexAccounts ?? []).map(account => account.id)); + for (const accountId of poolQuotaPrimeAttemptedAt.keys()) { + if (!configuredPoolIds.has(accountId)) poolQuotaPrimeAttemptedAt.delete(accountId); + } if (primeInFlight) return primeInFlight; primeInFlight = (async () => { - const runtimeConfig = getRuntimeConfig(config); const pool = (runtimeConfig.codexAccounts ?? []).filter(isSelectableCodexPoolAccount); const stale = pool.filter(a => { const q = getAccountQuota(a.id); @@ -1377,35 +1432,31 @@ export async function primeCodexPoolQuotas( primeMain(), mapWithConcurrency(stale, POOL_QUOTA_REFRESH_CONCURRENCY, async a => { if (!getCodexAccountCredential(a.id)) return; - const attemptedAt = Date.now(); - const startGeneration = readCodexAccountRecord(a.id)?.generation ?? 0; + let result: PoolQuotaResult; try { - const result = await fetchPoolAccountQuota(a.id, false, a.plan); - // Credential admission can settle without sending WHAM (for example while - // another refresh owns the grant). Do not turn that local deferral into a - // five-minute upstream backoff. - if (result.quotaProbeSkipped) return; - poolQuotaPrimeAttemptedAt.set(a.id, { - // getValidCodexToken may rotate the credential before WHAM is sent. - // Bind the backoff to the generation that actually made the request; - // otherwise the next prime sees a false generation change and retries - // the same failed WHAM call immediately. - generation: result.credentialGeneration ?? startGeneration, - at: attemptedAt, - }); + result = await fetchPoolAccountQuota(a.id, false, a.plan, getValidPoolTokenForPrime); } catch (error) { - // Global quota-flight admission rejected this account before any WHAM - // request existed. Leave it immediately eligible for the next prime. + // Local quota-flight saturation proves no WHAM request existed for this account. + // Consume it per item so sibling workers remain inside the shared prime lifetime. if (error instanceof PoolQuotaProbeBusyError) return; - // Unexpected failures are still bounded, but only against the credential - // whose attempt started; a concurrent replacement remains immediately - // eligible through the generation comparison above. - poolQuotaPrimeAttemptedAt.set(a.id, { - generation: startGeneration, - at: attemptedAt, - }); throw error; } + // Only the data-plane function knows whether upstream dispatch began. Any + // cache hit, credential deferral, or local admission failure remains eligible. + const attempted = result.quotaProbeAttempted; + if (!attempted) return; + if (!configuredPoolAccount(getRuntimeConfig(config), a.id)) { + poolQuotaPrimeAttemptedAt.delete(a.id); + return; + } + poolQuotaPrimeAttemptedAt.set(a.id, { + // getValidCodexToken may rotate the credential before WHAM is sent. + // Bind the backoff to the generation that actually made the request; + // otherwise the next prime sees a false generation change and retries + // the same failed WHAM call immediately. + generation: attempted.credentialGeneration, + at: attempted.at, + }); }), ]); } catch { @@ -1423,6 +1474,7 @@ export async function primeCodexPoolQuotas( export function clearCodexQuotaPrimeState(): void { primeInFlight = null; poolQuotaPrimeAttemptedAt.clear(); + getValidPoolTokenForPrime = getValidCodexToken; } /** Test-only: drop the shared single-flight promise while keeping the per-account diff --git a/tests/codex-quota-prime.test.ts b/tests/codex-quota-prime.test.ts index a9eb921a6b..8b5f78661d 100644 --- a/tests/codex-quota-prime.test.ts +++ b/tests/codex-quota-prime.test.ts @@ -10,8 +10,14 @@ import { clearCodexQuotaPrimeSingleFlightForTests, clearMainAccountInfoCache, seedCodexAuthAdmissionForTests, + setCodexPoolQuotaTokenResolverForTests, } from "../src/codex/auth-api"; -import { readCodexAccountRecord, saveCodexAccountCredential } from "../src/codex/account-store"; +import { + CodexCredentialGenerationConflictError, + CodexCredentialRefreshLockTimeoutError, + readCodexAccountRecord, + saveCodexAccountCredential, +} from "../src/codex/account-store"; import { resetMainCodexAccountIdentityTrackingForTests } from "../src/codex/account-lifecycle"; import { MAIN_CODEX_ACCOUNT_ID } from "../src/codex/main-account"; import { @@ -438,6 +444,129 @@ describe("primeCodexPoolQuotas", () => { } }); + test("a real failed pool quota probe becomes eligible after the TTL expires", async () => { + const originalNow = Date.now; + let now = 1_800_000_000_000; + Date.now = () => now; + const config = makeConfig(); + seedPoolAccount(config, "p1"); + saveCodexAccountCredential("p1", { + accessToken: "access-p1-long-lived", + refreshToken: "refresh-p1-long-lived", + expiresAt: now + 60 * 60_000, + chatgptAccountId: "acct-p1", + }); + const originalFetch = globalThis.fetch; + let calls = 0; + let upstreamHealthy = false; + try { + globalThis.fetch = async (input: RequestInfo | URL) => { + if (String(input).includes("/backend-api/wham/usage")) { + calls += 1; + return upstreamHealthy ? whamResponse(20) : new Response("upstream unavailable", { status: 503 }); + } + return originalFetch(input); + }; + + await primeCodexPoolQuotas(config, "test"); + expect(calls).toBe(1); + + now += 5 * 60_000 - 1; + clearCodexQuotaPrimeSingleFlightForTests(); + await primeCodexPoolQuotas(config, "test"); + expect(calls).toBe(1); + + now += 1; + upstreamHealthy = true; + clearCodexQuotaPrimeSingleFlightForTests(); + await primeCodexPoolQuotas(config, "test"); + expect(calls).toBe(2); + expect(getAccountQuota("p1")).toMatchObject({ weeklyPercent: 20 }); + } finally { + globalThis.fetch = originalFetch; + Date.now = originalNow; + } + }); + + test("removing an account from the pool purges its failed-prime backoff", async () => { + const config = makeConfig(); + seedPoolAccount(config, "p1"); + const originalPool = [...(config.codexAccounts ?? [])]; + const originalFetch = globalThis.fetch; + let calls = 0; + let upstreamHealthy = false; + try { + globalThis.fetch = async (input: RequestInfo | URL) => { + if (String(input).includes("/backend-api/wham/usage")) { + calls += 1; + return upstreamHealthy ? whamResponse(20) : new Response("upstream unavailable", { status: 503 }); + } + return originalFetch(input); + }; + + await primeCodexPoolQuotas(config, "test"); + expect(calls).toBe(1); + + config.codexAccounts = []; + clearCodexQuotaPrimeSingleFlightForTests(); + await primeCodexPoolQuotas(config, "removed"); + + config.codexAccounts = originalPool; + upstreamHealthy = true; + clearCodexQuotaPrimeSingleFlightForTests(); + await primeCodexPoolQuotas(config, "restored"); + expect(calls).toBe(2); + expect(getAccountQuota("p1")).toMatchObject({ weeklyPercent: 20 }); + } finally { + globalThis.fetch = originalFetch; + } + }); + + test("a late failed probe cannot restore backoff for an account removed in flight", async () => { + const config = makeConfig(); + seedPoolAccount(config, "p1"); + const originalPool = [...(config.codexAccounts ?? [])]; + const originalFetch = globalThis.fetch; + let calls = 0; + let upstreamHealthy = false; + let releaseFirst!: () => void; + const firstDispatched = new Promise(resolve => { releaseFirst = resolve; }); + let finishFirst!: () => void; + const firstGate = new Promise(resolve => { finishFirst = resolve; }); + try { + globalThis.fetch = async (input: RequestInfo | URL) => { + if (String(input).includes("/backend-api/wham/usage")) { + calls += 1; + if (calls === 1) { + releaseFirst(); + await firstGate; + return new Response("upstream unavailable", { status: 503 }); + } + return upstreamHealthy ? whamResponse(20) : new Response("upstream unavailable", { status: 503 }); + } + return originalFetch(input); + }; + + const firstPrime = primeCodexPoolQuotas(config, "test"); + await firstDispatched; + config.codexAccounts = []; + const removedPrime = primeCodexPoolQuotas(config, "removed"); + finishFirst(); + await Promise.all([firstPrime, removedPrime]); + + config.codexAccounts = originalPool; + upstreamHealthy = true; + clearCodexQuotaPrimeSingleFlightForTests(); + await primeCodexPoolQuotas(config, "restored"); + + expect(calls).toBe(2); + expect(getAccountQuota("p1")).toMatchObject({ weeklyPercent: 20 }); + } finally { + finishFirst(); + globalThis.fetch = originalFetch; + } + }); + test("re-authenticating a failed account retries without waiting out the backoff", async () => { const config = makeConfig(); seedPoolAccount(config, "p1"); @@ -509,6 +638,50 @@ describe("primeCodexPoolQuotas", () => { } }); + test.each([ + ["credential generation conflict", () => new CodexCredentialGenerationConflictError()], + ["refresh-lock timeout", () => new CodexCredentialRefreshLockTimeoutError()], + ] as const)("a %s before dispatch does not back off the next prime", async (_label, makeError) => { + const config = makeConfig(); + seedPoolAccount(config, "p1"); + const originalFetch = globalThis.fetch; + let tokenAttempts = 0; + let whamCalls = 0; + try { + globalThis.fetch = async (input: RequestInfo | URL) => { + if (String(input).includes("/backend-api/wham/usage")) { + whamCalls += 1; + return whamResponse(20); + } + return originalFetch(input); + }; + const getValidPoolToken = async () => { + tokenAttempts += 1; + if (tokenAttempts === 1) throw makeError(); + return { + accessToken: "access-p1", + chatgptAccountId: "acct-p1", + generation: readCodexAccountRecord("p1")!.generation, + }; + }; + const restoreTokenResolver = setCodexPoolQuotaTokenResolverForTests(getValidPoolToken); + + try { + await primeCodexPoolQuotas(config, "test"); + clearCodexQuotaPrimeSingleFlightForTests(); + await primeCodexPoolQuotas(config, "test"); + } finally { + restoreTokenResolver(); + } + + expect(tokenAttempts).toBe(2); + expect(whamCalls).toBe(1); + expect(getAccountQuota("p1")).toMatchObject({ weeklyPercent: 20 }); + } finally { + globalThis.fetch = originalFetch; + } + }); + test("a refreshed credential keeps the backoff earned by its failed WHAM request", async () => { const config = makeConfig(); seedPoolAccount(config, "p1"); @@ -557,6 +730,49 @@ describe("primeCodexPoolQuotas", () => { } }); + test("a failed 401 replay binds backoff to the replay credential generation", async () => { + const config = makeConfig(); + seedPoolAccount(config, "p1"); + const originalFetch = globalThis.fetch; + let oauthCalls = 0; + let whamCalls = 0; + try { + globalThis.fetch = async (input: RequestInfo | URL) => { + const url = String(input); + if (url.includes("/oauth/token")) { + oauthCalls += 1; + return Response.json({ + access_token: "fresh-p1", + refresh_token: "fresh-refresh-p1", + expires_in: 3600, + }); + } + if (url.includes("/backend-api/wham/usage")) { + whamCalls += 1; + if (whamCalls === 1) { + return Response.json({ error: { code: "transient_edge_rejection" } }, { status: 401 }); + } + throw new Error("replay transport unavailable"); + } + return originalFetch(input); + }; + + await primeCodexPoolQuotas(config, "test"); + expect(oauthCalls).toBe(1); + expect(whamCalls).toBe(2); + expect(getAccountQuota("p1")).toBeNull(); + + clearCodexQuotaPrimeSingleFlightForTests(); + await primeCodexPoolQuotas(config, "test"); + + expect(oauthCalls).toBe(1); + expect(whamCalls).toBe(2); + expect(getAccountQuota("p1")).toBeNull(); + } finally { + globalThis.fetch = originalFetch; + } + }); + test("one blocked account does not sink the rest", async () => { const config = makeConfig(); seedPoolAccount(config, "ok"); From 80101192ba35b9ca71d47dedc392b92cc440a15f Mon Sep 17 00:00:00 2001 From: jun Date: Wed, 2 Sep 2026 02:25:26 +0900 Subject: [PATCH 4/4] fix(codex): prune removed-account prime markers before the eligibility return The prune sat after the provider-eligibility early return, so an account removed while the provider was disabled kept its failed-attempt marker. Restoring the same id inside POOL_CACHE_TTL then read that stale failure as current and skipped the retry the restored credential was entitled to. The existing coverage removed an account with the provider enabled, which is why this survived review. Add the disabled-window case. --- src/codex/auth-api.ts | 15 +++++++---- tests/codex-quota-prime.test.ts | 46 +++++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+), 5 deletions(-) diff --git a/src/codex/auth-api.ts b/src/codex/auth-api.ts index aabc3b0a43..4f9c5efd28 100644 --- a/src/codex/auth-api.ts +++ b/src/codex/auth-api.ts @@ -1380,17 +1380,22 @@ export async function primeCodexPoolQuotas( options: PrimeCodexPoolQuotasOptions = {}, ): Promise { const openai = config.providers[OPENAI_CODEX_PROVIDER_ID]; + // Prune attempt markers for accounts that no longer exist BEFORE the eligibility + // return. A removal that happens while the provider is disabled or out of pool mode + // would otherwise leave a stale failure marker behind; restoring the same account id + // within POOL_CACHE_TTL would then read that old failure as current and skip the + // retry the restored credential is entitled to. + const runtimeConfig = getRuntimeConfig(config); + const configuredPoolIds = new Set((runtimeConfig.codexAccounts ?? []).map(account => account.id)); + for (const accountId of poolQuotaPrimeAttemptedAt.keys()) { + if (!configuredPoolIds.has(accountId)) poolQuotaPrimeAttemptedAt.delete(accountId); + } if ( !openai || openai.disabled === true || !isCanonicalOpenAiForwardProvider(openai) || providerCodexAccountMode(OPENAI_CODEX_PROVIDER_ID, openai) !== "pool" ) return; - const runtimeConfig = getRuntimeConfig(config); - const configuredPoolIds = new Set((runtimeConfig.codexAccounts ?? []).map(account => account.id)); - for (const accountId of poolQuotaPrimeAttemptedAt.keys()) { - if (!configuredPoolIds.has(accountId)) poolQuotaPrimeAttemptedAt.delete(accountId); - } if (primeInFlight) return primeInFlight; primeInFlight = (async () => { const pool = (runtimeConfig.codexAccounts ?? []).filter(isSelectableCodexPoolAccount); diff --git a/tests/codex-quota-prime.test.ts b/tests/codex-quota-prime.test.ts index 8b5f78661d..d9da22e691 100644 --- a/tests/codex-quota-prime.test.ts +++ b/tests/codex-quota-prime.test.ts @@ -522,6 +522,52 @@ describe("primeCodexPoolQuotas", () => { } }); + test("removal purges the backoff even while the provider is disabled", async () => { + // The prune used to sit AFTER the provider-eligibility early return, so a removal + // that happened while the provider was disabled (or out of pool mode) left the + // stale failure marker in place. Restoring the same account id within + // POOL_CACHE_TTL then read that old failure as current and skipped the retry the + // restored credential is entitled to. + const config = makeConfig(); + seedPoolAccount(config, "p1"); + const originalPool = [...(config.codexAccounts ?? [])]; + const originalFetch = globalThis.fetch; + let calls = 0; + let upstreamHealthy = false; + try { + globalThis.fetch = async (input: RequestInfo | URL) => { + if (String(input).includes("/backend-api/wham/usage")) { + calls += 1; + return upstreamHealthy ? whamResponse(20) : new Response("upstream unavailable", { status: 503 }); + } + return originalFetch(input); + }; + + // One failed prime records the backoff marker. + await primeCodexPoolQuotas(config, "test"); + expect(calls).toBe(1); + + // The account is removed while the provider is disabled: the prime returns early, + // but the marker must still be pruned. + config.codexAccounts = []; + config.providers.openai!.disabled = true; + clearCodexQuotaPrimeSingleFlightForTests(); + await primeCodexPoolQuotas(config, "removed-while-disabled"); + expect(calls).toBe(1); + + // Restored and re-enabled inside the TTL window: the prime must dispatch now. + config.codexAccounts = originalPool; + config.providers.openai!.disabled = false; + upstreamHealthy = true; + clearCodexQuotaPrimeSingleFlightForTests(); + await primeCodexPoolQuotas(config, "restored"); + expect(calls).toBe(2); + expect(getAccountQuota("p1")).toMatchObject({ weeklyPercent: 20 }); + } finally { + globalThis.fetch = originalFetch; + } + }); + test("a late failed probe cannot restore backoff for an account removed in flight", async () => { const config = makeConfig(); seedPoolAccount(config, "p1");