From 79d2164e261a24c750df86b6a650cd8824dbd233 Mon Sep 17 00:00:00 2001 From: Yuxin Qiao <104957188+Yuxin-Qiao@users.noreply.github.com> Date: Tue, 4 Aug 2026 04:45:46 +0800 Subject: [PATCH] fix(codex): accept monthly-classified recovery snapshots and probe-owned refresh generations Addresses the two unresolved Codex review threads on #955: - isCompleteCodexQuotaRecoverySnapshot() required weeklyPercent for every non-Go/Free plan by plan name, but the parser classifies windows by duration: a Team response with an explicitly monthly primary window parses to monthlyPercent only, so those accounts could never recover early and stayed cooled until their predicted expiry. - settleCodexQuotaRecoveryProbe() required the claim-time credential generation to match exactly. A probe-owned token refresh inside getValidCodexToken() advances the generation by one before WHAM completes, so a successful fresh reading was rejected and the account waited another probe interval. replacedAt is preserved by refresh and stamped by external replacement, so it fences the +1 transition. --- src/codex/quota.ts | 15 +++++-- src/codex/routing.ts | 22 ++++++++- tests/codex-cooldown-recovery.test.ts | 64 ++++++++++++++++++++++++--- 3 files changed, 90 insertions(+), 11 deletions(-) diff --git a/src/codex/quota.ts b/src/codex/quota.ts index 739baf986..2d397a51e 100644 --- a/src/codex/quota.ts +++ b/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 { diff --git a/src/codex/routing.ts b/src/codex/routing.ts index 3331bc7d0..3d62c9048 100644 --- a/src/codex/routing.ts +++ b/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); diff --git a/tests/codex-cooldown-recovery.test.ts b/tests/codex-cooldown-recovery.test.ts index 7ce26fd19..1c71975e2 100644 --- a/tests/codex-cooldown-recovery.test.ts +++ b/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.