diff --git a/src/domain/quotaCache.ts b/src/domain/quotaCache.ts index a0b0e2caad..e89f0849f5 100644 --- a/src/domain/quotaCache.ts +++ b/src/domain/quotaCache.ts @@ -101,6 +101,35 @@ function clampPercent(value: number): number { return Math.max(0, Math.min(100, value)); } +function normalizeWindowKey(value: string): string { + return value + .toLowerCase() + .replace(/[^a-z0-9]+/g, " ") + .trim(); +} + +function resolveQuotaWindow( + quotas: Record, + windowName: string +): QuotaInfo | null { + const direct = quotas[windowName]; + if (direct) return direct; + + const normalizedTarget = normalizeWindowKey(windowName); + if (!normalizedTarget) return null; + + for (const [key, quota] of Object.entries(quotas)) { + const normalizedKey = normalizeWindowKey(key); + if (!normalizedKey) continue; + if (normalizedKey === normalizedTarget) return quota; + // Support canonical selection of generic windows from labeled windows, + // e.g. "weekly" from "weekly (7d)" or "session" from "session (5h)". + if (normalizedKey.startsWith(`${normalizedTarget} `)) return quota; + } + + return null; +} + function earliestResetAt(quotas: Record): string | null { let earliest: string | null = null; let earliestMs = Infinity; @@ -201,7 +230,7 @@ export function getQuotaWindowStatus( const now = Date.now(); - const window = entry.quotas[windowName]; + const window = resolveQuotaWindow(entry.quotas, windowName); if (!window) return null; const remainingPercentage = clampPercent(window.remainingPercentage); diff --git a/src/sse/services/auth.ts b/src/sse/services/auth.ts index 37de3bcc00..e77916c413 100644 --- a/src/sse/services/auth.ts +++ b/src/sse/services/auth.ts @@ -132,12 +132,37 @@ function normalizeWindowName(windowName: unknown): string | null { return normalized.length > 0 ? normalized : null; } -function getLegacyCodexWindows(providerSpecificData: JsonRecord): string[] { +function uniqueWindows(windows: string[]): string[] { + return [...new Set(windows)]; +} + +function normalizeCodexWindowName(windowName: string): string { + const normalized = windowName.trim().toLowerCase(); + if (normalized === "session (5h)" || normalized === "5h" || normalized === "five_hour") { + return "session"; + } + if (normalized === "weekly (7d)" || normalized === "7d" || normalized === "seven_day") { + return "weekly"; + } + return normalized; +} + +function applyCodexWindowPolicy(rawWindows: string[], providerSpecificData: JsonRecord): string[] { const codexPolicy = getCodexLimitPolicy(providerSpecificData); - const windows: string[] = []; + const normalizedRaw = rawWindows.map(normalizeCodexWindowName).filter(Boolean); + + // Preserve explicitly configured custom windows, but enforce canonical Codex windows + // from toggles so weekly exhaustion is never skipped when useWeekly=true. + let windows = [...normalizedRaw]; + windows = windows.filter((windowName) => { + if (windowName === "session") return codexPolicy.use5h; + if (windowName === "weekly") return codexPolicy.useWeekly; + return true; + }); if (codexPolicy.use5h) windows.push("session"); if (codexPolicy.useWeekly) windows.push("weekly"); - return windows; + + return uniqueWindows(windows); } export function resolveQuotaLimitPolicy( @@ -149,8 +174,7 @@ export function resolveQuotaLimitPolicy( const windows = rawWindows.map(normalizeWindowName).filter(Boolean) as string[]; if (provider === "codex") { - const fallbackWindows = getLegacyCodexWindows(providerSpecificData); - const defaultWindows = windows.length > 0 ? windows : fallbackWindows; + const defaultWindows = applyCodexWindowPolicy(windows, providerSpecificData); const enabled = toBooleanOrDefault(rawPolicy.enabled, defaultWindows.length > 0); return { diff --git a/tests/unit/quota-policy-generalization.test.mjs b/tests/unit/quota-policy-generalization.test.mjs index 881480fef3..a99c27a172 100644 --- a/tests/unit/quota-policy-generalization.test.mjs +++ b/tests/unit/quota-policy-generalization.test.mjs @@ -21,6 +21,26 @@ test("resolveQuotaLimitPolicy keeps codex legacy defaults when generic policy is assert.equal(policy.thresholdPercent, 90); }); +test("resolveQuotaLimitPolicy enforces codex weekly window when weekly toggle is enabled", () => { + const policy = auth.resolveQuotaLimitPolicy("codex", { + codexLimitPolicy: { use5h: true, useWeekly: true }, + limitPolicy: { enabled: true, windows: ["session"] }, + }); + + assert.equal(policy.enabled, true); + assert.deepEqual(policy.windows.sort(), ["session", "weekly"]); +}); + +test("resolveQuotaLimitPolicy removes codex weekly window when weekly toggle is disabled", () => { + const policy = auth.resolveQuotaLimitPolicy("codex", { + codexLimitPolicy: { use5h: true, useWeekly: false }, + limitPolicy: { enabled: true, windows: ["session", "weekly"] }, + }); + + assert.equal(policy.enabled, true); + assert.deepEqual(policy.windows, ["session"]); +}); + test("resolveQuotaLimitPolicy disables non-codex policy by default", () => { const policy = auth.resolveQuotaLimitPolicy("openai", {}); assert.equal(policy.enabled, false); @@ -60,6 +80,26 @@ test("evaluateQuotaLimitPolicy blocks when configured window reaches threshold", assert.equal(result.resetAt, resetAt); }); +test("evaluateQuotaLimitPolicy matches canonical weekly window against labeled cache keys", () => { + const resetAt = new Date(Date.now() + 60_000).toISOString(); + quotaCache.setQuotaCache("conn-policy-weekly-label", "codex", { + "weekly (7d)": { remainingPercentage: 0, resetAt }, + }); + + const result = auth.evaluateQuotaLimitPolicy( + "codex", + buildConnection("conn-policy-weekly-label", { + codexLimitPolicy: { use5h: true, useWeekly: true }, + limitPolicy: { enabled: true, windows: ["weekly"] }, + }) + ); + + assert.equal(result.blocked, true); + assert.equal(result.reasons.length, 1); + assert.match(result.reasons[0], /weekly usage/i); + assert.equal(result.resetAt, resetAt); +}); + test("evaluateQuotaLimitPolicy does not block when no quota data exists", () => { const result = auth.evaluateQuotaLimitPolicy( "openai",