fix(codex): retain short-window quota end to end

src/codex/quota.ts

@@ -186,7 +186,7 @@ function normalizeResetAt(value: unknown): number | undefined {
 }
 
 function hasKnownQuotaValue(quota: Omit<StoredAccountQuota, "updatedAt">): 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<StoredAccountQuota, "updatedAt">): boole
   return quota.monthlyPercent !== undefined || quota.monthlyResetAt !== undefined;
 }
 
+function snapshotHasShort(quota: Omit<StoredAccountQuota, "updatedAt">): boolean {
+  return quota.shortPercent !== undefined
+    || quota.shortResetAt !== undefined
+    || quota.shortWindowSeconds !== undefined;
+}
+
 function snapshotHasUsage(quota: Omit<StoredAccountQuota, "updatedAt">): 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(),
   };

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;
 }

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<string, unknown> }>;
+      };
+      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, {

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