fix(codex): accept monthly-classified snapshots and probe-owned refresh generations (#955 review)

src/codex/quota.ts

@@ -103,10 +103,17 @@ export function isCompleteCodexQuotaRecoverySnapshot(
   // Recovery still fails closed on MISSING EVIDENCE — a credits-only or windowless payload
   // carries no usage reading at all and must never clear a cooldown. What it does not do is
   // fail closed on an unfamiliar plan NAME, which only ever meant "cooled forever".
-  const required = codexQuotaWindowForPlan(plan) === "monthly"
-    ? quota.monthlyPercent
-    : quota.weeklyPercent;
-  return typeof required === "number" && Number.isFinite(required);
+  //
+  // The parser classifies windows by DURATION, not by plan name: a Team response whose
+  // primary window is explicitly monthly parses to monthlyPercent only (no secondary
+  // window exists), so requiring weeklyPercent because the plan is not go/free would
+  // strand exactly those accounts until their predicted expiry. Accept whichever window(s)
+  // the parser actually wrote; Go/Free never carry a weekly value, so monthly-only is
+  // required there.
+  if (codexQuotaWindowForPlan(plan) === "monthly") {
+    return typeof quota.monthlyPercent === "number" && Number.isFinite(quota.monthlyPercent);
+  }
+  return hasKnownQuotaValue(quota);
 }
 
 export function normalizeUsagePercent(value: unknown): number | undefined {

src/codex/routing.ts

@@ -139,6 +139,8 @@ export type CodexQuotaRecoveryProbeClaim = {
   leaseId: string;
   cooldownGeneration: number;
   credentialGeneration: number;
+  /** Claim-time `replacedAt`; unchanged after a probe-owned refresh, stamped on external replacement. */
+  credentialReplacedAt?: number;
 };
 
 export type CodexQuotaRecoveryProbeProof = {
@@ -439,6 +441,7 @@ export function claimDueCodexQuotaRecoveryProbes(
     scope?: CodexQuotaScope;
     health: CodexUpstreamHealth;
     credentialGeneration: number;
+    credentialReplacedAt?: number;
     order: number;
   }> = [];
   for (const [order, account] of (config.codexAccounts ?? []).entries()) {
@@ -467,6 +470,7 @@ export function claimDueCodexQuotaRecoveryProbes(
       ...(candidate.scope ? { scope: candidate.scope } : {}),
       health: candidate.health,
       credentialGeneration: record.generation,
+      ...(record.replacedAt !== undefined ? { credentialReplacedAt: record.replacedAt } : {}),
       order,
     });
   }
@@ -491,6 +495,9 @@ export function claimDueCodexQuotaRecoveryProbes(
       leaseId,
       cooldownGeneration: candidate.health.cooldownGeneration ?? 0,
       credentialGeneration: candidate.credentialGeneration,
+      ...(candidate.credentialReplacedAt !== undefined
+        ? { credentialReplacedAt: candidate.credentialReplacedAt }
+        : {}),
     };
   });
 }
