From 77ebee3f76218929be7bb6f653fb7e063cdafca6 Mon Sep 17 00:00:00 2001 From: Ingwannu Date: Tue, 18 Aug 2026 15:52:33 +0000 Subject: [PATCH] fix(codex): retain short-window quota end to end --- src/codex/quota.ts | 31 ++++++++++-- src/codex/routing.ts | 10 ++-- tests/codex-auth-api.test.ts | 91 ++++++++++++++++++++++++++++++++++++ tests/codex-routing.test.ts | 16 +++++++ 4 files changed, 139 insertions(+), 9 deletions(-) diff --git a/src/codex/quota.ts b/src/codex/quota.ts index a4c226db93..104db68d5b 100644 --- a/src/codex/quota.ts +++ b/src/codex/quota.ts @@ -186,7 +186,7 @@ function normalizeResetAt(value: unknown): number | undefined { } function hasKnownQuotaValue(quota: Omit): boolean { - return [quota.weeklyPercent, quota.monthlyPercent] + return [quota.weeklyPercent, quota.monthlyPercent, quota.shortPercent] .some(value => typeof value === "number" && Number.isFinite(value)); } @@ -226,8 +226,14 @@ function snapshotHasMonthly(quota: Omit): boole return quota.monthlyPercent !== undefined || quota.monthlyResetAt !== undefined; } +function snapshotHasShort(quota: Omit): boolean { + return quota.shortPercent !== undefined + || quota.shortResetAt !== undefined + || quota.shortWindowSeconds !== undefined; +} + function snapshotHasUsage(quota: Omit): boolean { - return snapshotHasWeekly(quota) || snapshotHasMonthly(quota); + return snapshotHasWeekly(quota) || snapshotHasMonthly(quota) || snapshotHasShort(quota); } export function setAccountQuotaFromParsed( accountId: string, @@ -246,6 +252,9 @@ export function setAccountQuotaFromParsed( if (existing?.monthlyPercent !== undefined) next.monthlyPercent = existing.monthlyPercent; if (existing?.monthlyResetAt !== undefined) next.monthlyResetAt = existing.monthlyResetAt; if (existing?.monthlyIsPrimaryWindow === true) next.monthlyIsPrimaryWindow = true; + if (existing?.shortPercent !== undefined) next.shortPercent = existing.shortPercent; + if (existing?.shortResetAt !== undefined) next.shortResetAt = existing.shortResetAt; + if (existing?.shortWindowSeconds !== undefined) next.shortWindowSeconds = existing.shortWindowSeconds; next.resetCredits = quota.resetCredits; accountQuota.set(accountId, next); schedulePersistAccountQuotas(); @@ -270,12 +279,25 @@ export function setAccountQuotaFromParsed( // while silently dropping `monthlyIsPrimaryWindow` would look like tertiary-only data to // any future reader, and that failure would be invisible. if (quota.monthlyIsPrimaryWindow === true) next.monthlyIsPrimaryWindow = true; - } else if (snapshotHasWeekly(quota) && existing?.monthlyPercent !== undefined) { + } else if ((snapshotHasWeekly(quota) || snapshotHasShort(quota)) + && existing?.monthlyPercent !== undefined) { next.monthlyPercent = existing.monthlyPercent; if (existing.monthlyResetAt !== undefined) next.monthlyResetAt = existing.monthlyResetAt; if (existing.monthlyIsPrimaryWindow === true) next.monthlyIsPrimaryWindow = true; } + if (snapshotHasShort(quota)) { + if (quota.shortPercent !== undefined) next.shortPercent = quota.shortPercent; + if (quota.shortResetAt !== undefined) next.shortResetAt = quota.shortResetAt; + if (quota.shortWindowSeconds !== undefined) next.shortWindowSeconds = quota.shortWindowSeconds; + } else { + // Header and reset-credit updates are partial snapshots. Preserve the last full WHAM + // burst tuple when those updates do not carry enough window metadata to replace it. + if (existing?.shortPercent !== undefined) next.shortPercent = existing.shortPercent; + if (existing?.shortResetAt !== undefined) next.shortResetAt = existing.shortResetAt; + if (existing?.shortWindowSeconds !== undefined) next.shortWindowSeconds = existing.shortWindowSeconds; + } + if (quota.resetCredits !== undefined) next.resetCredits = quota.resetCredits; else if (existing?.resetCredits !== undefined) next.resetCredits = existing.resetCredits; @@ -369,6 +391,9 @@ export function updateAccountQuota( : {}), ...(existing?.weeklyResetAt !== undefined ? { weeklyResetAt: existing.weeklyResetAt } : {}), ...(existing?.monthlyResetAt !== undefined ? { monthlyResetAt: existing.monthlyResetAt } : {}), + ...(existing?.shortPercent !== undefined ? { shortPercent: existing.shortPercent } : {}), + ...(existing?.shortResetAt !== undefined ? { shortResetAt: existing.shortResetAt } : {}), + ...(existing?.shortWindowSeconds !== undefined ? { shortWindowSeconds: existing.shortWindowSeconds } : {}), ...(existing?.resetCredits !== undefined ? { resetCredits: existing.resetCredits } : {}), updatedAt: Date.now(), }; diff --git a/src/codex/routing.ts b/src/codex/routing.ts index 7e201ea58a..6103432dd0 100644 --- a/src/codex/routing.ts +++ b/src/codex/routing.ts @@ -322,14 +322,12 @@ function deleteScopedHealth(accountId: string, scope: CodexQuotaScope): void { export function computeCodexUsageScore(quota: { weeklyPercent?: number; monthlyPercent?: number; + shortPercent?: number; } | null, plan?: unknown): number { if (!quota) return CODEX_UNKNOWN_USAGE_SCORE; - if (isThirtyDayOnlyCodexPlan(plan)) { - return typeof quota.monthlyPercent === "number" && Number.isFinite(quota.monthlyPercent) - ? quota.monthlyPercent - : CODEX_UNKNOWN_USAGE_SCORE; - } - const values = [quota.weeklyPercent, quota.monthlyPercent] + const values = (isThirtyDayOnlyCodexPlan(plan) + ? [quota.monthlyPercent, quota.shortPercent] + : [quota.weeklyPercent, quota.monthlyPercent, quota.shortPercent]) .filter((value): value is number => typeof value === "number" && Number.isFinite(value)); return values.length > 0 ? Math.max(...values) : CODEX_UNKNOWN_USAGE_SCORE; } diff --git a/tests/codex-auth-api.test.ts b/tests/codex-auth-api.test.ts index f096c33c8d..d14cedf6be 100644 --- a/tests/codex-auth-api.test.ts +++ b/tests/codex-auth-api.test.ts @@ -1183,6 +1183,56 @@ describe("codex-auth API", () => { expect(getAccountQuota("preserve-valid")).toEqual(before); }); + test("quota cache rebuilds preserve the short-window tuple", () => { + setAccountQuotaFromParsed("short-cache", { + weeklyPercent: 1, + weeklyResetAt: 2_000_586_800, + monthlyPercent: 3, + monthlyResetAt: 2_002_592_000, + shortPercent: 0, + shortResetAt: 2_000_000_000, + shortWindowSeconds: 18_000, + }); + expect(getAccountQuota("short-cache")).toMatchObject({ + weeklyPercent: 1, + shortPercent: 0, + shortResetAt: 2_000_000_000, + shortWindowSeconds: 18_000, + }); + + setAccountQuotaFromParsed("short-cache", { + shortPercent: 4, + shortResetAt: 2_000_000_100, + shortWindowSeconds: 18_000, + }); + expect(getAccountQuota("short-cache")).toMatchObject({ + weeklyPercent: 1, + monthlyPercent: 3, + shortPercent: 4, + shortResetAt: 2_000_000_100, + shortWindowSeconds: 18_000, + }); + + updateAccountQuota("short-cache", 2, 2_000_586_900); + expect(getAccountQuota("short-cache")).toMatchObject({ + weeklyPercent: 2, + monthlyPercent: 3, + shortPercent: 4, + shortResetAt: 2_000_000_100, + shortWindowSeconds: 18_000, + }); + + setAccountQuotaFromParsed("short-cache", { resetCredits: 3 }); + expect(getAccountQuota("short-cache")).toMatchObject({ + weeklyPercent: 2, + monthlyPercent: 3, + shortPercent: 4, + shortResetAt: 2_000_000_100, + shortWindowSeconds: 18_000, + resetCredits: 3, + }); + }); + test("GET /api/codex-auth/quota returns stored quotas", async () => { updateAccountQuota("q-test", 30); const req = new Request("http://localhost/api/codex-auth/quota", { method: "GET" }); @@ -1230,6 +1280,47 @@ describe("codex-auth API", () => { } }); + test("GET /api/codex-auth/accounts preserves a parsed K12 short window through cache and DTO", async () => { + const config = makeConfig(); + seedPoolAccount(config, { + id: "pool-k12-short", + email: "pool-k12-short@example.com", + plan: "k12", + accessToken: "tok", + refreshToken: "ref", + chatgptAccountId: "acc-pool-k12-short", + }); + + const originalFetch = globalThis.fetch; + globalThis.fetch = (async () => Response.json({ + plan_type: "k12", + rate_limit: { + primary_window: { used_percent: 0, reset_at: 2_000_000_000, limit_window_seconds: 18_000 }, + secondary_window: { used_percent: 1, reset_at: 2_000_586_800, limit_window_seconds: 604_800 }, + }, + })) as typeof fetch; + + try { + const req = new Request("http://localhost/api/codex-auth/accounts?refresh=1", { method: "GET" }); + const resp = await handleCodexAuthAPI(req, new URL(req.url), config); + expect(resp!.status).toBe(200); + const data = await resp!.json() as { + accounts: Array<{ id: string; quota?: Record }>; + }; + const quota = data.accounts.find(account => account.id === "pool-k12-short")?.quota; + expect(quota).toMatchObject({ + weeklyPercent: 1, + weeklyResetAt: 2_000_586_800, + shortPercent: 0, + shortResetAt: 2_000_000_000, + shortWindowSeconds: 18_000, + }); + expect(getAccountQuota("pool-k12-short")).toMatchObject(quota!); + } finally { + globalThis.fetch = originalFetch; + } + }); + test("GET /api/codex-auth/accounts refresh=1 bypasses cached pool quota", async () => { const config = makeConfig(); seedPoolAccount(config, { diff --git a/tests/codex-routing.test.ts b/tests/codex-routing.test.ts index 18d7a21ca0..c16290d344 100644 --- a/tests/codex-routing.test.ts +++ b/tests/codex-routing.test.ts @@ -119,6 +119,8 @@ describe("codex routing", () => { test("usage score uses the hottest known quota window", () => { expect(computeCodexUsageScore({ weeklyPercent: 81 })).toBe(81); expect(computeCodexUsageScore({ weeklyPercent: 15, monthlyPercent: 91 })).toBe(91); + expect(computeCodexUsageScore({ weeklyPercent: 15, monthlyPercent: 20, shortPercent: 92 })).toBe(92); + expect(computeCodexUsageScore({ shortPercent: 0 })).toBe(0); expect(computeCodexUsageScore({ weeklyPercent: 15 })).toBe(15); }); @@ -179,6 +181,7 @@ describe("codex routing", () => { test("go and free plans use only the 30d quota window", () => { expect(computeCodexUsageScore({ weeklyPercent: 99, monthlyPercent: 12 }, "go")).toBe(12); expect(computeCodexUsageScore({ weeklyPercent: 99, monthlyPercent: 13 }, "free")).toBe(13); + expect(computeCodexUsageScore({ weeklyPercent: 99, monthlyPercent: 12, shortPercent: 14 }, "go")).toBe(14); expect(computeCodexUsageScore({ weeklyPercent: 1 }, "go")).toBe(CODEX_UNKNOWN_USAGE_SCORE); }); @@ -1274,6 +1277,19 @@ describe("codex routing", () => { }); }); + test("a zero-valued short-only WHAM snapshot remains known quota (#2047)", () => { + expect(parseUsageQuota({ + plan_type: "k12", + rate_limit: { + primary_window: { used_percent: 0, reset_at: 2000000000, limit_window_seconds: 18000 }, + }, + })).toMatchObject({ + shortPercent: 0, + shortResetAt: 2000000000, + shortWindowSeconds: 18000, + }); + }); + test("an exhausted burst window takes the account out of rotation (#1791)", () => { // Upstream enforces the 5-hour window independently, so an account at 100% there is // genuinely blocked even while its weekly quota is untouched. Reporting it as usable