From bf5e518db159dd332bfe69cfffc693f483ee1502 Mon Sep 17 00:00:00 2001 From: luvs01 Date: Sun, 30 Aug 2026 20:57:25 +0900 Subject: [PATCH 1/3] 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 cadffe8613912c0cff38d0a42298a9ccdde32d07 Mon Sep 17 00:00:00 2001 From: luvs01 Date: Sun, 30 Aug 2026 23:20:59 +0900 Subject: [PATCH 2/3] 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 7c9ce6cf72272031bacc1ef6b78fc52f78a01144 Mon Sep 17 00:00:00 2001 From: luvs01 Date: Tue, 1 Sep 2026 15:44:48 +0900 Subject: [PATCH 3/3] 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");