mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-08-03 05:45:04 +03:00
feat(limits): per-window quota cutoffs across all providers with usage data (#2267)
feat(limits): per-window quota cutoffs across all providers with usage data (#2267 — thanks @payne0420)
This commit is contained in:
@@ -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]);
|
||||
}
|
||||
|
||||
214
open-sse/services/genericQuotaFetcher.ts
Normal file
214
open-sse/services/genericQuotaFetcher.ts
Normal file
@@ -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<string, CacheEntry>();
|
||||
|
||||
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<string, unknown>;
|
||||
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<string, unknown>;
|
||||
return typeof q.resetAt === "string" ? q.resetAt : null;
|
||||
}
|
||||
|
||||
interface ConnectionInputs {
|
||||
id?: string;
|
||||
provider?: string;
|
||||
accessToken?: string;
|
||||
apiKey?: string;
|
||||
providerSpecificData?: Record<string, unknown>;
|
||||
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<string, unknown>;
|
||||
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<string, { percentUsed: number; resetAt: string | null }> = {};
|
||||
let worstPercent = 0;
|
||||
let worstResetAt: string | null = null;
|
||||
for (const [name, entry] of Object.entries(quotasObj as Record<string, unknown>)) {
|
||||
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<typeof getUsageForProvider>[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);
|
||||
}
|
||||
}
|
||||
@@ -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<string, QuotaWindowInfo>;
|
||||
}
|
||||
|
||||
export type QuotaFetcher = (
|
||||
@@ -26,8 +51,34 @@ export type QuotaFetcher = (
|
||||
connection?: Record<string, unknown>
|
||||
) => Promise<QuotaInfo | null>;
|
||||
|
||||
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<string, readonly string[]>();
|
||||
|
||||
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<string, readonly string[]> {
|
||||
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<string, QuotaFetcher>();
|
||||
|
||||
@@ -44,15 +95,51 @@ export function isQuotaPreflightEnabled(connection: Record<string, unknown>): 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<string, unknown>
|
||||
connection: Record<string, unknown>,
|
||||
thresholds?: PreflightQuotaThresholds
|
||||
): Promise<PreflightQuotaResult> {
|
||||
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`
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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<string, number> | null;
|
||||
/** Per-(provider, window) defaults from resilience settings. */
|
||||
providerDefaults: Record<string, number>;
|
||||
/** 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<string, number | null> | null) => Promise<void>;
|
||||
}
|
||||
|
||||
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<Record<string, string>>({});
|
||||
const [saving, setSaving] = useState(false);
|
||||
const [error, setError] = useState<string | null>(null);
|
||||
|
||||
// Reset drafts whenever the modal opens against a new connection.
|
||||
useEffect(() => {
|
||||
if (!isOpen) return;
|
||||
const initial: Record<string, string> = {};
|
||||
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<string, number | null> | "invalid" => {
|
||||
const patch: Record<string, number | null> = {};
|
||||
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 (
|
||||
<Modal
|
||||
isOpen={isOpen}
|
||||
onClose={onClose}
|
||||
title={t("quotaCutoffsTitle", { name: connectionName, provider })}
|
||||
size="md"
|
||||
footer={
|
||||
<>
|
||||
{hasAnyOverride && (
|
||||
<Button variant="ghost" onClick={handleResetAll} disabled={saving}>
|
||||
{t("quotaCutoffsResetAll")}
|
||||
</Button>
|
||||
)}
|
||||
<Button variant="ghost" onClick={onClose} disabled={saving}>
|
||||
{t("cancel")}
|
||||
</Button>
|
||||
<Button onClick={handleSave} loading={saving}>
|
||||
{t("save")}
|
||||
</Button>
|
||||
</>
|
||||
}
|
||||
>
|
||||
<p className="text-sm text-text-muted mb-4">{t("quotaCutoffsExplainer")}</p>
|
||||
<div className="space-y-3">
|
||||
{windows.length === 0 && (
|
||||
<div className="text-sm text-text-muted italic">{t("quotaCutoffsNoWindows")}</div>
|
||||
)}
|
||||
{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 (
|
||||
<div key={w.key} className="flex items-center justify-between gap-3">
|
||||
<div className="flex-1 min-w-0">
|
||||
<div className="text-sm font-medium text-text-main">{w.displayName}</div>
|
||||
<div className="text-[11px] text-text-muted">
|
||||
{t("quotaCutoffsDefaultHint", { default: resolvedDefault })}
|
||||
</div>
|
||||
</div>
|
||||
<div className="flex items-center gap-1">
|
||||
<input
|
||||
type="number"
|
||||
min={0}
|
||||
max={100}
|
||||
step={1}
|
||||
value={drafts[w.key] ?? ""}
|
||||
placeholder={placeholder}
|
||||
disabled={saving}
|
||||
onChange={(e) => 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"
|
||||
}`}
|
||||
/>
|
||||
<span className="text-xs text-text-muted">%</span>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
})}
|
||||
</div>
|
||||
{error && (
|
||||
<div className="mt-3 text-sm text-red-500 flex items-center gap-1.5">
|
||||
<span className="material-symbols-outlined text-[16px]">error</span>
|
||||
{error}
|
||||
</div>
|
||||
)}
|
||||
</Modal>
|
||||
);
|
||||
}
|
||||
@@ -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<string, string> = {
|
||||
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<any | null>(null);
|
||||
const [cutoffModalWindows, setCutoffModalWindows] = useState<any[]>([]);
|
||||
const [providerWindowDefaults, setProviderWindowDefaults] = useState<
|
||||
Record<string, Record<string, number>>
|
||||
>({});
|
||||
const [globalThresholdDefault, setGlobalThresholdDefault] = useState<number>(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<string, number | null> | 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 */}
|
||||
<div
|
||||
className="items-center px-4 py-2.5 border-b border-border text-[11px] font-semibold uppercase tracking-wider text-text-muted"
|
||||
style={{ display: "grid", gridTemplateColumns: "280px 1fr 128px 48px" }}
|
||||
style={{ display: "grid", gridTemplateColumns: "280px 1fr 128px 96px 48px" }}
|
||||
>
|
||||
<div>{t("account")}</div>
|
||||
<div>{t("modelQuotas")}</div>
|
||||
<div className="text-center">{t("lastUsed")}</div>
|
||||
<div className="text-center" title={t("quotaCutoffsColumnHelp")}>
|
||||
{t("quotaThresholdLabel")}
|
||||
</div>
|
||||
<div className="text-center">{t("actions")}</div>
|
||||
</div>
|
||||
|
||||
@@ -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() {
|
||||
})()}
|
||||
</div>
|
||||
|
||||
{/* Quota Threshold Cutoff — button opens modal */}
|
||||
<div className="flex justify-center items-center">
|
||||
{(() => {
|
||||
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 (
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => {
|
||||
setCutoffModalWindows(connectionWindows);
|
||||
setCutoffModalConn(conn);
|
||||
}}
|
||||
disabled={!connectionHasWindows}
|
||||
title={
|
||||
connectionHasWindows
|
||||
? t("quotaCutoffsButtonHelp")
|
||||
: t("quotaCutoffsButtonDisabled")
|
||||
}
|
||||
className={`px-2 py-1 rounded-md border text-[11px] font-medium tabular-nums transition-colors disabled:opacity-40 disabled:cursor-not-allowed ${
|
||||
hasOverrides
|
||||
? "border-primary/40 text-primary bg-primary/5"
|
||||
: "border-border text-text-muted hover:bg-black/[0.04] dark:hover:bg-white/[0.04]"
|
||||
}`}
|
||||
>
|
||||
{label}
|
||||
</button>
|
||||
);
|
||||
})()}
|
||||
</div>
|
||||
|
||||
{/* Actions */}
|
||||
<div className="flex justify-center gap-0.5">
|
||||
<button
|
||||
@@ -820,6 +937,49 @@ export default function ProviderLimits() {
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{cutoffModalConn && (
|
||||
<QuotaCutoffModal
|
||||
isOpen={!!cutoffModalConn}
|
||||
onClose={() => {
|
||||
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,
|
||||
};
|
||||
});
|
||||
}}
|
||||
/>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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<string, number>) }
|
||||
: {};
|
||||
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)
|
||||
|
||||
33
src/app/api/providers/quota-windows/route.ts
Normal file
33
src/app/api/providers/quota-windows/route.ts
Normal file
@@ -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 });
|
||||
}
|
||||
}
|
||||
@@ -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}",
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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):
|
||||
-- { "<windowName>": <integer 0-100>, ... }
|
||||
--
|
||||
-- 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;
|
||||
@@ -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<string, number> | 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<string, number> | null {
|
||||
if (value === null || value === undefined) return null;
|
||||
if (typeof value !== "object" || Array.isArray(value)) return null;
|
||||
const map: Record<string, number> = {};
|
||||
for (const [key, v] of Object.entries(value as Record<string, unknown>)) {
|
||||
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) {
|
||||
|
||||
@@ -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<string, Record<string, number>>;
|
||||
}
|
||||
|
||||
export interface ResilienceSettings {
|
||||
requestQueue: RequestQueueSettings;
|
||||
connectionCooldown: Record<AuthCategory, ConnectionCooldownProfileSettings>;
|
||||
providerBreaker: Record<AuthCategory, ProviderBreakerProfileSettings>;
|
||||
waitForCooldown: WaitForCooldownSettings;
|
||||
quotaPreflight: QuotaPreflightSettings;
|
||||
}
|
||||
|
||||
export interface ResilienceSettingsPatch {
|
||||
@@ -52,6 +80,7 @@ export interface ResilienceSettingsPatch {
|
||||
connectionCooldown?: Partial<Record<AuthCategory, Partial<ConnectionCooldownProfileSettings>>>;
|
||||
providerBreaker?: Partial<Record<AuthCategory, Partial<ProviderBreakerProfileSettings>>>;
|
||||
waitForCooldown?: Partial<WaitForCooldownSettings>;
|
||||
quotaPreflight?: Partial<QuotaPreflightSettings>;
|
||||
}
|
||||
|
||||
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<string, Record<string, number>>
|
||||
): Record<string, Record<string, number>> {
|
||||
// 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<string, Record<string, number>> = {};
|
||||
for (const [provider, windows] of Object.entries(rawProviders)) {
|
||||
if (!provider || typeof windows !== "object" || windows === null) continue;
|
||||
const windowMap: Record<string, number> = {};
|
||||
for (const [windowName, percent] of Object.entries(windows as Record<string, unknown>)) {
|
||||
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),
|
||||
};
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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<unknown[]> | null = null;
|
||||
let combosCacheTs = 0;
|
||||
const COMBOS_CACHE_TTL_MS = 10_000;
|
||||
|
||||
@@ -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<string, unknown>;
|
||||
|
||||
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<string, number> | 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<string, number> | null =
|
||||
rawThresholds && typeof rawThresholds === "object" && !Array.isArray(rawThresholds)
|
||||
? (rawThresholds as Record<string, number>)
|
||||
: 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<string, number> | 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;
|
||||
}
|
||||
|
||||
90
tests/unit/generic-quota-fetcher.test.ts
Normal file
90
tests/unit/generic-quota-fetcher.test.ts
Normal file
@@ -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.
|
||||
});
|
||||
164
tests/unit/provider-connections-quota-threshold.test.ts
Normal file
164
tests/unit/provider-connections-quota-threshold.test.ts
Normal file
@@ -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`);
|
||||
}
|
||||
});
|
||||
@@ -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")], []);
|
||||
});
|
||||
|
||||
128
tests/unit/resilience-settings-quota-preflight.test.ts
Normal file
128
tests/unit/resilience-settings-quota-preflight.test.ts
Normal file
@@ -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,
|
||||
});
|
||||
});
|
||||
@@ -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", {
|
||||
|
||||
Reference in New Issue
Block a user