From 9b0c5a02d95220161fc18c72aba5756fdcbf68e3 Mon Sep 17 00:00:00 2001 From: bitkyc08-arch Date: Thu, 20 Aug 2026 03:39:19 +0900 Subject: [PATCH] fix(codex): retain the K12 short-window quota end to end parseUsageQuota filled shortPercent and setAccountQuotaFromParsed dropped it, so the 5-hour burst window never reached the cache, the accounts DTO, the dashboard, or routing. A saturated short window was invisible to account selection. Carries @Ingwannu's #2056: shortPercent joins hasKnownQuotaValue, a new snapshotHasShort keeps a short-only snapshot from reading as empty, partial weekly/monthly snapshots no longer clobber a known short window, and updateAccountQuota carries the tuple. Also fixes the blocker raised in review on both #2056 and #2062: the scorer took Math.max over every finite window, so a snapshot carrying only shortPercent: 0 scored a flat 0 and made an account whose long windows were never observed look like the emptiest in the pool - pickLowestUsageAmong would then send every request to it. The burst window now refines a known long-window position instead of standing in for one, and returns CODEX_UNKNOWN_USAGE_SCORE until a governing window is actually observed. The ported test asserted the old behavior directly (computeCodexUsageScore({ shortPercent: 0 }) === 0); it is replaced by a case that pins the corrected contract in both directions. Closes #2047 --- src/codex/quota.ts | 31 ++++++++++-- src/codex/routing.ts | 22 +++++---- tests/codex-auth-api.test.ts | 91 ++++++++++++++++++++++++++++++++++++ tests/codex-routing.test.ts | 27 +++++++++++ 4 files changed, 160 insertions(+), 11 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..fa6fa63d83 100644 --- a/src/codex/routing.ts +++ b/src/codex/routing.ts @@ -322,16 +322,22 @@ 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] - .filter((value): value is number => typeof value === "number" && Number.isFinite(value)); - return values.length > 0 ? Math.max(...values) : CODEX_UNKNOWN_USAGE_SCORE; + const finite = (value: unknown): value is number => typeof value === "number" && Number.isFinite(value); + const longWindows = isThirtyDayOnlyCodexPlan(plan) + ? [quota.monthlyPercent] + : [quota.weeklyPercent, quota.monthlyPercent]; + const knownLong = longWindows.filter(finite); + // The short burst window only REFINES a known long-window position; it cannot stand in for + // one. A snapshot carrying just `shortPercent: 0` would otherwise score a flat 0 and make an + // account whose weekly/monthly usage is entirely unverified look like the emptiest in the + // pool, so `pickLowestUsageAmong` would send every request to it. Unknown has to stay + // unknown until a governing window is actually observed. + if (knownLong.length === 0) return CODEX_UNKNOWN_USAGE_SCORE; + const values = finite(quota.shortPercent) ? [...knownLong, quota.shortPercent] : knownLong; + return Math.max(...values); } export function classifyCodexUpstreamOutcome( 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..d84725722b 100644 --- a/tests/codex-routing.test.ts +++ b/tests/codex-routing.test.ts @@ -119,9 +119,22 @@ 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({ weeklyPercent: 15 })).toBe(15); }); + test("a short-only snapshot is unknown usage, not zero usage", () => { + // The burst window refines a known long-window position; it cannot stand in for one. + // Scoring a bare `shortPercent: 0` as 0 would make an account whose weekly/monthly usage + // was never observed look like the emptiest in the pool, and pickLowestUsageAmong would + // send every request to it. + expect(computeCodexUsageScore({ shortPercent: 0 })).toBe(CODEX_UNKNOWN_USAGE_SCORE); + expect(computeCodexUsageScore({ shortPercent: 87 })).toBe(CODEX_UNKNOWN_USAGE_SCORE); + // Once a governing window is known, the burst still wins when it is hotter. + expect(computeCodexUsageScore({ weeklyPercent: 1, shortPercent: 100 })).toBe(100); + expect(computeCodexUsageScore({ weeklyPercent: 40, shortPercent: 0 })).toBe(40); + }); + test("exact-account failures record health without rotating the active Pool account", () => { const transient = makeConfig({ upstreamFailoverThreshold: 1, activeCodexAccountId: "a" }); const transientThread = "fixed-transient-thread"; @@ -179,6 +192,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 +1288,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