fix(codex): throttle repeated failed pool quota primes

src/codex/auth-api.ts

@@ -1074,6 +1074,16 @@ async function fetchPoolAccountQuota(accountId: string, forceRefresh = false, co
 }
 
 let primeInFlight: Promise<void> | 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<string, { generation: number; at: number }>();
 let cooldownRecoveryInFlight: Promise<void> | null = null;
 
 export async function runCodexCooldownRecoveryProbes(config: OcxConfig, now = Date.now()): Promise<void> {
@@ -1162,7 +1172,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();
@@ -1190,7 +1208,35 @@ export async function primeCodexPoolQuotas(
         primeMain(),
         mapWithConcurrency(stale, POOL_QUOTA_REFRESH_CONCURRENCY, async a => {
           if (!getCodexAccountCredential(a.id)) return;
-          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 {
@@ -1207,6 +1253,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. */

tests/codex-quota-prime.test.ts

@@ -7,9 +7,11 @@ import {
   updateAccountQuota,
   clearAccountQuota,
   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 {
@@ -400,6 +402,161 @@ 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("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");