@@ -506,10 +513,21 @@ export function settleCodexQuotaRecoveryProbe(
     ? scopedHealthFor(claim.accountId, claim.scope)
     : upstreamHealth.get(claim.accountId);
   if (!health || health.probeLeaseId !== claim.leaseId) return false;
+  const currentRecord = readCodexAccountRecord(claim.accountId);
+  const proofGeneration = proof.credentialGeneration;
+  // A probe-owned token refresh (getValidCodexToken) advances the credential generation by
+  // exactly one while preserving `replacedAt`; an external credential replacement bumps the
+  // generation too but stamps a fresh `replacedAt`. Accept the +1 transition only when the
+  // claim-time lineage is intact AND the generation the fresh quota was proven under is live.
+  const generationFenced = proofGeneration !== undefined
+    && (proofGeneration === claim.credentialGeneration
+      ? isCodexAccountGenerationLive(claim.accountId, proofGeneration)
+      : proofGeneration === claim.credentialGeneration + 1
+        && currentRecord?.replacedAt === claim.credentialReplacedAt
+        && isCodexAccountGenerationLive(claim.accountId, proofGeneration));
   const fenced = (health.cooldownGeneration ?? 0) === claim.cooldownGeneration
     && (health.probeLeaseGeneration ?? 0) === claim.cooldownGeneration
-    && claim.credentialGeneration === proof.credentialGeneration
-    && isCodexAccountGenerationLive(claim.accountId, claim.credentialGeneration);
+    && generationFenced;
   if (!recovered || !fenced) {
     const released = withProbeLeaseReleased(health, now);
     if (claim.scope) setScopedHealth(claim.accountId, claim.scope, released);

tests/codex-cooldown-recovery.test.ts

@@ -117,6 +117,55 @@ describe("Codex cooldown recovery worker", () => {
     expect(routed).toEqual(["b", "b"]);
   });
 
+  test("recovers a Team account from a duration-classified monthly snapshot", async () => {
+    // WHAM can legitimately return only an explicitly monthly primary window for a Team
+    // plan (30.4-day window, no secondary). parseUsageQuota then writes monthlyPercent only,
+    // so recovery must accept the window the parser actually classified instead of demanding
+    // a weekly reading because the plan name is not "go"/"free".
+    const config = makeConfig(["a"]);
+    saveCredential("a");
+    cool(config, "a");
+    globalThis.fetch = async () => usageResponse(0, {
+      plan_type: "team",
+      rate_limit: {
+        primary_window: { used_percent: 6, reset_at: 1_900_000_000, limit_window_seconds: 2_628_000 },
+        secondary_window: null,
+        tertiary_window: null,
+      },
+      rate_limit_reset_credits: { available_count: 0 },
+    });
+    await runCodexCooldownRecoveryProbes(config, due());
+    expect(getCodexQuotaHealthSnapshot("a", "shared", due() + 1)).toBeNull();
+  });
+
+  test("recovers when the probe's own token refresh advances the credential generation", async () => {
+    // A near-expiry access token makes getValidCodexToken() refresh it inside the probe
+    // fetch, bumping the credential generation from 1 to 2 before WHAM completes. The fresh
+    // quota is proven under the new live generation, so settling against the claim-time
+    // generation must accept this probe's own refresh, not treat it as a replacement.
+    const config = makeConfig(["a"]);
+    saveCodexAccountCredential("a", {
+      accessToken: "access-a",
+      refreshToken: "refresh-a",
+      expiresAt: Date.now() + 30_000,
+      chatgptAccountId: "acct-a",
+    });
+    cool(config, "a");
+    globalThis.fetch = async input => {
+      const url = typeof input === "string" ? input : input instanceof URL ? input.href : input.url;
+      if (url.includes("/oauth/token")) {
+        return new Response(JSON.stringify({
+          access_token: "access-a-2",
+          refresh_token: "refresh-a-2",
+          expires_in: 3600,
+        }), { status: 200, headers: { "Content-Type": "application/json" } });
+      }
+      return usageResponse(12);
+    };
+    await runCodexCooldownRecoveryProbes(config, due());
+    expect(getCodexQuotaHealthSnapshot("a", "shared", due() + 1)).toBeNull();
+  });
+
   test.each([
     ["still exhausted", () => usageResponse(100)],
     ["credits only", () => usageResponse(0, { plan_type: "team", rate_limit_reset_credits: { available_count: 1 } })],
@@ -290,13 +339,18 @@ describe("Codex cooldown recovery worker", () => {
 
     for (const plan of snapshotPlans) {
       const monthly = codexQuotaWindowForPlan(plan) === "monthly";
-      const filled = monthly ? { monthlyPercent: 12 } : { weeklyPercent: 12 };
-      const empty = monthly ? { weeklyPercent: 12 } : { monthlyPercent: 12 };
-      expect(isCompleteCodexQuotaRecoverySnapshot(filled, plan)).toBe(true);
-      // The other window is not evidence for this plan, in either direction.
-      expect(isCompleteCodexQuotaRecoverySnapshot(empty, plan)).toBe(false);
+      // The parser classifies windows by duration, so a weekly-billed plan can carry a
+      // monthly-only reading (30-day primary, no secondary). Any window the parser actually
+      // wrote is evidence; only Go/Free never carry a weekly value.
+      expect(isCompleteCodexQuotaRecoverySnapshot({ weeklyPercent: 12 }, plan)).toBe(!monthly);
+      expect(isCompleteCodexQuotaRecoverySnapshot({ monthlyPercent: 12 }, plan)).toBe(true);
     }
 
+    // Monthly-billed Go/Free parse to monthlyPercent only; a weekly-only reading is not
+    // evidence for them.
+    expect(isCompleteCodexQuotaRecoverySnapshot({ weeklyPercent: 12 }, "go")).toBe(false);
+    expect(isCompleteCodexQuotaRecoverySnapshot({ weeklyPercent: 12 }, "free")).toBe(false);
+
     // Absent plan follows the parser's weekly default.
     expect(isCompleteCodexQuotaRecoverySnapshot({ weeklyPercent: 12 }, undefined)).toBe(true);
     // Missing EVIDENCE still fails closed — that is the guard that matters.