From ffad342b91e1c5657bf022ce906c5b3ed31bf888 Mon Sep 17 00:00:00 2001 From: Yash Singh Date: Sun, 6 Sep 2026 21:40:10 -0500 Subject: [PATCH] fix(codex): keep Spark limits from replacing the main allowance - Select the main Codex rate-limit bucket explicitly - Ignore model-specific Spark notifications --- .../src/provider/Layers/CodexProvider.ts | 6 ++ .../provider/Layers/codexUsageLimits.test.ts | 72 +++++++++++++++++++ .../src/provider/Layers/codexUsageLimits.ts | 12 +++- 3 files changed, 89 insertions(+), 1 deletion(-) diff --git a/apps/server/src/provider/Layers/CodexProvider.ts b/apps/server/src/provider/Layers/CodexProvider.ts index d65a09c6d4f9..48f67c993e15 100644 --- a/apps/server/src/provider/Layers/CodexProvider.ts +++ b/apps/server/src/provider/Layers/CodexProvider.ts @@ -49,6 +49,10 @@ const RATE_LIMITS_PROBE_TIMEOUT_MS = 3_000; type CodexRateLimitsProbe = | { readonly snapshot: CodexRateLimitSnapshot; + readonly rateLimitsByLimitId?: + | Readonly> + | null + | undefined; readonly resetCredits: CodexResetCreditsSummary | null | undefined; } | { readonly failure: string }; @@ -434,6 +438,7 @@ const probeCodexAppServerProvider = Effect.fn("probeCodexAppServerProvider")(fun client.request("account/rateLimits/read", undefined).pipe( Effect.map((response): CodexRateLimitsProbe => ({ snapshot: response.rateLimits, + rateLimitsByLimitId: response.rateLimitsByLimitId, resetCredits: response.rateLimitResetCredits, })), Effect.timeoutOption(Duration.millis(RATE_LIMITS_PROBE_TIMEOUT_MS)), @@ -652,6 +657,7 @@ export const checkCodexProviderStatus = Effect.fn("checkCodexProviderStatus")(fu }) : codexRateLimitsToLimits({ snapshot: snapshot.rateLimits.snapshot, + rateLimitsByLimitId: snapshot.rateLimits.rateLimitsByLimitId, resetCredits: snapshot.rateLimits.resetCredits, checkedAt, }); diff --git a/apps/server/src/provider/Layers/codexUsageLimits.test.ts b/apps/server/src/provider/Layers/codexUsageLimits.test.ts index a86d5a9ab6b6..9de27ef75c25 100644 --- a/apps/server/src/provider/Layers/codexUsageLimits.test.ts +++ b/apps/server/src/provider/Layers/codexUsageLimits.test.ts @@ -60,6 +60,54 @@ describe("codexRateLimitsToLimits", () => { }, ]); }); + + it("selects the main Codex allowance and leaves Spark out", () => { + const spark = { + limitId: "codex_bengalfox", + primary: { usedPercent: 0, windowDurationMins: 300 }, + secondary: { usedPercent: 90, windowDurationMins: 10080 }, + }; + expect( + codexRateLimitsToLimits({ + checkedAt, + snapshot: spark, + rateLimitsByLimitId: { + codex_bengalfox: spark, + codex: { secondary: { usedPercent: 42, windowDurationMins: 10080 } }, + }, + }).windows, + ).toEqual([ + { + id: "secondary", + kind: "weekly", + label: "Weekly", + usedPercent: 42, + windowDurationMins: 10080, + }, + ]); + }); + + it.each([undefined, null, {}])( + "supports legacy reads with no bucket map: %j", + (rateLimitsByLimitId) => { + const snapshot = { primary: { usedPercent: 12, windowDurationMins: 300 } }; + expect(codexRateLimitsToLimits({ checkedAt, snapshot, rateLimitsByLimitId })).toEqual( + codexRateLimitsToLimits({ checkedAt, snapshot }), + ); + }, + ); + + it("does not show a model-specific legacy snapshot as the main allowance", () => { + expect( + codexRateLimitsToLimits({ + checkedAt, + snapshot: { + limitId: "codex_bengalfox", + secondary: { usedPercent: 90 }, + }, + }).windows, + ).toEqual([]); + }); }); describe("codexRateLimitsToUpdate", () => { @@ -81,6 +129,30 @@ describe("codexRateLimitsToUpdate", () => { }); expect(codexRateLimitsToUpdate({ planType: "plus" })).toBeUndefined(); }); + + it("ignores Spark notifications so they cannot overwrite the main allowance", () => { + expect( + codexRateLimitsToUpdate({ + limitId: "codex_bengalfox", + primary: { usedPercent: 0, windowDurationMins: 300 }, + secondary: { usedPercent: 90, windowDurationMins: 10080 }, + }), + ).toBeUndefined(); + expect( + codexRateLimitsToUpdate({ + limitId: "codex", + secondary: { usedPercent: 42, windowDurationMins: 10080 }, + })?.windows, + ).toEqual([ + { + id: "secondary", + kind: "weekly", + label: "Weekly", + usedPercent: 42, + windowDurationMins: 10080, + }, + ]); + }); }); describe("codexRateLimitsFailureMessage", () => { diff --git a/apps/server/src/provider/Layers/codexUsageLimits.ts b/apps/server/src/provider/Layers/codexUsageLimits.ts index a9e03681fafc..32862cdb337d 100644 --- a/apps/server/src/provider/Layers/codexUsageLimits.ts +++ b/apps/server/src/provider/Layers/codexUsageLimits.ts @@ -26,6 +26,7 @@ interface CodexRateLimitWindow { /** Structural view of the generated `RateLimitSnapshot`; both messages satisfy it. */ export interface CodexRateLimitSnapshot { + readonly limitId?: string | null; readonly planType?: string | null; readonly primary?: CodexRateLimitWindow | null; readonly secondary?: CodexRateLimitWindow | null; @@ -68,6 +69,9 @@ function labelForKind(kind: ServerProviderUsageWindow["kind"]): string { export function codexRateLimitsToWindows( snapshot: CodexRateLimitSnapshot, ): ReadonlyArray { + // Show the main allowance only. Model-specific notifications (such as Spark) + // must not replace its primary/secondary rows. Older CLIs omit the limit id. + if (snapshot.limitId && snapshot.limitId !== "codex") return []; const isMonthlyPlan = snapshot.planType === "free" || snapshot.planType === "go"; const positions = [ ["primary", snapshot.primary, isMonthlyPlan ? MONTH_MINS : SESSION_MINS], @@ -110,14 +114,20 @@ export function codexResetCreditsToContract( export function codexRateLimitsToLimits(input: { readonly snapshot: CodexRateLimitSnapshot; + readonly rateLimitsByLimitId?: + | Readonly> + | null + | undefined; readonly resetCredits?: CodexResetCreditsSummary | null | undefined; readonly checkedAt: string; }): ServerProviderUsageLimits { const resetCredits = codexResetCreditsToContract(input.resetCredits); + // Select the main bucket explicitly; the legacy snapshot can name another limit. + const windows = codexRateLimitsToWindows(input.rateLimitsByLimitId?.codex ?? input.snapshot); return { ...makeUsageLimits({ checkedAt: input.checkedAt, - windows: codexRateLimitsToWindows(input.snapshot), + windows, }), ...(resetCredits ? { resetCredits } : {}), };