From aa0e312d8a6ec70b47445a61cb58aeb448b64156 Mon Sep 17 00:00:00 2001 From: payne Date: Fri, 15 May 2026 02:19:55 +0300 Subject: [PATCH] feat(limits): per-window quota cutoffs across all providers with usage data (#2267) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit feat(limits): per-window quota cutoffs across all providers with usage data (#2267 — thanks @payne0420) --- open-sse/services/codexQuotaFetcher.ts | 36 ++- open-sse/services/genericQuotaFetcher.ts | 214 +++++++++++++++++ open-sse/services/quotaPreflight.ts | 183 +++++++++++++-- open-sse/services/usage.ts | 33 +++ .../ProviderLimits/QuotaCutoffModal.tsx | 201 ++++++++++++++++ .../usage/components/ProviderLimits/index.tsx | 164 ++++++++++++- src/app/api/providers/[id]/route.ts | 25 ++ src/app/api/providers/quota-windows/route.ts | 33 +++ src/i18n/messages/en.json | 11 + src/lib/db/core.ts | 10 + ...der_connection_quota_window_thresholds.sql | 14 ++ src/lib/db/providers.ts | 94 +++++++- src/lib/resilience/settings.ts | 104 +++++++++ src/shared/validation/schemas.ts | 17 ++ src/sse/handlers/chat.ts | 7 + src/sse/services/auth.ts | 96 +++++++- tests/unit/generic-quota-fetcher.test.ts | 90 ++++++++ ...ovider-connections-quota-threshold.test.ts | 164 +++++++++++++ tests/unit/quota-preflight.test.ts | 215 +++++++++++++++--- ...esilience-settings-quota-preflight.test.ts | 128 +++++++++++ tests/unit/sse-auth.test.ts | 108 +++++++++ 21 files changed, 1887 insertions(+), 60 deletions(-) create mode 100644 open-sse/services/genericQuotaFetcher.ts create mode 100644 src/app/(dashboard)/dashboard/usage/components/ProviderLimits/QuotaCutoffModal.tsx create mode 100644 src/app/api/providers/quota-windows/route.ts create mode 100644 src/lib/db/migrations/056_provider_connection_quota_window_thresholds.sql create mode 100644 tests/unit/generic-quota-fetcher.test.ts create mode 100644 tests/unit/provider-connections-quota-threshold.test.ts create mode 100644 tests/unit/resilience-settings-quota-preflight.test.ts diff --git a/open-sse/services/codexQuotaFetcher.ts b/open-sse/services/codexQuotaFetcher.ts index ec31cc9215..fe2d421d9b 100644 --- a/open-sse/services/codexQuotaFetcher.ts +++ b/open-sse/services/codexQuotaFetcher.ts @@ -16,9 +16,20 @@ * Registration: call registerCodexQuotaFetcher() once at server startup. */ -import { registerQuotaFetcher, type QuotaInfo } from "./quotaPreflight.ts"; +import { registerQuotaFetcher, registerQuotaWindows, type QuotaInfo } from "./quotaPreflight.ts"; import { registerMonitorFetcher } from "./quotaMonitor.ts"; +/** + * Stable identifiers for Codex's quota windows. These match the quota keys + * surfaced by `getCodexUsage` (in usage.ts) and rendered by the dashboard, + * so per-window thresholds set in the UI line up with the keys persisted + * in `provider_connections.quota_window_thresholds_json`. The dedicated + * Codex fetcher exposes only session + weekly today; the plan-dependent + * code_review window is surfaced by the generic path when present. + */ +export const CODEX_WINDOW_SESSION = "session"; // primary 5-hour window +export const CODEX_WINDOW_WEEKLY = "weekly"; // secondary 7-day window + // Codex usage endpoint (same as usage.ts CODEX_CONFIG) const CODEX_USAGE_URL = "https://chatgpt.com/backend-api/wham/usage"; @@ -249,16 +260,26 @@ function parseCodexUsageResponse(data: unknown): CodexDualWindowQuota | null { const limitReached = Boolean(rateLimit["limit_reached"] ?? rateLimit["limitReached"]); + const window5h = { percentUsed: usedPercent5h / 100, resetAt: resetAt5h }; + const window7d = { percentUsed: usedPercent7d / 100, resetAt: resetAt7d }; + return { used: worstPercentUsed, total: 100, percentUsed: percentUsedNormalized, - resetAt: getDominantResetAt({ - window5h: { percentUsed: usedPercent5h / 100, resetAt: resetAt5h }, - window7d: { percentUsed: usedPercent7d / 100, resetAt: resetAt7d }, - }), - window5h: { percentUsed: usedPercent5h / 100, resetAt: resetAt5h }, - window7d: { percentUsed: usedPercent7d / 100, resetAt: resetAt7d }, + resetAt: getDominantResetAt({ window5h, window7d }), + // Per-window breakdown for the preflight evaluator. Keys match what the + // dashboard renders (session = 5h, weekly = 7d) so user-set cutoffs and + // displayed quotas refer to the same windows. + windows: { + ...(hasPrimary ? { [CODEX_WINDOW_SESSION]: window5h } : {}), + ...(hasSecondary ? { [CODEX_WINDOW_WEEKLY]: window7d } : {}), + }, + // Legacy fields preserved for existing consumers (quotaMonitor, cooldown + // computation in accountFallback). These mirror the new windows entries + // but keep the historical names — do not remove without checking callers. + window5h, + window7d, limitReached, }; } @@ -314,4 +335,5 @@ export function invalidateCodexQuotaCache(connectionId: string): void { export function registerCodexQuotaFetcher(): void { registerQuotaFetcher("codex", fetchCodexQuota); registerMonitorFetcher("codex", fetchCodexQuota); + registerQuotaWindows("codex", [CODEX_WINDOW_SESSION, CODEX_WINDOW_WEEKLY]); } diff --git a/open-sse/services/genericQuotaFetcher.ts b/open-sse/services/genericQuotaFetcher.ts new file mode 100644 index 0000000000..139bd13ac7 --- /dev/null +++ b/open-sse/services/genericQuotaFetcher.ts @@ -0,0 +1,214 @@ +/** + * genericQuotaFetcher.ts — Generic preflight quota fetcher + * + * Wraps the existing per-provider usage fetchers in `usage.ts` so that any + * provider with a `getUsageForProvider` implementation gets per-window + * preflight enforcement automatically. This is the bridge between the + * dashboard's "Provider Limits" data (which already supports ~16 providers) + * and the quotaPreflight system (which previously only had Codex). + * + * For providers that ship their own custom QuotaFetcher (Codex, CROF, + * DeepSeek, Bailian Coding Plan, etc.) the registrar skips them — their + * bespoke fetchers stay in charge. + * + * Each provider's first successful response also populates the static + * `registerQuotaWindows` registry so other callers (UI window catalog, + * tests) can discover which windows that provider exposes. + */ + +import { getUsageForProvider, USAGE_FETCHER_PROVIDERS } from "./usage.ts"; +import { + getQuotaFetcher, + registerQuotaFetcher, + registerQuotaWindows, + type QuotaFetcher, + type QuotaInfo, +} from "./quotaPreflight.ts"; + +// 60s — matches Codex's TTL. Long enough to avoid hammering upstream usage +// endpoints on every routing decision, short enough that a near-exhausted +// account is skipped within one minute of crossing its threshold. +const CACHE_TTL_MS = 60_000; + +interface CacheEntry { + quota: QuotaInfo; + fetchedAt: number; +} + +const cache = new Map(); + +function cacheKey(provider: string, connectionId: string): string { + return `${provider}::${connectionId}`; +} + +// Auto-cleanup stale entries — same shape as codexQuotaFetcher. +const _cacheCleanup = setInterval(() => { + const now = Date.now(); + for (const [key, entry] of cache) { + if (now - entry.fetchedAt > CACHE_TTL_MS * 5) cache.delete(key); + } +}, 5 * 60_000); +if (typeof _cacheCleanup === "object" && "unref" in _cacheCleanup) { + (_cacheCleanup as { unref?: () => void }).unref?.(); +} + +function toNumber(value: unknown): number | null { + if (typeof value === "number" && Number.isFinite(value)) return value; + if (typeof value === "string") { + const parsed = parseFloat(value); + if (Number.isFinite(parsed)) return parsed; + } + return null; +} + +/** + * Compute percentUsed (0-1) for a single quota entry. Prefers the explicit + * remainingPercentage / used / total fields surfaced by per-provider + * fetchers (see `usage.ts`). Returns null when the entry is unlimited or + * doesn't expose enough data to compute a percent — preflight ignores + * those windows. + */ +function percentUsedForQuota(entry: unknown): number | null { + if (!entry || typeof entry !== "object") return null; + const q = entry as Record; + if (q.unlimited === true) return null; + + const remainingPercentage = toNumber(q.remainingPercentage); + if (remainingPercentage !== null) { + // remainingPercentage is 0-100 in the usage.ts contract. + const used = (100 - Math.max(0, Math.min(100, remainingPercentage))) / 100; + return used; + } + + const used = toNumber(q.used); + const total = toNumber(q.total); + if (used !== null && total !== null && total > 0) { + return Math.max(0, Math.min(1, used / total)); + } + + return null; +} + +function resetAtForQuota(entry: unknown): string | null { + if (!entry || typeof entry !== "object") return null; + const q = entry as Record; + return typeof q.resetAt === "string" ? q.resetAt : null; +} + +interface ConnectionInputs { + id?: string; + provider?: string; + accessToken?: string; + apiKey?: string; + providerSpecificData?: Record; + projectId?: string; + email?: string; +} + +/** + * Reshape a raw `getUsageForProvider` response into the preflight `QuotaInfo` + * contract. Returns `null` if there are no measurable windows (all unlimited + * / shape-unknown / missing). Exported for unit testing — the production path + * is `fetchGenericQuota`, which adds caching + the upstream call. + */ +export function convertUsageToQuotaInfo(usage: unknown): QuotaInfo | null { + if (!usage || typeof usage !== "object") return null; + const usageRecord = usage as Record; + if ( + typeof usageRecord.message === "string" && + (!usageRecord.quotas || typeof usageRecord.quotas !== "object") + ) { + // Provider explicitly told us it couldn't fetch (auth expired, etc.). + // Fail open — let the request proceed and surface the failure through + // its normal error path. + return null; + } + + const quotasObj = usageRecord.quotas; + if (!quotasObj || typeof quotasObj !== "object" || Array.isArray(quotasObj)) { + return null; + } + + const windows: Record = {}; + let worstPercent = 0; + let worstResetAt: string | null = null; + for (const [name, entry] of Object.entries(quotasObj as Record)) { + const percentUsed = percentUsedForQuota(entry); + if (percentUsed === null) continue; + const resetAt = resetAtForQuota(entry); + windows[name] = { percentUsed, resetAt }; + if (percentUsed > worstPercent) { + worstPercent = percentUsed; + worstResetAt = resetAt; + } + } + + if (Object.keys(windows).length === 0) return null; + + return { + used: 0, + total: 0, + percentUsed: worstPercent, + resetAt: worstResetAt, + windows, + }; +} + +/** + * Fetch quota for a connection by delegating to the appropriate + * provider-specific usage fetcher and reshaping its output into the + * preflight `QuotaInfo` contract (with a `windows` map for per-window + * threshold evaluation). + */ +export const fetchGenericQuota: QuotaFetcher = async (connectionId, connection) => { + if (!connection) return null; + const conn = connection as ConnectionInputs; + const provider = typeof conn.provider === "string" ? conn.provider : null; + if (!provider) return null; + + const key = cacheKey(provider, connectionId); + const cached = cache.get(key); + if (cached && Date.now() - cached.fetchedAt < CACHE_TTL_MS) { + return cached.quota; + } + + let usage: unknown; + try { + usage = await getUsageForProvider(conn as Parameters[0]); + } catch { + return null; + } + + const quota = convertUsageToQuotaInfo(usage); + if (!quota) return null; + + // Refresh the static window catalog so the dashboard can render the right + // modal inputs without waiting for the user to open the page. + registerQuotaWindows(provider, Object.keys(quota.windows || {})); + + cache.set(key, { quota, fetchedAt: Date.now() }); + return quota; +}; + +/** + * Force-invalidate the cache for a connection — call after the connection + * receives an upstream 429 / quota-reset event so the next preflight gets + * fresh data instead of a 60s stale window. + */ +export function invalidateGenericQuotaCache(provider: string, connectionId: string): void { + cache.delete(cacheKey(provider, connectionId)); +} + +/** + * Register the generic fetcher for every provider that has a usage + * implementation. Providers with bespoke fetchers (Codex, CROF, DeepSeek, + * Bailian Coding Plan) MUST be registered before this runs so the defensive + * `getQuotaFetcher` check below preserves them — see `src/sse/handlers/chat.ts` + * for the registration order. Idempotent: re-running this is a no-op. + */ +export function registerGenericQuotaFetchers(): void { + for (const provider of USAGE_FETCHER_PROVIDERS) { + if (getQuotaFetcher(provider)) continue; // bespoke fetcher already registered — leave it alone + registerQuotaFetcher(provider, fetchGenericQuota); + } +} diff --git a/open-sse/services/quotaPreflight.ts b/open-sse/services/quotaPreflight.ts index 3e5740cdfe..ed5cdb8d8e 100644 --- a/open-sse/services/quotaPreflight.ts +++ b/open-sse/services/quotaPreflight.ts @@ -2,9 +2,20 @@ * quotaPreflight.ts — Feature 04 * Quota Preflight & Troca Proativa de Conta * - * Toggle: providerSpecificData.quotaPreflightEnabled (default: false) - * Providers register quota fetchers via registerQuotaFetcher(). - * Graceful degradation when no fetcher registered. + * Providers register quota fetchers via registerQuotaFetcher(). The caller + * (`src/sse/services/auth.ts::getProviderCredentialsWithQuotaPreflight`) is + * responsible for deciding WHEN to invoke preflight — calling it adds the + * latency of an upstream usage fetch, so it should only run when there's + * something to enforce (per-connection overrides, per-(provider, window) + * defaults, or the legacy `quotaPreflightEnabled` flag). + * + * Threshold semantics are "minimum remaining %" — matching the dashboard's + * quota bars, which show remaining (not used). A cutoff of 10 means "stop + * using this connection when it has 10% or less remaining." + * + * `isQuotaPreflightEnabled` remains exported for back-compat so the caller + * can honor the legacy flag, but `preflightQuota` itself no longer gates on + * it — once you invoke preflight, it runs the fetcher and evaluates. */ export interface PreflightQuotaResult { @@ -14,11 +25,25 @@ export interface PreflightQuotaResult { resetAt?: string | null; } +export interface QuotaWindowInfo { + percentUsed: number; + resetAt?: string | null; +} + export interface QuotaInfo { used: number; total: number; + /** Worst-case percentUsed across all known windows (legacy, single-signal). */ percentUsed: number; resetAt?: string | null; + /** + * Optional per-window breakdown. When present, preflight evaluates each + * window against its own threshold (block if ANY window has dropped to or + * below its min-remaining cutoff) instead of using `percentUsed`. Keys are + * window names that match the quota keys surfaced by getUsageForProvider + * (e.g. "session", "weekly", "monthly"). + */ + windows?: Record; } export type QuotaFetcher = ( @@ -26,8 +51,34 @@ export type QuotaFetcher = ( connection?: Record ) => Promise; -const EXHAUSTION_THRESHOLD = 0.98; -const WARN_THRESHOLD = 0.8; +/** + * Registry of named quota windows per provider. Used by the dashboard to + * discover which inputs to render in the cutoffs modal. Providers without + * multiple windows can skip registration — preflight falls back to the + * single-signal `percentUsed` path in that case. + */ +const quotaWindowsRegistry = new Map(); + +export function registerQuotaWindows(provider: string, windows: readonly string[]): void { + quotaWindowsRegistry.set(provider, [...windows]); +} + +export function getQuotaWindows(provider: string): readonly string[] { + return ( + quotaWindowsRegistry.get(provider) || quotaWindowsRegistry.get(provider.toLowerCase()) || [] + ); +} + +export function getAllProviderQuotaWindows(): Record { + return Object.fromEntries(quotaWindowsRegistry); +} + +// Thresholds use "minimum remaining %" semantics so the numbers match the +// dashboard's quota bars (which show remaining %). A cutoff of 2 means +// "block when only 2% remaining" (= 98% used). Warn fires earlier — at +// 20% remaining (= 80% used) by default. +const DEFAULT_MIN_REMAINING_PERCENT = 2; +const DEFAULT_WARN_REMAINING_PERCENT = 20; const quotaFetcherRegistry = new Map(); @@ -44,15 +95,51 @@ export function isQuotaPreflightEnabled(connection: Record): bo return psd?.quotaPreflightEnabled === true; } +export interface PreflightQuotaThresholds { + /** + * Resolve the minimum-remaining cutoff (0-100 integer) for a given window + * name. The connection is blocked when its remaining quota drops to this + * value or below — e.g. returning 10 means "stop when only 10% remaining." + * Resolution order, low-to-high precedence: + * global default → per-(provider, window) default → connection override + * Window name is `null` when the underlying fetcher only exposes a single- + * signal `percentUsed` (legacy path). + */ + resolveMinRemainingPercent?: (window: string | null) => number; + /** + * Resolve the warning threshold (0-100 integer remaining %) for a window. + * Warn fires when remaining quota drops to this value or below — should be + * HIGHER than the min-remaining cutoff so warnings appear before the block + * point. + */ + resolveWarnRemainingPercent?: (window: string | null) => number; +} + +function resolveOrDefault( + resolver: ((window: string | null) => number) | undefined, + window: string | null, + fallbackPercent: number +): number { + if (!resolver) return fallbackPercent; + const raw = resolver(window); + if (typeof raw === "number" && Number.isFinite(raw) && raw >= 0 && raw <= 100) { + return raw; + } + return fallbackPercent; +} + +function remainingPercentFrom(percentUsed: number): number { + return Math.max(0, (1 - percentUsed) * 100); +} + export async function preflightQuota( provider: string, connectionId: string, - connection: Record + connection: Record, + thresholds?: PreflightQuotaThresholds ): Promise { - if (!isQuotaPreflightEnabled(connection)) { - return { proceed: true }; - } - + // No legacy enable-flag gate here — the caller decides when to invoke us + // (see file-level docstring). When there's no fetcher we proceed silently. const fetcher = getQuotaFetcher(provider); if (!fetcher) { return { proceed: true }; @@ -69,11 +156,77 @@ export async function preflightQuota( return { proceed: true }; } - const { percentUsed } = quota; + // Per-window evaluation — only when the fetcher surfaces a windows map. + // We block as soon as ANY single window's remaining quota drops to its + // configured cutoff or below; warnings are logged independently per window. + if (quota.windows && Object.keys(quota.windows).length > 0) { + let worstUsedPercent = 0; + let worstWindow: string | null = null; + let worstResetAt: string | null = null; + for (const [windowName, windowInfo] of Object.entries(quota.windows)) { + const minRemainingPercent = resolveOrDefault( + thresholds?.resolveMinRemainingPercent, + windowName, + DEFAULT_MIN_REMAINING_PERCENT + ); + const warnRemainingPercent = resolveOrDefault( + thresholds?.resolveWarnRemainingPercent, + windowName, + DEFAULT_WARN_REMAINING_PERCENT + ); + const remainingPercent = remainingPercentFrom(windowInfo.percentUsed); - if (percentUsed >= EXHAUSTION_THRESHOLD) { + if (remainingPercent <= minRemainingPercent) { + // Track the most-depleted blocking window so the response can name it. + if (windowInfo.percentUsed > worstUsedPercent) { + worstUsedPercent = windowInfo.percentUsed; + worstWindow = windowName; + worstResetAt = windowInfo.resetAt ?? null; + } else if (worstWindow === null) { + worstWindow = windowName; + worstResetAt = windowInfo.resetAt ?? null; + } + } else if (remainingPercent <= warnRemainingPercent) { + console.warn( + `[QuotaPreflight] ${provider}/${connectionId} ${windowName}: ${remainingPercent.toFixed(1)}% remaining — approaching cutoff` + ); + } + } + + if (worstWindow !== null) { + const worstRemaining = remainingPercentFrom(worstUsedPercent); + console.info( + `[QuotaPreflight] ${provider}/${connectionId} ${worstWindow}: ${worstRemaining.toFixed(1)}% remaining — switching` + ); + return { + proceed: false, + reason: "quota_exhausted", + quotaPercent: worstUsedPercent, + resetAt: worstResetAt, + }; + } + + return { proceed: true, quotaPercent: quota.percentUsed }; + } + + // Legacy single-signal path for fetchers that don't expose per-window data. + const minRemainingPercent = resolveOrDefault( + thresholds?.resolveMinRemainingPercent, + null, + DEFAULT_MIN_REMAINING_PERCENT + ); + const warnRemainingPercent = resolveOrDefault( + thresholds?.resolveWarnRemainingPercent, + null, + DEFAULT_WARN_REMAINING_PERCENT + ); + + const { percentUsed } = quota; + const remainingPercent = remainingPercentFrom(percentUsed); + + if (remainingPercent <= minRemainingPercent) { console.info( - `[QuotaPreflight] ${provider}/${connectionId}: ${(percentUsed * 100).toFixed(1)}% used — switching` + `[QuotaPreflight] ${provider}/${connectionId}: ${remainingPercent.toFixed(1)}% remaining — switching (cutoff ${minRemainingPercent}%)` ); return { proceed: false, @@ -83,9 +236,9 @@ export async function preflightQuota( }; } - if (percentUsed >= WARN_THRESHOLD) { + if (remainingPercent <= warnRemainingPercent) { console.warn( - `[QuotaPreflight] ${provider}/${connectionId}: ${(percentUsed * 100).toFixed(1)}% used — approaching limit` + `[QuotaPreflight] ${provider}/${connectionId}: ${remainingPercent.toFixed(1)}% remaining — approaching cutoff` ); } diff --git a/open-sse/services/usage.ts b/open-sse/services/usage.ts index 21d85889f6..8a7f790223 100644 --- a/open-sse/services/usage.ts +++ b/open-sse/services/usage.ts @@ -1021,6 +1021,39 @@ async function getCursorUsage(accessToken: string, providerSpecificData?: unknow } } +/** + * Single source of truth for which providers have a `getUsageForProvider` + * implementation. Consumers like `genericQuotaFetcher.ts` reference this so + * the registration list can't drift from the switch statement below. + * + * If you add a new provider to the switch, add it here too. + */ +export const USAGE_FETCHER_PROVIDERS = [ + "github", + "gemini-cli", + "antigravity", + "claude", + "codex", + "cursor", + "kiro", + "amazon-q", + "kimi-coding", + "qwen", + "qoder", + "glm", + "glm-cn", + "zai", + "glmt", + "minimax", + "minimax-cn", + "crof", + "bailian-coding-plan", + "nanogpt", + "deepseek", +] as const; + +export type UsageFetcherProvider = (typeof USAGE_FETCHER_PROVIDERS)[number]; + /** * Get usage data for a provider connection * @param {Object} connection - Provider connection with accessToken diff --git a/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/QuotaCutoffModal.tsx b/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/QuotaCutoffModal.tsx new file mode 100644 index 0000000000..075a2b8788 --- /dev/null +++ b/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/QuotaCutoffModal.tsx @@ -0,0 +1,201 @@ +"use client"; + +import { useEffect, useState } from "react"; +import { useTranslations } from "next-intl"; +import Modal from "@/shared/components/Modal"; +import Button from "@/shared/components/Button"; + +export interface QuotaCutoffModalWindow { + /** Stable key — must match the quota name surfaced by the usage fetcher. */ + key: string; + /** Human-readable label rendered next to the input. */ + displayName: string; +} + +interface QuotaCutoffModalProps { + isOpen: boolean; + onClose: () => void; + /** Label shown in the modal title. */ + connectionName: string; + /** Used in the modal title for context (e.g. "(codex)"). */ + provider: string; + /** + * Windows this connection exposes — discovered from its live quota cache + * so the modal works for any provider with usage data, not just providers + * that registered with quotaPreflight at startup. + */ + windows: QuotaCutoffModalWindow[]; + /** Currently persisted per-window overrides on the connection. */ + current: Record | null; + /** Per-(provider, window) defaults from resilience settings. */ + providerDefaults: Record; + /** Global fallback used when no provider/window default exists. */ + globalDefaultPercent: number; + /** + * Called when the user clicks Save. Receives the patch in the same shape + * the API expects: each window key is either a number (set override) or + * null (clear that window's override). `null` for the whole patch means + * "clear every override" — currently invoked via the "Reset all" button. + */ + onSave: (patch: Record | null) => Promise; +} + +export default function QuotaCutoffModal({ + isOpen, + onClose, + connectionName, + provider, + windows, + current, + providerDefaults, + globalDefaultPercent, + onSave, +}: QuotaCutoffModalProps) { + const t = useTranslations("usage"); + // Local draft: string per window so empty-string means "inherit". + const [drafts, setDrafts] = useState>({}); + const [saving, setSaving] = useState(false); + const [error, setError] = useState(null); + + // Reset drafts whenever the modal opens against a new connection. + useEffect(() => { + if (!isOpen) return; + const initial: Record = {}; + for (const w of windows) { + const persisted = current?.[w.key]; + initial[w.key] = typeof persisted === "number" ? String(persisted) : ""; + } + setDrafts(initial); + setError(null); + }, [isOpen, windows, current]); + + const resolveDefaultFor = (windowKey: string): number => + typeof providerDefaults[windowKey] === "number" + ? providerDefaults[windowKey] + : globalDefaultPercent; + + const buildPatch = (): Record | "invalid" => { + const patch: Record = {}; + for (const w of windows) { + const raw = (drafts[w.key] ?? "").trim(); + if (raw === "") { + // Only emit an explicit null when there was previously an override + // to clear; otherwise just omit the key. + if (current?.[w.key] !== undefined) patch[w.key] = null; + continue; + } + const n = Number(raw); + if (!Number.isInteger(n) || n < 0 || n > 100) return "invalid"; + if (current?.[w.key] !== n) patch[w.key] = n; + } + return patch; + }; + + const handleSave = async () => { + const patch = buildPatch(); + if (patch === "invalid") { + setError(t("quotaThresholdInvalid")); + return; + } + if (Object.keys(patch).length === 0) { + onClose(); + return; + } + setSaving(true); + setError(null); + try { + await onSave(patch); + onClose(); + } catch (err: any) { + setError(err?.message || "Failed to save"); + } finally { + setSaving(false); + } + }; + + const handleResetAll = async () => { + setSaving(true); + setError(null); + try { + await onSave(null); + onClose(); + } catch (err: any) { + setError(err?.message || "Failed to save"); + } finally { + setSaving(false); + } + }; + + const hasAnyOverride = + current !== null && current !== undefined && Object.keys(current).length > 0; + + return ( + + {hasAnyOverride && ( + + )} + + + + } + > +

{t("quotaCutoffsExplainer")}

+
+ {windows.length === 0 && ( +
{t("quotaCutoffsNoWindows")}
+ )} + {windows.map((w) => { + const persisted = current?.[w.key]; + const resolvedDefault = resolveDefaultFor(w.key); + const placeholder = `${resolvedDefault}`; + const isOverride = + typeof persisted === "number" && (drafts[w.key] ?? "") === String(persisted); + return ( +
+
+
{w.displayName}
+
+ {t("quotaCutoffsDefaultHint", { default: resolvedDefault })} +
+
+
+ setDrafts((prev) => ({ ...prev, [w.key]: e.target.value }))} + className={`w-20 px-2 py-1 text-sm text-center rounded-md border bg-transparent text-text-main focus:outline-none focus:border-primary/60 disabled:opacity-50 ${ + isOverride ? "border-primary/40" : "border-border" + }`} + /> + % +
+
+ ); + })} +
+ {error && ( +
+ error + {error} +
+ )} +
+ ); +} diff --git a/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/index.tsx b/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/index.tsx index 5ecbeaab2f..4df3ef897d 100644 --- a/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/index.tsx +++ b/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/index.tsx @@ -18,6 +18,7 @@ import { pickMaskedDisplayValue, pickDisplayValue } from "@/shared/utils/maskEma import useEmailPrivacyStore from "@/store/emailPrivacyStore"; import EmailPrivacyToggle from "@/shared/components/EmailPrivacyToggle"; import ProviderIcon from "@/shared/components/ProviderIcon"; +import QuotaCutoffModal from "./QuotaCutoffModal"; const LS_GROUP_BY = "omniroute:limits:groupBy"; const LS_EXPANDED_GROUPS = "omniroute:limits:expandedGroups"; @@ -80,6 +81,18 @@ function getBarColor(remainingPercentage) { return { bar: "#ef4444", text: "#ef4444", bg: "rgba(239,68,68,0.12)" }; } +// Short label for a quota-window key, used in the inline cutoff summary +// ("session:90% · weekly:80%"). Unknown keys fall back to the key itself, +// shortened to keep the button compact. +function shortWindowLabel(key: string): string { + const map: Record = { + session: "5h", + weekly: "7d", + code_review: "review", + }; + return map[key] || (key.length > 8 ? `${key.slice(0, 7)}…` : key); +} + // Format countdown function formatCountdown(resetAt) { if (!resetAt) return null; @@ -127,6 +140,56 @@ export default function ProviderLimits() { const lastFetchTimeRef = useRef({}); const staleProbeRef = useRef({}); + // Cutoff modal state: connection being edited, the window list captured at + // open time (from quotaData), and the resilience-settings defaults the + // modal renders as placeholders. Kept as separate slices instead of + // mutating the connection object — the window list is UI state, not part + // of the domain. + const [cutoffModalConn, setCutoffModalConn] = useState(null); + const [cutoffModalWindows, setCutoffModalWindows] = useState([]); + const [providerWindowDefaults, setProviderWindowDefaults] = useState< + Record> + >({}); + const [globalThresholdDefault, setGlobalThresholdDefault] = useState(98); + + // Load the resilience-settings defaults once. The endpoint also returns a + // per-provider window registry but we ignore it here — the modal uses the + // connection's live quota cache for window discovery instead. + useEffect(() => { + let alive = true; + fetch("/api/providers/quota-windows") + .then((r) => (r.ok ? r.json() : null)) + .then((data) => { + if (!alive || !data) return; + setProviderWindowDefaults(data.defaults?.providerWindowDefaults || {}); + if (typeof data.defaults?.globalThresholdPercent === "number") { + setGlobalThresholdDefault(data.defaults.globalThresholdPercent); + } + }) + .catch(() => { + /* fail silent — modal still works with empty defaults */ + }); + return () => { + alive = false; + }; + }, []); + + const saveQuotaWindowThresholds = useCallback( + async (connectionId: string, patch: Record | null) => { + const res = await fetch(`/api/providers/${connectionId}`, { + method: "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ quotaWindowThresholds: patch }), + }); + if (!res.ok) throw new Error(`HTTP ${res.status}`); + const data = await res.json(); + const newValue = data?.connection?.quotaWindowThresholds ?? null; + setConnections((prev) => + prev.map((c) => (c.id === connectionId ? { ...c, quotaWindowThresholds: newValue } : c)) + ); + }, + [] + ); const fetchConnections = useCallback(async () => { try { @@ -534,11 +597,14 @@ export default function ProviderLimits() { {/* Table header */}
{t("account")}
{t("modelQuotas")}
{t("lastUsed")}
+
+ {t("quotaThresholdLabel")} +
{t("actions")}
@@ -561,7 +627,7 @@ export default function ProviderLimits() { className="items-center px-4 py-3.5 transition-[background] duration-150 hover:bg-black/[0.03] dark:hover:bg-white/[0.02]" style={{ display: "grid", - gridTemplateColumns: "280px 1fr 128px 48px", + gridTemplateColumns: "280px 1fr 128px 96px 48px", borderBottom: !isLast ? "1px solid var(--color-border)" : "none", }} > @@ -755,6 +821,57 @@ export default function ProviderLimits() { })()} + {/* Quota Threshold Cutoff — button opens modal */} +
+ {(() => { + const overrides = (conn.quotaWindowThresholds || null) as Record< + string, + number + > | null; + const hasOverrides = overrides && Object.keys(overrides).length > 0; + // Window list comes from the connection's own quota cache + // (the same data that drives the Model Quotas bars), so the + // button works for every provider with usage data — not + // just providers that registered with quotaPreflight. + const connectionWindows = (quota?.quotas || []).filter( + (q: any) => q && typeof q.name === "string" && !q.isCredits + ); + const connectionHasWindows = connectionWindows.length > 0; + // Summary: up to 2 entries with short labels; "+N" for the rest. + let label: string = t("quotaCutoffsButtonDefault"); + if (hasOverrides && overrides) { + const entries = Object.entries(overrides); + const visible = entries + .slice(0, 2) + .map(([k, v]) => `${shortWindowLabel(k)}:${v}%`) + .join(" · "); + label = entries.length > 2 ? `${visible} +${entries.length - 2}` : visible; + } + return ( + + ); + })()} +
+ {/* Actions */}
+ + {cutoffModalConn && ( + { + setCutoffModalConn(null); + setCutoffModalWindows([]); + }} + connectionName={ + pickDisplayValue( + [cutoffModalConn.name, cutoffModalConn.displayName, cutoffModalConn.email], + emailsVisible, + cutoffModalConn.provider + ) || cutoffModalConn.provider + } + provider={cutoffModalConn.provider} + windows={cutoffModalWindows.map((q: any) => ({ + key: q.name, + displayName: q.displayName || formatQuotaLabel(q.name), + }))} + current={cutoffModalConn.quotaWindowThresholds || null} + providerDefaults={providerWindowDefaults[cutoffModalConn.provider] || {}} + globalDefaultPercent={globalThresholdDefault} + onSave={async (patch) => { + await saveQuotaWindowThresholds(cutoffModalConn.id, patch); + // Reflect the new state in the modal-open connection ref so the + // button summary updates without closing/reopening. + setCutoffModalConn((prev: any) => { + if (!prev) return prev; + if (patch === null) return { ...prev, quotaWindowThresholds: null }; + const next = { ...(prev.quotaWindowThresholds || {}) }; + for (const [k, v] of Object.entries(patch)) { + if (v === null) delete next[k]; + else next[k] = v; + } + return { + ...prev, + quotaWindowThresholds: Object.keys(next).length === 0 ? null : next, + }; + }); + }} + /> + )} ); } diff --git a/src/app/api/providers/[id]/route.ts b/src/app/api/providers/[id]/route.ts index 4e954bc917..e66555f64f 100644 --- a/src/app/api/providers/[id]/route.ts +++ b/src/app/api/providers/[id]/route.ts @@ -131,6 +131,7 @@ export async function PUT(request: Request, { params }: { params: Promise<{ id: healthCheckInterval, group, maxConcurrent, + quotaWindowThresholds: incomingWindowThresholds, projectId, providerSpecificData: incomingPsd, } = body; @@ -158,6 +159,30 @@ export async function PUT(request: Request, { params }: { params: Promise<{ id: if (healthCheckInterval !== undefined) updateData.healthCheckInterval = healthCheckInterval; if (group !== undefined) updateData.group = group; if (maxConcurrent !== undefined) updateData.maxConcurrent = maxConcurrent; + if (incomingWindowThresholds !== undefined) { + // PATCH semantics: + // • null → clear every per-window override on this connection + // • {} (empty map) → no-op (no keys to merge); existing overrides preserved + // • partial map → merge into the existing map; a `null` value at any + // key clears just that window's override + if (incomingWindowThresholds === null) { + updateData.quotaWindowThresholds = null; + } else { + const existingMap = + existing.quotaWindowThresholds && typeof existing.quotaWindowThresholds === "object" + ? { ...(existing.quotaWindowThresholds as Record) } + : {}; + for (const [window, value] of Object.entries(incomingWindowThresholds)) { + if (value === null) { + delete existingMap[window]; + } else if (typeof value === "number") { + existingMap[window] = value; + } + } + updateData.quotaWindowThresholds = + Object.keys(existingMap).length === 0 ? null : existingMap; + } + } if (projectId !== undefined) updateData.projectId = projectId; // Merge providerSpecificData (partial update — preserve existing keys not sent by caller) diff --git a/src/app/api/providers/quota-windows/route.ts b/src/app/api/providers/quota-windows/route.ts new file mode 100644 index 0000000000..d4a0d842f5 --- /dev/null +++ b/src/app/api/providers/quota-windows/route.ts @@ -0,0 +1,33 @@ +import { NextResponse } from "next/server"; +import { getAllProviderQuotaWindows } from "@omniroute/open-sse/services/quotaPreflight.ts"; +import { getCachedSettings } from "@/lib/localDb"; +import { resolveResilienceSettings } from "@/lib/resilience/settings"; +import { requireManagementAuth } from "@/lib/api/requireManagementAuth"; + +// GET /api/providers/quota-windows +// Returns the named quota windows registered by each provider's quota fetcher, +// plus the resolved per-(provider, window) default thresholds from resilience +// settings. The Provider Limits cutoff modal uses this to know which inputs to +// render per connection and which placeholders to show. Gated by the same +// management-auth middleware as the rest of /api/providers/* because it +// exposes operational routing policy (provider defaults, global cutoff). +export async function GET(request: Request) { + const authError = await requireManagementAuth(request); + if (authError) return authError; + + try { + const windows = getAllProviderQuotaWindows(); + const settings = await getCachedSettings(); + const resilience = resolveResilienceSettings(settings); + return NextResponse.json({ + windows, + defaults: { + globalThresholdPercent: resilience.quotaPreflight.defaultThresholdPercent, + providerWindowDefaults: resilience.quotaPreflight.providerWindowDefaults, + }, + }); + } catch (error) { + console.log("Error fetching quota windows:", error); + return NextResponse.json({ error: "Failed to fetch quota windows" }, { status: 500 }); + } +} diff --git a/src/i18n/messages/en.json b/src/i18n/messages/en.json index 7955339c41..33abde2e32 100644 --- a/src/i18n/messages/en.json +++ b/src/i18n/messages/en.json @@ -4202,6 +4202,17 @@ "lastUsed": "Last Refreshed", "actions": "Actions", "refreshQuota": "Refresh quota", + "quotaThresholdLabel": "Cutoff", + "quotaThresholdInvalid": "Enter an integer 0–100, or leave blank to inherit the default.", + "quotaCutoffsColumnHelp": "Stop using this account when any window's remaining quota drops to its cutoff. Numbers match the dashboard bars (remaining %).", + "quotaCutoffsButtonDefault": "Default", + "quotaCutoffsButtonHelp": "Click to set per-window cutoffs for this account (in remaining %).", + "quotaCutoffsButtonDisabled": "Quota data hasn't loaded yet — refresh this row to enable cutoffs.", + "quotaCutoffsTitle": "Quota cutoffs — {name} ({provider})", + "quotaCutoffsExplainer": "Stop using this account when any window's REMAINING quota drops to the value below — same units as the dashboard bars. Empty inherits the resilience-settings default.", + "quotaCutoffsDefaultHint": "Default: stop at {default}% remaining", + "quotaCutoffsNoWindows": "This provider has no registered quota windows.", + "quotaCutoffsResetAll": "Reset all", "today": "Today", "tomorrow": "Tomorrow", "dayTimeFormat": "{day}, {time}", diff --git a/src/lib/db/core.ts b/src/lib/db/core.ts index 1308c6cb61..d245aa1c6a 100644 --- a/src/lib/db/core.ts +++ b/src/lib/db/core.ts @@ -443,6 +443,16 @@ export function rowToCamel(row: unknown): JsonRecord | null { } catch { result[camelKey] = v; } + } else if (camelKey.endsWith("Json") && typeof v === "string") { + // Convention: any column with a `_json` suffix is JSON-encoded TEXT. + // Surface the parsed object under the friendlier name (key minus the + // "Json" suffix) — e.g. quotaWindowThresholdsJson → quotaWindowThresholds. + const baseKey = camelKey.slice(0, -"Json".length); + try { + result[baseKey] = JSON.parse(v); + } catch { + result[baseKey] = null; + } } else { result[camelKey] = v; } diff --git a/src/lib/db/migrations/056_provider_connection_quota_window_thresholds.sql b/src/lib/db/migrations/056_provider_connection_quota_window_thresholds.sql new file mode 100644 index 0000000000..01823dbf58 --- /dev/null +++ b/src/lib/db/migrations/056_provider_connection_quota_window_thresholds.sql @@ -0,0 +1,14 @@ +-- 056_provider_connection_quota_window_thresholds.sql +-- Per-window quota cutoffs on provider connections. +-- +-- Shape of quota_window_thresholds_json (when set): +-- { "": , ... } +-- +-- A NULL column or missing key means "inherit the resilience-settings default +-- for that provider+window (or the global default if no per-window default)". +-- +-- Window names match the quota keys surfaced by `getUsageForProvider` +-- (open-sse/services/usage.ts) and rendered by the Dashboard › Limits page, +-- so user-set cutoffs and displayed quotas refer to the same windows. + +ALTER TABLE provider_connections ADD COLUMN quota_window_thresholds_json TEXT; diff --git a/src/lib/db/providers.ts b/src/lib/db/providers.ts index 9a023f40e2..4c922f06fb 100644 --- a/src/lib/db/providers.ts +++ b/src/lib/db/providers.ts @@ -45,10 +45,51 @@ function withNullableMaxConcurrent( }; } +// Always surface `quotaWindowThresholds` (possibly null) on the returned +// object — `cleanNulls` strips null values, but the UI needs to see null so +// it can distinguish "no overrides on this connection" from "field was +// never read." Mirrors `withNullableMaxConcurrent`'s contract so create and +// update return the same shape regardless of whether the source had the key +// stripped or carried forward. +function withNullableQuotaWindowThresholds( + record: JsonRecord, + source: JsonRecord | null | undefined +): JsonRecord { + return { + ...record, + quotaWindowThresholds: (source?.quotaWindowThresholds ?? null) as Record | null, + }; +} + function toRecord(value: unknown): JsonRecord { return value && typeof value === "object" ? (value as JsonRecord) : {}; } +// Sanitize the per-window threshold map: keep only 0-100 integer values. +// Called once at each write-path boundary (createProviderConnection + +// updateProviderConnection) so both the in-memory return and the persisted +// row share the same shape. Serialization below trusts this output. +function sanitizeQuotaWindowThresholds(value: unknown): Record | null { + if (value === null || value === undefined) return null; + if (typeof value !== "object" || Array.isArray(value)) return null; + const map: Record = {}; + for (const [key, v] of Object.entries(value as Record)) { + if (typeof v === "number" && Number.isInteger(v) && v >= 0 && v <= 100) { + map[key] = v; + } + } + return Object.keys(map).length === 0 ? null : map; +} + +// Serialize an already-sanitized map for SQLite TEXT storage. Pass `null` to +// store a NULL column; anything else is expected to be the output of +// sanitizeQuotaWindowThresholds above. +function serializeQuotaWindowThresholds(value: unknown): string | null { + if (value === null || value === undefined) return null; + if (typeof value !== "object" || Array.isArray(value)) return null; + return JSON.stringify(value); +} + function toStringOrNull(value: unknown): string | null { return typeof value === "string" ? value : null; } @@ -82,7 +123,12 @@ export async function getProviderConnections(filter: JsonRecord = {}) { const rows = db.prepare(sql).all(params); return rows.map((r) => { const camelRow = rowToCamel(r); - return decryptConnectionFields(withNullableMaxConcurrent(cleanNulls(camelRow), camelRow)); + return decryptConnectionFields( + withNullableQuotaWindowThresholds( + withNullableMaxConcurrent(cleanNulls(camelRow), camelRow), + camelRow + ) + ); }); } @@ -92,7 +138,12 @@ export async function getProviderConnectionById(id: string) { if (!row) return null; const camelRow = rowToCamel(row); - return decryptConnectionFields(withNullableMaxConcurrent(cleanNulls(camelRow), camelRow)); + return decryptConnectionFields( + withNullableQuotaWindowThresholds( + withNullableMaxConcurrent(cleanNulls(camelRow), camelRow), + camelRow + ) + ); } export async function createProviderConnection(data: JsonRecord) { @@ -163,7 +214,10 @@ export async function createProviderConnection(data: JsonRecord) { ); _updateConnectionRow(db, existingId, merged); backupDbFile("pre-write"); - return withNullableMaxConcurrent(cleanNulls(merged), merged); + return withNullableQuotaWindowThresholds( + withNullableMaxConcurrent(cleanNulls(merged), merged), + merged + ); } // Generate name: prefer explicit name, then email, then a stable short-ID label. @@ -226,6 +280,7 @@ export async function createProviderConnection(data: JsonRecord) { "rateLimitProtection", "group", "maxConcurrent", + "quotaWindowThresholds", ]; for (const field of optionalFields) { if (data[field] !== undefined && data[field] !== null) { @@ -235,6 +290,16 @@ export async function createProviderConnection(data: JsonRecord) { if (normalizedProviderSpecificData && Object.keys(normalizedProviderSpecificData).length > 0) { connection.providerSpecificData = normalizedProviderSpecificData; } + // Sanitize the window-thresholds map up front so the in-memory `connection` + // matches the row we're about to insert. The serialize path runs the same + // sanitizer on the way to SQLite. Assigning null (when sanitize collapses + // to no-overrides) keeps the field present on the returned object so the + // UI can tell "field was read, no overrides" apart from "field absent." + if ("quotaWindowThresholds" in connection) { + connection.quotaWindowThresholds = sanitizeQuotaWindowThresholds( + connection.quotaWindowThresholds + ); + } _insertConnectionRow(db, encryptConnectionFields({ ...connection })); const providerId = toStringOrNull(data.provider); @@ -244,7 +309,10 @@ export async function createProviderConnection(data: JsonRecord) { backupDbFile("pre-write"); invalidateDbCache("connections"); // Bust connections read cache - return withNullableMaxConcurrent(cleanNulls(connection), connection); + return withNullableQuotaWindowThresholds( + withNullableMaxConcurrent(cleanNulls(connection), connection), + connection + ); } function _insertConnectionRow(db: DbLike, conn: JsonRecord) { @@ -259,6 +327,7 @@ function _insertConnectionRow(db: DbLike, conn: JsonRecord) { last_tested, api_key, id_token, provider_specific_data, expires_in, display_name, global_priority, default_model, token_type, consecutive_use_count, rate_limit_protection, last_used_at, "group", max_concurrent, + quota_window_thresholds_json, created_at, updated_at ) VALUES ( @id, @provider, @authType, @name, @email, @priority, @isActive, @@ -269,6 +338,7 @@ function _insertConnectionRow(db: DbLike, conn: JsonRecord) { @lastTested, @apiKey, @idToken, @providerSpecificData, @expiresIn, @displayName, @globalPriority, @defaultModel, @tokenType, @consecutiveUseCount, @rateLimitProtection, @lastUsedAt, @group, @maxConcurrent, + @quotaWindowThresholdsJson, @createdAt, @updatedAt ) ` @@ -313,6 +383,7 @@ function _insertConnectionRow(db: DbLike, conn: JsonRecord) { lastUsedAt: conn.lastUsedAt || null, group: conn.group || null, maxConcurrent: conn.maxConcurrent ?? null, + quotaWindowThresholdsJson: serializeQuotaWindowThresholds(conn.quotaWindowThresholds), createdAt: conn.createdAt, updatedAt: conn.updatedAt, }); @@ -339,6 +410,7 @@ function _updateConnectionRow(db: DbLike, id: string, data: JsonRecord) { last_used_at = @lastUsedAt, "group" = @group, max_concurrent = @maxConcurrent, + quota_window_thresholds_json = @quotaWindowThresholdsJson, updated_at = @updatedAt WHERE id = @id ` @@ -383,6 +455,7 @@ function _updateConnectionRow(db: DbLike, id: string, data: JsonRecord) { lastUsedAt: data.lastUsedAt || null, group: data.group || null, maxConcurrent: data.maxConcurrent ?? null, + quotaWindowThresholdsJson: serializeQuotaWindowThresholds(data.quotaWindowThresholds), updatedAt: now, }); } @@ -401,6 +474,14 @@ export async function updateProviderConnection(id: string, data: JsonRecord) { toStringOrNull(merged.provider), merged.providerSpecificData ); + // Mirror the sanitization the create path applies — keep the returned + // object in lockstep with what we persist. + if ("quotaWindowThresholds" in merged) { + const sanitized = sanitizeQuotaWindowThresholds(merged.quotaWindowThresholds); + // For updates we always carry the key forward (even as null) so the read + // path surfaces the cleared state to callers that just patched it. + merged.quotaWindowThresholds = sanitized; + } _updateConnectionRow(db, id, encryptConnectionFields({ ...merged })); backupDbFile("pre-write"); invalidateDbCache("connections"); // Bust connections read cache @@ -414,7 +495,10 @@ export async function updateProviderConnection(id: string, data: JsonRecord) { _reorderConnections(db, providerId); } - return withNullableMaxConcurrent(cleanNulls(merged), merged); + return withNullableQuotaWindowThresholds( + withNullableMaxConcurrent(cleanNulls(merged), merged), + merged + ); } export async function deleteProviderConnection(id: string) { diff --git a/src/lib/resilience/settings.ts b/src/lib/resilience/settings.ts index 23842e5f47..9883c5a9bc 100644 --- a/src/lib/resilience/settings.ts +++ b/src/lib/resilience/settings.ts @@ -40,11 +40,39 @@ export interface WaitForCooldownSettings { maxRetryWaitMs: number; } +export interface QuotaPreflightSettings { + /** + * Global minimum-remaining cutoff (percent, 0-100). A connection is skipped + * when its remaining quota drops to this value or below. Matches the + * dashboard's quota bars (which show REMAINING %, not used %), so the + * number means the same thing in both places. Default: 2 (stop at 2% + * remaining = 98% used). + */ + defaultThresholdPercent: number; + /** + * Global warn threshold (percent, 0-100 remaining %). Fires when remaining + * quota drops to this value or below. Must be HIGHER than the cutoff so + * warnings appear before the block point. Default: 20 (warn at 20% + * remaining = 80% used). + */ + warnThresholdPercent: number; + /** + * Per-(provider, window) defaults for providers that expose multiple quota + * windows (e.g. Codex's session + weekly). Values are minimum-remaining % + * cutoffs. Resolution order, low-to-high precedence: + * defaultThresholdPercent + * → providerWindowDefaults[provider][window] + * → connection.quotaWindowThresholds[window] + */ + providerWindowDefaults: Record>; +} + export interface ResilienceSettings { requestQueue: RequestQueueSettings; connectionCooldown: Record; providerBreaker: Record; waitForCooldown: WaitForCooldownSettings; + quotaPreflight: QuotaPreflightSettings; } export interface ResilienceSettingsPatch { @@ -52,6 +80,7 @@ export interface ResilienceSettingsPatch { connectionCooldown?: Partial>>; providerBreaker?: Partial>>; waitForCooldown?: Partial; + quotaPreflight?: Partial; } function asRecord(value: unknown): JsonRecord { @@ -124,6 +153,16 @@ export const DEFAULT_RESILIENCE_SETTINGS: ResilienceSettings = { maxRetryWaitSec: 30, maxRetryWaitMs: 30000, }, + quotaPreflight: { + // Remaining-% semantics. 2 = "stop when only 2% remaining" (= 98% used). + // Uniform across all providers and windows; operators set per-window + // overrides per connection via the Cutoff modal in Dashboard › Limits, + // or per-(provider, window) globally via the providerWindowDefaults map + // below (no factory seeds — keep behavior consistent across providers). + defaultThresholdPercent: 2, + warnThresholdPercent: 20, + providerWindowDefaults: {}, + }, }; function normalizeRequestQueueSettings( @@ -250,6 +289,65 @@ function normalizeProviderBreakerProfile( }; } +function normalizeProviderWindowDefaults( + next: unknown, + fallback: Record> +): Record> { + // Accept either an explicit object or fall back. Drop providers/windows + // whose values are not a valid 0-100 integer so a malformed setting can't + // accidentally disable cutoffs entirely. + const rawProviders = asRecord(next ?? fallback); + const out: Record> = {}; + for (const [provider, windows] of Object.entries(rawProviders)) { + if (!provider || typeof windows !== "object" || windows === null) continue; + const windowMap: Record = {}; + for (const [windowName, percent] of Object.entries(windows as Record)) { + if (!windowName) continue; + const parsed = + typeof percent === "number" + ? percent + : typeof percent === "string" && percent.trim() !== "" + ? Number(percent) + : NaN; + if (Number.isFinite(parsed)) { + const clamped = Math.min(100, Math.max(0, Math.trunc(parsed))); + windowMap[windowName] = clamped; + } + } + if (Object.keys(windowMap).length > 0) { + out[provider] = windowMap; + } + } + return out; +} + +function normalizeQuotaPreflightSettings( + next: unknown, + fallback: QuotaPreflightSettings +): QuotaPreflightSettings { + const record = asRecord(next); + // Remaining-% semantics: cutoff is the lowest acceptable remaining %, warn + // is the higher "you're getting close" remaining %. So warn MUST be greater + // than cutoff — otherwise the warn log would only fire after the request + // is already blocked. + const defaultThresholdPercent = toInteger( + record.defaultThresholdPercent, + fallback.defaultThresholdPercent, + { min: 0, max: 99 } + ); + const warnRaw = toInteger(record.warnThresholdPercent, fallback.warnThresholdPercent, { + min: 0, + max: 100, + }); + const warnThresholdPercent = + warnRaw <= defaultThresholdPercent ? Math.min(100, defaultThresholdPercent + 1) : warnRaw; + const providerWindowDefaults = normalizeProviderWindowDefaults( + record.providerWindowDefaults, + fallback.providerWindowDefaults + ); + return { defaultThresholdPercent, warnThresholdPercent, providerWindowDefaults }; +} + function normalizeWaitForCooldownSettings( next: unknown, fallback: WaitForCooldownSettings @@ -351,6 +449,7 @@ function buildLegacyFallback(settings: JsonRecord): ResilienceSettings { maxRetryWaitSec: waitMaxRetrySec, maxRetryWaitMs: waitMaxRetrySec * 1000, }, + quotaPreflight: DEFAULT_RESILIENCE_SETTINGS.quotaPreflight, }; } @@ -387,6 +486,10 @@ export function resolveResilienceSettings( current.waitForCooldown, fallback.waitForCooldown ), + quotaPreflight: normalizeQuotaPreflightSettings( + current.quotaPreflight, + fallback.quotaPreflight + ), }; } @@ -420,6 +523,7 @@ export function mergeResilienceSettings( updates.waitForCooldown, current.waitForCooldown ), + quotaPreflight: normalizeQuotaPreflightSettings(updates.quotaPreflight, current.quotaPreflight), }; } diff --git a/src/shared/validation/schemas.ts b/src/shared/validation/schemas.ts index 2939886594..c728787782 100644 --- a/src/shared/validation/schemas.ts +++ b/src/shared/validation/schemas.ts @@ -1589,6 +1589,23 @@ export const updateProviderConnectionSchema = z healthCheckInterval: z.coerce.number().int().min(0).optional(), group: z.union([z.string().max(100), z.null()]).optional(), maxConcurrent: z.union([z.null(), z.coerce.number().int().min(0)]).optional(), + // Per-window quota cutoffs. Map keys are window names (e.g. "window5h", + // "window7d"); values are 0-100 integers, or null to clear that window's + // override (the API route merges this into the existing map and prunes + // null entries before persisting). The whole field set to null clears + // every override on the connection. + quotaWindowThresholds: z + .union([ + z.null(), + z.record( + // Window keys mirror the quota names from getUsageForProvider — + // bound for defense-in-depth so a malicious payload can't ship + // megabyte-long keys that would bloat the DB row. + z.string().min(1).max(64), + z.union([z.null(), z.coerce.number().int().min(0).max(100)]) + ), + ]) + .optional(), projectId: z.union([z.string(), z.null()]).optional(), // Partial patch of per-connection provider-specific settings (e.g. quota toggles) providerSpecificData: z diff --git a/src/sse/handlers/chat.ts b/src/sse/handlers/chat.ts index 03d257337c..dc8b473e69 100644 --- a/src/sse/handlers/chat.ts +++ b/src/sse/handlers/chat.ts @@ -79,6 +79,7 @@ import { import { registerBailianCodingPlanQuotaFetcher } from "@omniroute/open-sse/services/bailianQuotaFetcher.ts"; import { registerCrofUsageFetcher } from "@omniroute/open-sse/services/crofUsageFetcher.ts"; import { registerDeepseekQuotaFetcher } from "@omniroute/open-sse/services/deepseekQuotaFetcher.ts"; +import { registerGenericQuotaFetchers } from "@omniroute/open-sse/services/genericQuotaFetcher.ts"; import { getCooldownAwareRetryDecision, resolveCooldownAwareRetrySettings, @@ -100,6 +101,12 @@ registerCrofUsageFetcher(); // Register DeepSeek balance quota fetcher. // Hooks into quotaPreflight + quotaMonitor so combos can switch accounts before balance is exhausted. registerDeepseekQuotaFetcher(); + +// Register the generic quota fetcher for every other provider that has a +// usage implementation in usage.ts but no bespoke preflight fetcher. This is +// what lets the per-window cutoff modal in Dashboard › Limits actually +// enforce thresholds for Claude / GLM / Cursor / etc., not just Codex. +registerGenericQuotaFetchers(); let combosCachePromise: Promise | null = null; let combosCacheTs = 0; const COMBOS_CACHE_TTL_MS = 10_000; diff --git a/src/sse/services/auth.ts b/src/sse/services/auth.ts index 6ef32dd173..75be7b3d27 100644 --- a/src/sse/services/auth.ts +++ b/src/sse/services/auth.ts @@ -31,7 +31,11 @@ import { } from "@omniroute/open-sse/services/accountFallback.ts"; import { isLocalProvider } from "@omniroute/open-sse/config/providerRegistry.ts"; import { COOLDOWN_MS } from "@omniroute/open-sse/config/constants.ts"; -import { preflightQuota } from "@omniroute/open-sse/services/quotaPreflight.ts"; +import { + preflightQuota, + isQuotaPreflightEnabled, +} from "@omniroute/open-sse/services/quotaPreflight.ts"; +import { resolveResilienceSettings } from "@/lib/resilience/settings"; import { classifyProviderError, PROVIDER_ERROR_TYPES, @@ -46,6 +50,8 @@ type JsonRecord = Record; interface ProviderConnectionView { id: string; + provider: string; + email: string | null; isActive: boolean; rateLimitedUntil: string | null; testStatus: string | null; @@ -65,6 +71,10 @@ interface ProviderConnectionView { errorCode: string | number | null; backoffLevel: number; maxConcurrent: number | null; + // Per-window quota cutoff overrides — null means "no overrides, inherit + // resilience-settings defaults." Read by getProviderCredentialsWithQuotaPreflight + // to decide whether to invoke the upstream usage fetcher. + quotaWindowThresholds: Record | null; } interface RecoverableConnectionState { @@ -121,8 +131,18 @@ function toNullableNumber(value: unknown): number | null { function toProviderConnection(value: unknown): ProviderConnectionView { const row = asRecord(value); + // Only accept the per-window override map when it's a plain object — + // anything else collapses to null so the preflight gate treats it as "no + // overrides set." + const rawThresholds = row.quotaWindowThresholds; + const quotaWindowThresholds: Record | null = + rawThresholds && typeof rawThresholds === "object" && !Array.isArray(rawThresholds) + ? (rawThresholds as Record) + : null; return { id: toStringOrNull(row.id) || "", + provider: toStringOrNull(row.provider) || "", + email: toStringOrNull(row.email), isActive: row.isActive === true, rateLimitedUntil: toStringOrNull(row.rateLimitedUntil), testStatus: toStringOrNull(row.testStatus), @@ -143,6 +163,7 @@ function toProviderConnection(value: unknown): ProviderConnectionView { typeof row.errorCode === "string" || typeof row.errorCode === "number" ? row.errorCode : null, backoffLevel: toNumber(row.backoffLevel, 0), maxConcurrent: toNullableNumber(row.maxConcurrent), + quotaWindowThresholds, }; } @@ -1261,6 +1282,13 @@ export async function getProviderCredentials( ? connection.providerSpecificData.copilotToken : null, providerSpecificData: connection.providerSpecificData, + // Fields the generic quota fetcher (open-sse/services/genericQuotaFetcher.ts) + // needs to delegate to getUsageForProvider for any provider — kept aliased + // (`id` + `connectionId`) for back-compat with callers that already use the + // connectionId name. + id: connection.id, + provider: connection.provider, + email: connection.email, connectionId: connection.id, // Include current status for optimization check testStatus: connection.testStatus, @@ -1270,6 +1298,10 @@ export async function getProviderCredentials( errorCode: connection.errorCode, rateLimitedUntil: connection.rateLimitedUntil, maxConcurrent: connection.maxConcurrent, + // Surface per-window quota overrides so the preflight latency gate in + // getProviderCredentialsWithQuotaPreflight can see them. Without this, + // user-set cutoffs would silently never enforce. + quotaWindowThresholds: connection.quotaWindowThresholds ?? null, }; } finally { if (resolveMutex) resolveMutex(); @@ -1303,6 +1335,19 @@ export async function getProviderCredentialsWithQuotaPreflight( options.excludeConnectionIds ); + const resilience = resolveResilienceSettings(await getCachedSettings()); + const { defaultThresholdPercent, warnThresholdPercent, providerWindowDefaults } = + resilience.quotaPreflight; + const providerWindowMap = providerWindowDefaults[provider] || {}; + const providerHasDefaults = Object.keys(providerWindowMap).length > 0; + // The factory default is "block at 2% remaining" — effectively "right + // before 429." Skipping preflight at that level is a clean no-op. If an + // operator has raised the global to anything stricter (e.g. 20% remaining + // = stop at 80% used), preflight needs to run for every connection so the + // tighter floor is honored. + const FACTORY_NO_OP_REMAINING_PERCENT = 2; + const globalDefaultIsRestrictive = defaultThresholdPercent > FACTORY_NO_OP_REMAINING_PERCENT; + while (true) { const credentials = await getProviderCredentials( provider, @@ -1331,7 +1376,54 @@ export async function getProviderCredentialsWithQuotaPreflight( return credentials; } - const preflight = await preflightQuota(provider, connectionId, credentials); + // Cascading resolver: per-connection override → per-(provider, window) + // default → global default. Used per-window when the fetcher exposes + // multiple windows, and once (with window=null) for single-signal + // fetchers. The warn fallback is uniform — windows don't need their own + // warn levels in v1. + const perConnectionWindowOverrides = + (credentials as { quotaWindowThresholds?: Record | null }) + .quotaWindowThresholds || {}; + + // Latency gate: skip the upstream usage fetch entirely when there's + // nothing to enforce. Preflight is only worth its cost when at least + // one of the following is true: + // • a per-connection override on this row + // • a per-(provider, window) default in resilience settings + // • the legacy `quotaPreflightEnabled` flag in providerSpecificData + // • the global default is stricter than the factory no-op level + // (factory = 2% remaining, basically "right before 429" — anything + // stricter means the operator wants enforcement everywhere) + // Otherwise the resolver would return the factory default for every + // window, and a near-exhausted account would still be caught by the + // normal 429 → cooldown path. + const hasConnectionOverrides = Object.keys(perConnectionWindowOverrides).length > 0; + const legacyForceEnable = isQuotaPreflightEnabled(credentials); + if ( + !hasConnectionOverrides && + !providerHasDefaults && + !legacyForceEnable && + !globalDefaultIsRestrictive + ) { + return credentials; + } + + // Returns the minimum-remaining cutoff for a window — matches the + // dashboard's quota bars so the number the user types in the modal + // means the same thing as the percentage rendered on the bar. + const resolveMinRemainingPercent = (windowName: string | null): number => { + if (windowName !== null) { + const override = perConnectionWindowOverrides[windowName]; + if (typeof override === "number") return override; + const providerDefault = providerWindowMap[windowName]; + if (typeof providerDefault === "number") return providerDefault; + } + return defaultThresholdPercent; + }; + const preflight = await preflightQuota(provider, connectionId, credentials, { + resolveMinRemainingPercent, + resolveWarnRemainingPercent: () => warnThresholdPercent, + }); if (preflight.proceed) { return credentials; } diff --git a/tests/unit/generic-quota-fetcher.test.ts b/tests/unit/generic-quota-fetcher.test.ts new file mode 100644 index 0000000000..f8c219fc43 --- /dev/null +++ b/tests/unit/generic-quota-fetcher.test.ts @@ -0,0 +1,90 @@ +import test from "node:test"; +import assert from "node:assert/strict"; + +const genericModule = await import("../../open-sse/services/genericQuotaFetcher.ts"); +const preflightModule = await import("../../open-sse/services/quotaPreflight.ts"); + +const { convertUsageToQuotaInfo, registerGenericQuotaFetchers } = genericModule; +const { getQuotaFetcher } = preflightModule; + +test("convertUsageToQuotaInfo returns null on null/undefined input", () => { + assert.equal(convertUsageToQuotaInfo(null), null); + assert.equal(convertUsageToQuotaInfo(undefined), null); +}); + +test("convertUsageToQuotaInfo returns null when only an error message is present", () => { + // Auth-expired-style response from getUsageForProvider — fail open. + assert.equal(convertUsageToQuotaInfo({ message: "auth expired" }), null); +}); + +test("convertUsageToQuotaInfo maps remainingPercentage into per-window percentUsed", () => { + const result = convertUsageToQuotaInfo({ + quotas: { + session: { remainingPercentage: 30, resetAt: "2026-05-14T20:00:00Z" }, + weekly: { remainingPercentage: 10, resetAt: "2026-05-21T00:00:00Z" }, + }, + }); + assert.ok(result); + assert.deepEqual(result!.windows, { + session: { percentUsed: 0.7, resetAt: "2026-05-14T20:00:00Z" }, + weekly: { percentUsed: 0.9, resetAt: "2026-05-21T00:00:00Z" }, + }); + // Worst-case percentUsed mirrors what the legacy single-signal field needs. + assert.equal(result!.percentUsed, 0.9); + // Reset time should track the worst-case window so preflight can surface it. + assert.equal(result!.resetAt, "2026-05-21T00:00:00Z"); +}); + +test("convertUsageToQuotaInfo falls back to used/total when remainingPercentage is absent", () => { + const result = convertUsageToQuotaInfo({ + quotas: { session: { used: 45, total: 100, resetAt: null } }, + }); + assert.ok(result); + assert.equal(result!.windows!.session.percentUsed, 0.45); +}); + +test("convertUsageToQuotaInfo skips unlimited and unmeasurable windows", () => { + const result = convertUsageToQuotaInfo({ + quotas: { + session: { remainingPercentage: 50, resetAt: null }, + // No percentage and no used/total → skipped. + unknown_shape: { resetAt: null }, + // Unlimited windows are intentionally ignored — preflight can't block on them. + unlimited_credits: { unlimited: true, remainingPercentage: 99 }, + }, + }); + assert.ok(result); + assert.deepEqual(Object.keys(result!.windows || {}), ["session"]); +}); + +test("convertUsageToQuotaInfo returns null when no windows are measurable", () => { + const result = convertUsageToQuotaInfo({ + quotas: { unlimited_thing: { unlimited: true } }, + }); + assert.equal(result, null); +}); + +test("convertUsageToQuotaInfo clamps remainingPercentage outside 0-100", () => { + const result = convertUsageToQuotaInfo({ + quotas: { + a: { remainingPercentage: 150, resetAt: null }, // clamped to 100 → 0% used + b: { remainingPercentage: -10, resetAt: null }, // clamped to 0 → 100% used + }, + }); + assert.ok(result); + assert.equal(result!.windows!.a.percentUsed, 0); + assert.equal(result!.windows!.b.percentUsed, 1); +}); + +test("registerGenericQuotaFetchers registers Claude and GLM via the generic adapter", () => { + registerGenericQuotaFetchers(); + // Claude has no bespoke fetcher → should be registered. + assert.ok(getQuotaFetcher("claude"), "claude should be registered"); + assert.ok(getQuotaFetcher("glm"), "glm should be registered"); + assert.ok(getQuotaFetcher("zai"), "zai should be registered"); + // Codex has its own dedicated fetcher (registered by codexQuotaFetcher.ts, + // not by the generic registrar) — the generic registrar skips it. We can't + // assert "codex" here without first calling registerCodexQuotaFetcher, + // which would couple this test to chat.ts startup wiring. The skip list + // semantics are exercised by the source code review. +}); diff --git a/tests/unit/provider-connections-quota-threshold.test.ts b/tests/unit/provider-connections-quota-threshold.test.ts new file mode 100644 index 0000000000..d9bfebe935 --- /dev/null +++ b/tests/unit/provider-connections-quota-threshold.test.ts @@ -0,0 +1,164 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-db-quota-windows-")); +process.env.DATA_DIR = TEST_DATA_DIR; + +const core = await import("../../src/lib/db/core.ts"); +const providersDb = await import("../../src/lib/db/providers.ts"); +const { updateProviderConnectionSchema } = await import("../../src/shared/validation/schemas.ts"); + +async function resetStorage() { + core.resetDbInstance(); + + for (let attempt = 0; attempt < 10; attempt++) { + try { + if (fs.existsSync(TEST_DATA_DIR)) { + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + } + break; + } catch (error: any) { + if ((error?.code === "EBUSY" || error?.code === "EPERM") && attempt < 9) { + await new Promise((resolve) => setTimeout(resolve, 50 * (attempt + 1))); + } else { + throw error; + } + } + } + + fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); +} + +test.beforeEach(async () => { + await resetStorage(); +}); + +test.after(async () => { + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +test("createProviderConnection persists quotaWindowThresholds map", async () => { + const created = await providersDb.createProviderConnection({ + provider: "codex", + authType: "apikey", + name: "Codex A", + apiKey: "sk-a", + quotaWindowThresholds: { window5h: 95, window7d: 80 }, + }); + assert.deepEqual(created.quotaWindowThresholds, { window5h: 95, window7d: 80 }); + + const fetched = await providersDb.getProviderConnectionById(created.id); + assert.deepEqual(fetched.quotaWindowThresholds, { window5h: 95, window7d: 80 }); +}); + +test("createProviderConnection with no map yields null on re-read", async () => { + const created = await providersDb.createProviderConnection({ + provider: "codex", + authType: "apikey", + name: "Codex Default", + apiKey: "sk-default", + }); + const fetched = await providersDb.getProviderConnectionById(created.id); + // null/undefined are both acceptable signals for "no overrides". + assert.ok( + fetched.quotaWindowThresholds === null || fetched.quotaWindowThresholds === undefined, + `expected null/undefined, got ${JSON.stringify(fetched.quotaWindowThresholds)}` + ); +}); + +test("updateProviderConnection persists a partial map", async () => { + const created = await providersDb.createProviderConnection({ + provider: "codex", + authType: "apikey", + name: "Codex B", + apiKey: "sk-b", + }); + + const updated = await providersDb.updateProviderConnection(created.id, { + quotaWindowThresholds: { window5h: 50 }, + }); + assert.deepEqual(updated.quotaWindowThresholds, { window5h: 50 }); + + const reread = await providersDb.getProviderConnectionById(created.id); + assert.deepEqual(reread.quotaWindowThresholds, { window5h: 50 }); +}); + +test("updateProviderConnection with explicit null clears the column entirely", async () => { + const created = await providersDb.createProviderConnection({ + provider: "codex", + authType: "apikey", + name: "Codex Clearable", + apiKey: "sk-clear", + quotaWindowThresholds: { window5h: 90 }, + }); + assert.deepEqual(created.quotaWindowThresholds, { window5h: 90 }); + + const cleared = await providersDb.updateProviderConnection(created.id, { + quotaWindowThresholds: null, + }); + // After a clear, the read path should not return a stray map. + assert.ok(cleared.quotaWindowThresholds === null || cleared.quotaWindowThresholds === undefined); + const reread = await providersDb.getProviderConnectionById(created.id); + assert.ok(reread.quotaWindowThresholds === null || reread.quotaWindowThresholds === undefined); +}); + +test("DB serializer drops out-of-range values silently", async () => { + // The DB module sanitizes the map on the way in; values outside 0-100 or + // non-integers are pruned. This is a defense in depth — the Zod schema + // already rejects them at the API boundary, but the DB shouldn't trust. + const created = await providersDb.createProviderConnection({ + provider: "codex", + authType: "apikey", + name: "Codex Sanitize", + apiKey: "sk-san", + quotaWindowThresholds: { window5h: 95, bogus: 999, fractional: 1.5 }, + }); + assert.deepEqual(created.quotaWindowThresholds, { window5h: 95 }); +}); + +test("updateProviderConnectionSchema accepts a valid window map", () => { + const result = updateProviderConnectionSchema.safeParse({ + quotaWindowThresholds: { window5h: 95, window7d: 80 }, + }); + assert.equal(result.success, true); + if (result.success) { + assert.deepEqual(result.data.quotaWindowThresholds, { window5h: 95, window7d: 80 }); + } +}); + +test("updateProviderConnectionSchema accepts null to clear all overrides", () => { + const result = updateProviderConnectionSchema.safeParse({ quotaWindowThresholds: null }); + assert.equal(result.success, true); +}); + +test("updateProviderConnectionSchema accepts null at individual window keys", () => { + // The API route uses key=null as "clear that window's override" while + // preserving the others. + const result = updateProviderConnectionSchema.safeParse({ + quotaWindowThresholds: { window5h: null, window7d: 80 }, + }); + assert.equal(result.success, true); +}); + +test("updateProviderConnectionSchema coerces numeric strings inside the map", () => { + const result = updateProviderConnectionSchema.safeParse({ + quotaWindowThresholds: { window5h: "85" }, + }); + assert.equal(result.success, true); + if (result.success) { + assert.equal(result.data.quotaWindowThresholds?.window5h, 85); + } +}); + +test("updateProviderConnectionSchema rejects out-of-range values", () => { + for (const v of [-1, 101, 150, 1.5]) { + const result = updateProviderConnectionSchema.safeParse({ + quotaWindowThresholds: { window5h: v }, + }); + assert.equal(result.success, false, `expected window5h=${v} to be rejected`); + } +}); diff --git a/tests/unit/quota-preflight.test.ts b/tests/unit/quota-preflight.test.ts index c073a11bff..85bf79b401 100644 --- a/tests/unit/quota-preflight.test.ts +++ b/tests/unit/quota-preflight.test.ts @@ -3,7 +3,13 @@ import assert from "node:assert/strict"; const quotaPreflight = await import("../../open-sse/services/quotaPreflight.ts"); -const { registerQuotaFetcher, isQuotaPreflightEnabled, preflightQuota } = quotaPreflight; +const { + registerQuotaFetcher, + registerQuotaWindows, + getQuotaWindows, + isQuotaPreflightEnabled, + preflightQuota, +} = quotaPreflight; function createConnection(providerSpecificData = {}) { return { providerSpecificData }; @@ -19,18 +25,16 @@ async function withPatchedConsole(methodName, replacement, fn) { } } -test("isQuotaPreflightEnabled reads the provider flag strictly", () => { +test("isQuotaPreflightEnabled reads the provider flag strictly (back-compat helper)", () => { + // The flag itself no longer gates preflightQuota internally — the caller + // in auth.ts decides whether to invoke it. The helper is still exported + // so the caller can honor the legacy force-on flag. assert.equal(isQuotaPreflightEnabled(createConnection({ quotaPreflightEnabled: true })), true); assert.equal(isQuotaPreflightEnabled(createConnection({ quotaPreflightEnabled: "true" })), false); assert.equal(isQuotaPreflightEnabled(createConnection()), false); }); -test("preflightQuota passes through when the feature is disabled", async () => { - const result = await preflightQuota("provider-disabled", "conn-1", createConnection()); - assert.deepEqual(result, { proceed: true }); -}); - -test("preflightQuota passes through when no fetcher is registered", async () => { +test("preflightQuota passes through when no fetcher is registered for the provider", async () => { const result = await preflightQuota( "provider-missing-fetcher", "conn-2", @@ -56,8 +60,11 @@ test("preflightQuota passes through when the fetcher throws or returns null", as }); }); -test("preflightQuota warns but proceeds when usage is above the warning threshold", async () => { - const warnings = []; +// ─── Legacy single-signal path (no windows map on QuotaInfo) ────────────── + +test("preflightQuota (legacy single-signal): warns at 20% remaining by default", async () => { + const warnings: string[] = []; + // 80% used = 20% remaining → hits the default 20% warn threshold. registerQuotaFetcher("provider-warn", async () => ({ used: 80, total: 100, @@ -66,44 +73,194 @@ test("preflightQuota warns but proceeds when usage is above the warning threshol const result = await withPatchedConsole( "warn", - (message) => warnings.push(message), + (message: string) => warnings.push(message), async () => preflightQuota("provider-warn", "conn-5", createConnection({ quotaPreflightEnabled: true })) ); - assert.deepEqual(result, { - proceed: true, - quotaPercent: 0.8, - }); + assert.deepEqual(result, { proceed: true, quotaPercent: 0.8 }); assert.equal(warnings.length, 1); - assert.match(warnings[0], /approaching limit/i); + assert.match(warnings[0], /approaching cutoff/i); + assert.match(warnings[0], /20\.0% remaining/); }); -test("preflightQuota blocks when usage reaches the exhaustion threshold", async () => { - const infos = []; +test("preflightQuota (legacy single-signal): blocks at 2% remaining by default", async () => { + // 99% used = 1% remaining → below the default 2% cutoff → block. registerQuotaFetcher("provider-exhausted", async () => ({ used: 99, total: 100, percentUsed: 0.99, })); + const result = await preflightQuota( + "provider-exhausted", + "conn-6", + createConnection({ quotaPreflightEnabled: true }) + ); + + assert.equal(result.proceed, false); + assert.equal(result.reason, "quota_exhausted"); + assert.equal(result.quotaPercent, 0.99); +}); + +test("preflightQuota (legacy single-signal): resolver override drives the decision (remaining %)", async () => { + // 91% used = 9% remaining. Cutoff = 10 (remaining %) → block (9 ≤ 10). + registerQuotaFetcher("provider-override-block", async () => ({ + used: 91, + total: 100, + percentUsed: 0.91, + })); + + const result = await preflightQuota( + "provider-override-block", + "conn-override-1", + createConnection({ quotaPreflightEnabled: true }), + { + resolveMinRemainingPercent: () => 10, + resolveWarnRemainingPercent: () => 20, + } + ); + assert.equal(result.proceed, false); + assert.equal(result.reason, "quota_exhausted"); +}); + +test("preflightQuota (legacy single-signal): proceeds when remaining is above the cutoff", async () => { + // 89% used = 11% remaining. Cutoff = 10 → proceed (11 > 10). + registerQuotaFetcher("provider-override-pass", async () => ({ + used: 89, + total: 100, + percentUsed: 0.89, + })); + + const result = await preflightQuota( + "provider-override-pass", + "conn-override-2", + createConnection({ quotaPreflightEnabled: true }), + { resolveMinRemainingPercent: () => 10 } + ); + + assert.equal(result.proceed, true); +}); + +// ─── New per-window path (windows map on QuotaInfo) ─────────────────────── + +test("preflightQuota (per-window): blocks if ANY window falls to its cutoff", async () => { + // session: 50% used = 50% remaining (cutoff 5 → ok) + // weekly: 82% used = 18% remaining (cutoff 20 → BLOCK, 18 ≤ 20) + const infos: string[] = []; + registerQuotaFetcher("provider-windows-block", async () => ({ + used: 82, + total: 100, + percentUsed: 0.82, + windows: { + session: { percentUsed: 0.5, resetAt: "2026-05-14T20:00:00Z" }, + weekly: { percentUsed: 0.82, resetAt: "2026-05-21T00:00:00Z" }, + }, + })); + const result = await withPatchedConsole( "info", - (message) => infos.push(message), + (message: string) => infos.push(message), async () => preflightQuota( - "provider-exhausted", - "conn-6", - createConnection({ quotaPreflightEnabled: true }) + "provider-windows-block", + "conn-windows-1", + createConnection({ quotaPreflightEnabled: true }), + { + resolveMinRemainingPercent: (window) => + window === "session" ? 5 : window === "weekly" ? 20 : 2, + } ) ); - assert.deepEqual(result, { - proceed: false, - reason: "quota_exhausted", - quotaPercent: 0.99, - resetAt: null, - }); + assert.equal(result.proceed, false); + assert.equal(result.reason, "quota_exhausted"); + assert.equal(result.quotaPercent, 0.82); + assert.equal(result.resetAt, "2026-05-21T00:00:00Z"); assert.equal(infos.length, 1); - assert.match(infos[0], /switching/i); + assert.match(infos[0], /weekly/); + assert.match(infos[0], /18\.0% remaining/); +}); + +test("preflightQuota (per-window): both above cutoffs → proceed", async () => { + // session: 70% used = 30% remaining (cutoff 5 → ok) + // weekly: 40% used = 60% remaining (cutoff 20 → ok) + registerQuotaFetcher("provider-windows-pass", async () => ({ + used: 70, + total: 100, + percentUsed: 0.7, + windows: { + session: { percentUsed: 0.7, resetAt: null }, + weekly: { percentUsed: 0.4, resetAt: null }, + }, + })); + + const result = await preflightQuota( + "provider-windows-pass", + "conn-windows-2", + createConnection({ quotaPreflightEnabled: true }), + { + resolveMinRemainingPercent: (window) => (window === "session" ? 5 : 20), + } + ); + + assert.equal(result.proceed, true); +}); + +test("preflightQuota (per-window): resolver receives the window name, not null", async () => { + const seenWindows: (string | null)[] = []; + registerQuotaFetcher("provider-windows-resolver-witness", async () => ({ + used: 10, + total: 100, + percentUsed: 0.1, + windows: { + session: { percentUsed: 0.1, resetAt: null }, + weekly: { percentUsed: 0.05, resetAt: null }, + }, + })); + + await preflightQuota( + "provider-windows-resolver-witness", + "conn-windows-3", + createConnection({ quotaPreflightEnabled: true }), + { + resolveMinRemainingPercent: (window) => { + seenWindows.push(window); + return 2; + }, + } + ); + + assert.deepEqual(seenWindows.sort(), ["session", "weekly"]); +}); + +test("preflightQuota (per-window): omitted resolver falls back to the 2% remaining default", async () => { + // weekly at 99% used = 1% remaining < 2% default → block. + registerQuotaFetcher("provider-windows-default", async () => ({ + used: 99, + total: 100, + percentUsed: 0.99, + windows: { + session: { percentUsed: 0.1, resetAt: null }, + weekly: { percentUsed: 0.99, resetAt: null }, + }, + })); + + const result = await preflightQuota( + "provider-windows-default", + "conn-windows-4", + createConnection({ quotaPreflightEnabled: true }) + ); + + assert.equal(result.proceed, false); + assert.equal(result.quotaPercent, 0.99); +}); + +// ─── Window registry ───────────────────────────────────────────────────── + +test("registerQuotaWindows / getQuotaWindows round-trips", () => { + registerQuotaWindows("test-provider", ["a", "b"]); + assert.deepEqual([...getQuotaWindows("test-provider")], ["a", "b"]); + // Unknown provider returns an empty list rather than undefined. + assert.deepEqual([...getQuotaWindows("provider-with-no-registration-anywhere")], []); }); diff --git a/tests/unit/resilience-settings-quota-preflight.test.ts b/tests/unit/resilience-settings-quota-preflight.test.ts new file mode 100644 index 0000000000..bfbea8298e --- /dev/null +++ b/tests/unit/resilience-settings-quota-preflight.test.ts @@ -0,0 +1,128 @@ +import { test } from "node:test"; +import assert from "node:assert/strict"; +import { + DEFAULT_RESILIENCE_SETTINGS, + mergeResilienceSettings, + resolveResilienceSettings, + type ResilienceSettings, +} from "../../src/lib/resilience/settings.ts"; + +function cloneDefaults(): ResilienceSettings { + return structuredClone(DEFAULT_RESILIENCE_SETTINGS); +} + +test("default quotaPreflight thresholds use remaining-% semantics (matches dashboard)", () => { + const settings = cloneDefaults(); + // Block when only 2% remaining (= 98% used). + assert.equal(settings.quotaPreflight.defaultThresholdPercent, 2); + // Warn at 20% remaining (= 80% used). + assert.equal(settings.quotaPreflight.warnThresholdPercent, 20); + // Warn fires earlier than block, so warn % > block %. + assert.ok( + settings.quotaPreflight.warnThresholdPercent > settings.quotaPreflight.defaultThresholdPercent + ); +}); + +test("default providerWindowDefaults is empty — all providers share the global default", () => { + // No factory per-provider seeds; the per-window overrides apply only when + // operators explicitly set them. This keeps the modal placeholder + // consistent across providers (always shows the global default). + const settings = cloneDefaults(); + assert.deepEqual(settings.quotaPreflight.providerWindowDefaults, {}); +}); + +test("resolveResilienceSettings returns defaults when nothing is stored", () => { + const resolved = resolveResilienceSettings({}); + assert.equal(resolved.quotaPreflight.defaultThresholdPercent, 2); + assert.equal(resolved.quotaPreflight.warnThresholdPercent, 20); + assert.deepEqual(resolved.quotaPreflight.providerWindowDefaults, {}); +}); + +test("mergeResilienceSettings: partial defaultThresholdPercent update preserves warnThresholdPercent", () => { + const current = cloneDefaults(); + const next = mergeResilienceSettings(current, { + quotaPreflight: { defaultThresholdPercent: 10 }, + }); + assert.equal(next.quotaPreflight.defaultThresholdPercent, 10); + assert.equal(next.quotaPreflight.warnThresholdPercent, 20); +}); + +test("mergeResilienceSettings clamps defaultThresholdPercent above 99 to 99", () => { + // Block at 100% remaining would mean "always block" — clamp to 99 so it's + // at least conceivable to use the account when it's exactly full. + const next = mergeResilienceSettings(cloneDefaults(), { + quotaPreflight: { defaultThresholdPercent: 150 }, + }); + assert.equal(next.quotaPreflight.defaultThresholdPercent, 99); +}); + +test("warnThresholdPercent is forced ABOVE defaultThresholdPercent when sent in conflict", () => { + // In remaining-% semantics, warn must be > cutoff so warnings fire BEFORE + // the block point (more remaining = warn first, less remaining = block). + const next = mergeResilienceSettings(cloneDefaults(), { + quotaPreflight: { defaultThresholdPercent: 30, warnThresholdPercent: 10 }, + }); + assert( + next.quotaPreflight.warnThresholdPercent > next.quotaPreflight.defaultThresholdPercent, + `expected warn > default, got warn=${next.quotaPreflight.warnThresholdPercent} default=${next.quotaPreflight.defaultThresholdPercent}` + ); + assert.equal(next.quotaPreflight.defaultThresholdPercent, 30); + assert.equal(next.quotaPreflight.warnThresholdPercent, 31); +}); + +test("providerWindowDefaults: arbitrary new provider/window pairs are normalized and stored", () => { + const next = mergeResilienceSettings(cloneDefaults(), { + quotaPreflight: { + providerWindowDefaults: { + codex: { session: 10, weekly: 30 }, + someprovider: { monthly: 40 }, + }, + }, + }); + assert.deepEqual(next.quotaPreflight.providerWindowDefaults.codex, { + session: 10, + weekly: 30, + }); + assert.deepEqual(next.quotaPreflight.providerWindowDefaults.someprovider, { monthly: 40 }); +}); + +test("providerWindowDefaults: out-of-range values are clamped, garbage is pruned", () => { + const next = mergeResilienceSettings(cloneDefaults(), { + quotaPreflight: { + providerWindowDefaults: { + codex: { + session: 150, // clamped to 100 + weekly: -20, // clamped to 0 + // @ts-expect-error: intentionally bogus to ensure pruning + junk: "not a number", + }, + }, + }, + }); + assert.equal(next.quotaPreflight.providerWindowDefaults.codex.session, 100); + assert.equal(next.quotaPreflight.providerWindowDefaults.codex.weekly, 0); + assert.equal( + "junk" in next.quotaPreflight.providerWindowDefaults.codex, + false, + "non-numeric entries should be pruned" + ); +}); + +test("resolveResilienceSettings round-trips a stored providerWindowDefaults map", () => { + const stored = { + resilienceSettings: { + quotaPreflight: { + defaultThresholdPercent: 15, + warnThresholdPercent: 30, + providerWindowDefaults: { codex: { session: 12, weekly: 40 } }, + }, + }, + }; + const resolved = resolveResilienceSettings(stored); + assert.equal(resolved.quotaPreflight.defaultThresholdPercent, 15); + assert.equal(resolved.quotaPreflight.warnThresholdPercent, 30); + assert.deepEqual(resolved.quotaPreflight.providerWindowDefaults.codex, { + session: 12, + weekly: 40, + }); +}); diff --git a/tests/unit/sse-auth.test.ts b/tests/unit/sse-auth.test.ts index 2e13bea7ee..4408d869d8 100644 --- a/tests/unit/sse-auth.test.ts +++ b/tests/unit/sse-auth.test.ts @@ -203,6 +203,114 @@ test("getProviderCredentialsWithQuotaPreflight returns allRateLimited when a for assert.match(selected.lastError, /quota preflight/i); }); +test("getProviderCredentialsWithQuotaPreflight skips the upstream fetcher when no limits are configured", async () => { + // Latency gate regression test. When a connection has no per-window + // overrides AND its provider has no per-(provider, window) defaults seeded + // AND the legacy quotaPreflightEnabled flag isn't set, the dispatch loop + // must NOT call the registered quota fetcher. We assert by registering a + // fetcher that throws if invoked — any invocation surfaces as a test + // failure either via the thrown error or via the connection getting + // skipped (we expect it to pass through cleanly). + const conn = await seedConnection("openai", { + name: "quota-preflight-no-limits", + apiKey: "sk-no-limits", + // Crucially: no quotaPreflightEnabled flag, no overrides. + }); + + const quotaPreflight = await import("../../open-sse/services/quotaPreflight.ts"); + let fetcherCalls = 0; + quotaPreflight.registerQuotaFetcher("openai", async () => { + fetcherCalls++; + throw new Error( + "quota fetcher must not run when no per-window overrides or provider defaults are set" + ); + }); + + const selected = await auth.getProviderCredentialsWithQuotaPreflight("openai"); + + assert.equal((selected as any).connectionId, conn.id); + assert.equal(fetcherCalls, 0, "fetcher should not have been invoked"); +}); + +test("getProviderCredentialsWithQuotaPreflight invokes the fetcher when the global default is restrictive", async () => { + // No per-connection override and no provider-window defaults — but the + // operator has raised the global default cutoff above the factory no-op + // level (2% remaining). Preflight must run so the tighter floor applies. + const conn = await seedConnection("openai", { + name: "quota-preflight-restrictive-global", + apiKey: "sk-restrictive-global", + }); + await settingsDb.updateSettings({ + resilienceSettings: { + quotaPreflight: { + defaultThresholdPercent: 20, // stop at 20% remaining = 80% used + warnThresholdPercent: 30, + providerWindowDefaults: {}, + }, + }, + }); + + const quotaPreflight = await import("../../open-sse/services/quotaPreflight.ts"); + let fetcherCalls = 0; + quotaPreflight.registerQuotaFetcher("openai", async () => { + fetcherCalls++; + return null; + }); + + await auth.getProviderCredentialsWithQuotaPreflight("openai"); + assert.equal( + fetcherCalls, + 1, + "fetcher should run when global default is stricter than the factory no-op level" + ); + + // Reset settings so subsequent tests see factory defaults. + await settingsDb.updateSettings({ resilienceSettings: {} }); + // Verify the gate immediately returns to skip mode. + fetcherCalls = 0; + quotaPreflight.registerQuotaFetcher("openai", async () => { + fetcherCalls++; + throw new Error("must not run with factory global default"); + }); + await auth.getProviderCredentialsWithQuotaPreflight("openai"); + assert.equal(fetcherCalls, 0, "fetcher should not run after settings reset to factory default"); +}); + +test("getProviderCredentialsWithQuotaPreflight invokes the fetcher when an override IS set", async () => { + // Counterpart to the no-limits test: if the connection has a + // quotaWindowThresholds override, preflight must run. + const conn = await seedConnection("openai", { + name: "quota-preflight-with-override", + apiKey: "sk-with-override", + }); + const updated = await providersDb.updateProviderConnection(conn.id, { + quotaWindowThresholds: { primary: 50 }, + }); + // Sanity: the override must be readable on the connection row (this is + // what the dispatch loop reads through getProviderCredentials). + assert.deepEqual( + (updated as any)?.quotaWindowThresholds, + { primary: 50 }, + "override must be persisted on the connection row" + ); + const refetched = await providersDb.getProviderConnectionById(conn.id); + assert.deepEqual( + (refetched as any)?.quotaWindowThresholds, + { primary: 50 }, + "override must round-trip through getProviderConnectionById" + ); + + const quotaPreflight = await import("../../open-sse/services/quotaPreflight.ts"); + let fetcherCalls = 0; + quotaPreflight.registerQuotaFetcher("openai", async () => { + fetcherCalls++; + return null; // null → preflight proceeds, no skip + }); + + await auth.getProviderCredentialsWithQuotaPreflight("openai"); + assert.equal(fetcherCalls, 1, "fetcher should have been invoked exactly once"); +}); + test("getProviderCredentials keeps separate codex affinity per session", async () => { await settingsDb.updateSettings({ fallbackStrategy: "round-robin", stickyRoundRobinLimit: 10 }); const first = await seedConnection("codex", {