From 11bcdd810a1c971a54aab43563e672a02a4d6c7d Mon Sep 17 00:00:00 2001 From: DavyMassoneto Date: Thu, 5 Mar 2026 10:54:36 -0300 Subject: [PATCH] feat: quota-aware account selection + fix premature model unavailability - Move setModelUnavailable from per-account loop to all-accounts-exhausted path - Clear model unavailability on successful fallback - Add in-memory quota cache with background refresh (5min active, 20min exhausted) - Integrate quota cache in account selection to skip exhausted accounts - Mark accounts as exhausted from 429 when no cached quota data exists - Populate quota cache from dashboard usage endpoint --- src/app/api/usage/[connectionId]/route.ts | 7 + src/domain/quotaCache.ts | 248 ++++++++++++++++++++++ src/instrumentation.ts | 4 + src/sse/handlers/chat.ts | 34 ++- src/sse/services/auth.ts | 40 ++-- 5 files changed, 310 insertions(+), 23 deletions(-) create mode 100644 src/domain/quotaCache.ts diff --git a/src/app/api/usage/[connectionId]/route.ts b/src/app/api/usage/[connectionId]/route.ts index 9655ea3f62..41149c9669 100644 --- a/src/app/api/usage/[connectionId]/route.ts +++ b/src/app/api/usage/[connectionId]/route.ts @@ -4,6 +4,7 @@ import { getUsageForProvider } from "@omniroute/open-sse/services/usage.ts"; import { getExecutor } from "@omniroute/open-sse/executors/index.ts"; import { syncToCloud } from "@/lib/cloudSync"; import { runWithProxyContext } from "@omniroute/open-sse/utils/proxyFetch.ts"; +import { setQuotaCache } from "@/domain/quotaCache"; /** * Sync to cloud if enabled @@ -147,6 +148,12 @@ export async function GET(request: Request, { params }: { params: Promise<{ conn const usage = await runWithProxyContext(proxyInfo?.proxy || null, () => getUsageForProvider(connection) ); + + // Populate quota cache for quota-aware account selection + if (usage?.quotas) { + setQuotaCache(connectionId, connection.provider, usage.quotas); + } + return Response.json(usage); } catch (error) { console.error("[Usage API] Error fetching usage:", error); diff --git a/src/domain/quotaCache.ts b/src/domain/quotaCache.ts new file mode 100644 index 0000000000..ed7d48a20d --- /dev/null +++ b/src/domain/quotaCache.ts @@ -0,0 +1,248 @@ +/** + * Quota Cache — Domain Layer + * + * In-memory cache of provider quota data per connectionId. + * Populated by: + * - Dashboard usage endpoint (GET /api/usage/[connectionId]) + * - 429 responses marking account as exhausted + * + * Background refresh runs every 1 minute: + * - Active accounts (quota > 0%): refetch every 5 minutes + * - Exhausted accounts: refetch every 20 minutes (or immediately after resetAt passes) + * + * @module domain/quotaCache + */ + +import { getUsageForProvider } from "@omniroute/open-sse/services/usage.ts"; +import { getProviderConnectionById, resolveProxyForConnection } from "@/lib/localDb"; +import { runWithProxyContext } from "@omniroute/open-sse/utils/proxyFetch.ts"; + +// ─── Types ────────────────────────────────────────────────────────────────── + +interface QuotaInfo { + remainingPercentage: number; + resetAt: string | null; +} + +interface QuotaCacheEntry { + connectionId: string; + provider: string; + quotas: Record; + fetchedAt: number; + exhausted: boolean; + nextResetAt: string | null; +} + +// ─── Constants ────────────────────────────────────────────────────────────── + +const ACTIVE_TTL_MS = 5 * 60 * 1000; // 5 minutes for active accounts +const EXHAUSTED_TTL_MS = 5 * 60 * 1000; // 5 minutes for 429-sourced entries (no resetAt) +const EXHAUSTED_REFRESH_MS = 20 * 60 * 1000; // 20 minutes: recheck exhausted accounts +const REFRESH_INTERVAL_MS = 60 * 1000; // Background tick every 1 minute + +// ─── State ────────────────────────────────────────────────────────────────── + +const cache = new Map(); +let refreshTimer: ReturnType | null = null; + +// ─── Helpers ──────────────────────────────────────────────────────────────── + +function isExhausted(quotas: Record): boolean { + const entries = Object.values(quotas); + if (entries.length === 0) return false; + return entries.every((q) => q.remainingPercentage <= 0); +} + +function earliestResetAt(quotas: Record): string | null { + let earliest: string | null = null; + for (const q of Object.values(quotas)) { + if (q.resetAt && (!earliest || new Date(q.resetAt) < new Date(earliest))) { + earliest = q.resetAt; + } + } + return earliest; +} + +function normalizeQuotas(rawQuotas: Record): Record { + const result: Record = {}; + for (const [key, q] of Object.entries(rawQuotas)) { + if (q && typeof q === "object") { + result[key] = { + remainingPercentage: + q.remainingPercentage ?? + (q.total ? Math.round(((q.total - (q.used || 0)) / q.total) * 100) : 100), + resetAt: q.resetAt || null, + }; + } + } + return result; +} + +// ─── Public API ───────────────────────────────────────────────────────────── + +/** + * Store quota data for a connection (called by usage endpoint and background refresh). + */ +export function setQuotaCache( + connectionId: string, + provider: string, + rawQuotas: Record +) { + const quotas = normalizeQuotas(rawQuotas); + const exhausted = isExhausted(quotas); + cache.set(connectionId, { + connectionId, + provider, + quotas, + fetchedAt: Date.now(), + exhausted, + nextResetAt: exhausted ? earliestResetAt(quotas) : null, + }); +} + +/** + * Get cached quota entry (returns null if not cached). + */ +export function getQuotaCache(connectionId: string): QuotaCacheEntry | null { + return cache.get(connectionId) || null; +} + +/** + * Check if an account's quota is exhausted based on cached data. + * Returns false if no cache entry exists (unknown = assume available). + */ +export function isAccountQuotaExhausted(connectionId: string): boolean { + const entry = cache.get(connectionId); + if (!entry) return false; + + // If exhausted and we have a resetAt that has passed, consider it available + if (entry.exhausted && entry.nextResetAt) { + if (new Date(entry.nextResetAt).getTime() <= Date.now()) { + return false; // Reset time passed, assume available until refresh confirms + } + } + + // Check TTL based on state + const age = Date.now() - entry.fetchedAt; + if (entry.exhausted) { + // Exhausted entries without resetAt use fixed TTL + if (!entry.nextResetAt && age > EXHAUSTED_TTL_MS) return false; + // Exhausted entries with resetAt stay valid until resetAt or refresh + return true; + } + + // Active entries expire after ACTIVE_TTL + if (age > ACTIVE_TTL_MS) return false; + + return entry.exhausted; +} + +/** + * Mark an account as quota-exhausted from a 429 response (no quota data available). + * Uses 5-minute fixed TTL since we don't know the actual resetAt. + */ +export function markAccountExhaustedFrom429(connectionId: string, provider: string) { + cache.set(connectionId, { + connectionId, + provider, + quotas: {}, + fetchedAt: Date.now(), + exhausted: true, + nextResetAt: null, + }); +} + +// ─── Background Refresh ───────────────────────────────────────────────────── + +async function refreshEntry(entry: QuotaCacheEntry) { + try { + const connection = await getProviderConnectionById(entry.connectionId); + if (!connection || connection.authType !== "oauth" || !connection.isActive) { + cache.delete(entry.connectionId); + return; + } + + const proxyInfo = await resolveProxyForConnection(entry.connectionId); + const usage = await runWithProxyContext(proxyInfo?.proxy || null, () => + getUsageForProvider(connection) + ); + + if (usage?.quotas) { + setQuotaCache(entry.connectionId, entry.provider, usage.quotas); + } + } catch { + // Refresh failed silently — keep stale entry + } +} + +async function backgroundRefreshTick() { + const now = Date.now(); + + for (const entry of cache.values()) { + const age = now - entry.fetchedAt; + + if (entry.exhausted) { + // If resetAt has passed, refetch immediately + if (entry.nextResetAt && new Date(entry.nextResetAt).getTime() <= now) { + refreshEntry(entry); + continue; + } + // Recheck exhausted accounts every 20 minutes + if (age >= EXHAUSTED_REFRESH_MS) { + refreshEntry(entry); + } + } else { + // Refresh active accounts every 5 minutes + if (age >= ACTIVE_TTL_MS) { + refreshEntry(entry); + } + } + } +} + +/** + * Start the background refresh timer. + */ +export function startBackgroundRefresh() { + if (refreshTimer) return; + refreshTimer = setInterval(backgroundRefreshTick, REFRESH_INTERVAL_MS); + // Don't prevent process exit + if (refreshTimer && typeof refreshTimer === "object" && "unref" in refreshTimer) { + (refreshTimer as any).unref(); + } +} + +/** + * Stop the background refresh timer. + */ +export function stopBackgroundRefresh() { + if (refreshTimer) { + clearInterval(refreshTimer); + refreshTimer = null; + } +} + +/** + * Get cache stats (for debugging/dashboard). + */ +export function getQuotaCacheStats() { + const entries: Array<{ + connectionId: string; + provider: string; + exhausted: boolean; + nextResetAt: string | null; + ageMs: number; + }> = []; + + for (const entry of cache.values()) { + entries.push({ + connectionId: entry.connectionId.slice(0, 8) + "...", + provider: entry.provider, + exhausted: entry.exhausted, + nextResetAt: entry.nextResetAt, + ageMs: Date.now() - entry.fetchedAt, + }); + } + + return { total: cache.size, entries }; +} diff --git a/src/instrumentation.ts b/src/instrumentation.ts index 5e6bdc8d4e..9f56ff0f13 100644 --- a/src/instrumentation.ts +++ b/src/instrumentation.ts @@ -34,6 +34,10 @@ export async function register() { const { initApiBridgeServer } = await import("@/lib/apiBridgeServer"); initApiBridgeServer(); + // Quota cache: start background refresh for quota-aware account selection + const { startBackgroundRefresh } = await import("@/domain/quotaCache"); + startBackgroundRefresh(); + // Compliance: Initialize audit_log table + cleanup expired logs try { const { initAuditLog, cleanupExpiredLogs } = await import("@/lib/compliance/index"); diff --git a/src/sse/handlers/chat.ts b/src/sse/handlers/chat.ts index d2b83cdd3b..f67d59c39b 100644 --- a/src/sse/handlers/chat.ts +++ b/src/sse/handlers/chat.ts @@ -31,7 +31,12 @@ import { sanitizeRequest } from "../../shared/utils/inputSanitizer"; // Pipeline integration — wired modules import { getCircuitBreaker, CircuitBreakerOpenError } from "../../shared/utils/circuitBreaker"; -import { isModelAvailable, setModelUnavailable } from "../../domain/modelAvailability"; +import { + isModelAvailable, + setModelUnavailable, + clearModelUnavailability, +} from "../../domain/modelAvailability"; +import { getQuotaCache, markAccountExhaustedFrom429 } from "../../domain/quotaCache"; import { RequestTelemetry, recordTelemetry } from "../../shared/utils/requestTelemetry"; import { generateRequestId } from "../../shared/utils/requestId"; import { recordCost } from "../../domain/costRules"; @@ -127,7 +132,10 @@ export async function handleChat(request: any, clientRawRequest: any = null) { telemetry.startPhase("policy"); const policy = await enforceApiKeyPolicy(request, modelStr); if (policy.rejection) { - log.warn("POLICY", `API key policy rejected: ${modelStr} (key=${policy.apiKeyInfo?.id || "unknown"})`); + log.warn( + "POLICY", + `API key policy rejected: ${modelStr} (key=${policy.apiKeyInfo?.id || "unknown"})` + ); return policy.rejection; } const apiKeyInfo = policy.apiKeyInfo; @@ -243,6 +251,13 @@ async function handleSingleModelChat( const credentials = await getProviderCredentials(provider, excludeConnectionId); if (!credentials || credentials.allRateLimited) { + if (lastStatus === 429 || lastStatus === 503) { + setModelUnavailable(provider, model, 60000, `HTTP ${lastStatus}`); + log.info( + "AVAILABILITY", + `${provider}/${model} marked unavailable — all accounts exhausted (HTTP ${lastStatus})` + ); + } return handleNoCredentials( credentials, excludeConnectionId, @@ -296,22 +311,21 @@ async function handleSingleModelChat( }); if (result.success) { + if (excludeConnectionId) { + clearModelUnavailability(provider, model); + } recordCostIfNeeded(apiKeyInfo, result); if (telemetry) telemetry.startPhase("finalize"); if (telemetry) telemetry.endPhase(); return result.response; } - // Pipeline: Mark model unavailable on repeated failures - if (result.status === 429 || result.status === 503) { - setModelUnavailable(provider, model, 60000, `HTTP ${result.status}`); - log.info( - "AVAILABILITY", - `${provider}/${model} marked unavailable for 60s (HTTP ${result.status})` - ); + // 6. Mark quota-exhausted from 429 if no cached quota data + if (result.status === 429 && !getQuotaCache(credentials.connectionId)) { + markAccountExhaustedFrom429(credentials.connectionId, provider); } - // 6. Fallback to next account + // 7. Fallback to next account const { shouldFallback } = await markAccountUnavailable( credentials.connectionId, result.status, diff --git a/src/sse/services/auth.ts b/src/sse/services/auth.ts index 06d032c319..3f76fc956b 100644 --- a/src/sse/services/auth.ts +++ b/src/sse/services/auth.ts @@ -4,6 +4,7 @@ import { updateProviderConnection, getSettings, } from "@/lib/localDb"; +import { isAccountQuotaExhausted } from "@/domain/quotaCache"; import { isAccountUnavailable, getUnavailableUntil, @@ -197,6 +198,19 @@ export async function getProviderCredentials( return null; } + // Quota-aware: prioritize accounts with available quota + const withQuota = availableConnections.filter((c) => !isAccountQuotaExhausted(c.id)); + const exhaustedQuota = availableConnections.filter((c) => isAccountQuotaExhausted(c.id)); + const orderedConnections = + withQuota.length > 0 ? [...withQuota, ...exhaustedQuota] : availableConnections; + + if (exhaustedQuota.length > 0) { + log.debug( + "AUTH", + `${provider} | quota-aware: ${withQuota.length} with quota, ${exhaustedQuota.length} exhausted` + ); + } + const settings = await getSettings(); const strategy = settings.fallbackStrategy || "fill-first"; @@ -205,7 +219,7 @@ export async function getProviderCredentials( const stickyLimit = toNumber((settings as Record).stickyRoundRobinLimit, 3); // Sort by lastUsed (most recent first) to find current candidate - const byRecency = [...availableConnections].sort((a: any, b: any) => { + const byRecency = [...orderedConnections].sort((a: any, b: any) => { if (!a.lastUsedAt && !b.lastUsedAt) return (a.priority || 999) - (b.priority || 999); if (!a.lastUsedAt) return 1; if (!b.lastUsedAt) return -1; @@ -225,7 +239,7 @@ export async function getProviderCredentials( }); } else { // Pick the least recently used (excluding current if possible) - const sortedByOldest = [...availableConnections].sort((a: any, b: any) => { + const sortedByOldest = [...orderedConnections].sort((a: any, b: any) => { if (!a.lastUsedAt && !b.lastUsedAt) return (a.priority || 999) - (b.priority || 999); if (!a.lastUsedAt) return -1; if (!b.lastUsedAt) return 1; @@ -242,14 +256,14 @@ export async function getProviderCredentials( } } else if (strategy === "p2c") { // Power of Two Choices: pick 2 random, choose the one with fewer failures - if (availableConnections.length <= 2) { - connection = availableConnections[0]; + if (orderedConnections.length <= 2) { + connection = orderedConnections[0]; } else { - const i = Math.floor(Math.random() * availableConnections.length); - let j = Math.floor(Math.random() * (availableConnections.length - 1)); + const i = Math.floor(Math.random() * orderedConnections.length); + let j = Math.floor(Math.random() * (orderedConnections.length - 1)); if (j >= i) j++; - const a = availableConnections[i]; - const b = availableConnections[j]; + const a = orderedConnections[i]; + const b = orderedConnections[j]; // Prefer the one with fewer consecutive uses / better health const scoreA = (a.consecutiveUseCount || 0) + (a.lastError ? 10 : 0); const scoreB = (b.consecutiveUseCount || 0) + (b.lastError ? 10 : 0); @@ -257,11 +271,11 @@ export async function getProviderCredentials( } } else if (strategy === "random") { // Random: Fisher-Yates-inspired random pick - const idx = Math.floor(Math.random() * availableConnections.length); - connection = availableConnections[idx]; + const idx = Math.floor(Math.random() * orderedConnections.length); + connection = orderedConnections[idx]; } else if (strategy === "least-used") { // Least Used: pick the one with oldest lastUsedAt - const sorted = [...availableConnections].sort((a, b) => { + const sorted = [...orderedConnections].sort((a, b) => { if (!a.lastUsedAt && !b.lastUsedAt) return (a.priority || 999) - (b.priority || 999); if (!a.lastUsedAt) return -1; if (!b.lastUsedAt) return 1; @@ -271,13 +285,13 @@ export async function getProviderCredentials( } else if (strategy === "cost-optimized") { // Cost Optimized: sort by priority ascending (lower = cheaper/preferred) // Future: can be enhanced with actual cost data per provider - const sorted = [...availableConnections].sort( + const sorted = [...orderedConnections].sort( (a, b) => (a.priority || 999) - (b.priority || 999) ); connection = sorted[0]; } else { // Default: fill-first (already sorted by priority in getProviderConnections) - connection = availableConnections[0]; + connection = orderedConnections[0]; } return {