From 471052b9047ca21706d171c03140484650e7bbfa Mon Sep 17 00:00:00 2001 From: Minxi Hou Date: Mon, 24 Aug 2026 10:35:02 -0400 Subject: [PATCH] fix(combo): enable genuine quota-aware routing for generic providers (antigravity, claude, etc.) --- open-sse/services/combo/quotaScoring.ts | 6 +- open-sse/services/genericQuotaFetcher.ts | 49 ++++ open-sse/services/quotaPreflight.ts | 10 + src/instrumentation-node.ts | 55 +++++ src/lib/quota/saturationSignals.ts | 11 + .../universal-quota-aware-routing.test.ts | 229 ++++++++++++++++++ 6 files changed, 357 insertions(+), 3 deletions(-) create mode 100644 tests/unit/universal-quota-aware-routing.test.ts diff --git a/open-sse/services/combo/quotaScoring.ts b/open-sse/services/combo/quotaScoring.ts index a107768278..f99adbc561 100644 --- a/open-sse/services/combo/quotaScoring.ts +++ b/open-sse/services/combo/quotaScoring.ts @@ -261,7 +261,7 @@ function getWindowsMapQuotaWindow( ); } -function resolveQuotaWindowByName( +export function resolveQuotaWindowByName( quota: unknown, windowName: ResetWindowName ): QuotaWindowSnapshot | null { @@ -321,8 +321,8 @@ export function scoreResetAwareQuota( if (quota.limitReached === true) return { score: -Infinity }; const overallPercentUsed = clamp01(finiteNumberOrNull(quota.percentUsed) ?? 0.5); - const sessionWindow = getQuotaWindow(quota, "window5h"); - const weeklyWindow = getQuotaWindow(quota, "window7d") || getQuotaWindow(quota, "windowWeekly"); + const sessionWindow = resolveQuotaWindowByName(quota, "session"); + const weeklyWindow = resolveQuotaWindowByName(quota, "weekly"); const sessionRemaining = clamp01(1 - (sessionWindow?.percentUsed ?? overallPercentUsed)); const weeklyRemaining = clamp01(1 - (weeklyWindow?.percentUsed ?? overallPercentUsed)); const sessionScore = scoreQuotaWindow( diff --git a/open-sse/services/genericQuotaFetcher.ts b/open-sse/services/genericQuotaFetcher.ts index bed1409596..80c26dfb1c 100644 --- a/open-sse/services/genericQuotaFetcher.ts +++ b/open-sse/services/genericQuotaFetcher.ts @@ -150,16 +150,65 @@ export function convertUsageToQuotaInfo(usage: unknown): QuotaInfo | null { if (Object.keys(windows).length === 0) return null; + const normalized = normalizeQuotaWindows(windows); + Object.assign(windows, normalized); + return { used: 0, total: 0, percentUsed: worstPercent, resetAt: worstResetAt, windows, + ...normalized, limitReached: worstPercent >= 1 - 1e-9, }; } +/** + * Map provider-native window keys to canonical structural windows so that + * reset-aware / reset-window scoring works without knowing every provider's + * naming convention. + * + * - Claude: "session (5h)" → window5h, "weekly (7d)" → window7d + * - Antigravity: worst per-model quota → window5h; worst *_weekly quota → window7d + */ +function normalizeQuotaWindows( + windows: Record +): Record { + const normalized: Record = {}; + + // Claude-style explicit time windows. + if (windows["session (5h)"] && !normalized.window5h) { + normalized.window5h = windows["session (5h)"]; + } + if (windows["weekly (7d)"] && !normalized.window7d) { + normalized.window7d = windows["weekly (7d)"]; + } + + // Antigravity-style per-model 5h windows: pick the worst (most used) model quota. + const modelWindows = Object.entries(windows).filter( + ([key]) => + key !== "credits" && + !key.endsWith("_weekly") && + !key.startsWith("window") && + !key.includes("(5h)") && + !key.includes("(7d)") + ); + if (modelWindows.length > 0 && !normalized.window5h) { + const worst = modelWindows.reduce((a, b) => (a[1].percentUsed > b[1].percentUsed ? a : b)); + normalized.window5h = worst[1]; + } + + // Antigravity-style weekly family buckets: pick the worst *_weekly quota. + const weeklyWindows = Object.entries(windows).filter(([key]) => key.endsWith("_weekly")); + if (weeklyWindows.length > 0 && !normalized.window7d) { + const worst = weeklyWindows.reduce((a, b) => (a[1].percentUsed > b[1].percentUsed ? a : b)); + normalized.window7d = worst[1]; + } + + return normalized; +} + /** * Fetch quota for a connection by delegating to the appropriate * provider-specific usage fetcher and reshaping its output into the diff --git a/open-sse/services/quotaPreflight.ts b/open-sse/services/quotaPreflight.ts index a7db736c29..a6c7d99aef 100644 --- a/open-sse/services/quotaPreflight.ts +++ b/open-sse/services/quotaPreflight.ts @@ -48,6 +48,16 @@ export interface QuotaInfo { * (e.g. "session", "weekly", "monthly"). */ windows?: Record; + /** + * Structural, canonical window snapshots used by reset-aware / reset-window + * scoring. Providers that expose time-based windows (5h, weekly, monthly) + * populate these in addition to the provider-native `windows` map so the + * scorer does not need to know every provider's key naming convention. + */ + window5h?: QuotaWindowInfo; + window7d?: QuotaWindowInfo; + windowWeekly?: QuotaWindowInfo; + windowMonthly?: QuotaWindowInfo; /** True when the upstream usage endpoint explicitly reports exhausted quota. */ limitReached?: boolean; } diff --git a/src/instrumentation-node.ts b/src/instrumentation-node.ts index c622f4cfa7..37ab73e781 100755 --- a/src/instrumentation-node.ts +++ b/src/instrumentation-node.ts @@ -260,6 +260,57 @@ export async function warmAdaptiveVirtualLanesIntoRuntime(): Promise { } } +/** + * Register bespoke + generic quota fetchers once at Node.js boot. The legacy + * `src/sse/handlers/chat.ts` path registered these at module load, but the + * Next.js App Router production entry (`registerNodejs`) never did, leaving + * `quotaFetcherRegistry` empty for generic providers (antigravity, claude, + * etc.) and causing reset-aware scoring to fall back to the 0.5 dead score. + * + * Each registration call is idempotent; bespoke fetchers are registered first + * so the generic registrar skips providers that already have a dedicated + * fetcher. + */ +export async function registerQuotaFetchers(): Promise { + // Side-effect registrations for agentrouter, freeModel, grokCli, xaiOauth, + // firecrawl (same ordering as the legacy chat.ts path). + await import("@omniroute/open-sse/services/quotaTrackersBatch.ts"); + + const [ + { registerCodexQuotaFetcher }, + { registerBailianCodingPlanQuotaFetcher }, + { registerQwenTokenPlanQuotaFetcher }, + { registerCrofUsageFetcher }, + { registerDeepseekQuotaFetcher }, + { registerOpenrouterQuotaFetcher }, + { registerOpencodeQuotaFetcher }, + { registerGrokWebQuotaFetcher }, + { registerGenericQuotaFetchers }, + ] = await Promise.all([ + import("@omniroute/open-sse/services/codexQuotaFetcher"), + import("@omniroute/open-sse/services/bailianQuotaFetcher"), + import("@omniroute/open-sse/services/qwenTokenPlanQuotaFetcher"), + import("@omniroute/open-sse/services/crofUsageFetcher"), + import("@omniroute/open-sse/services/deepseekQuotaFetcher"), + import("@omniroute/open-sse/services/openrouterQuotaFetcher"), + import("@omniroute/open-sse/services/opencodeQuotaFetcher"), + import("@omniroute/open-sse/services/grokQuotaFetcher"), + import("@omniroute/open-sse/services/genericQuotaFetcher"), + ]); + + registerCodexQuotaFetcher(); + registerBailianCodingPlanQuotaFetcher(); + registerQwenTokenPlanQuotaFetcher(); + registerCrofUsageFetcher(); + registerDeepseekQuotaFetcher(); + registerOpenrouterQuotaFetcher(); + registerOpencodeQuotaFetcher(); + registerGrokWebQuotaFetcher(); + registerGenericQuotaFetchers(); + + console.log("[STARTUP] Quota fetchers registered"); +} + export async function registerNodejs(): Promise { markServerStarting(); @@ -271,6 +322,10 @@ export async function registerNodejs(): Promise { await import("@omniroute/open-sse/index.ts"); console.log("[STARTUP] Global fetch proxy patch initialized"); + // Register quota fetchers early so combo routing can use real quota-aware + // scoring for generic providers in the App Router production runtime. + await registerQuotaFetchers(); + // Guarantee the SQLite singleton — including a sql.js WASM pre-init when // both synchronous drivers (better-sqlite3, node:sqlite) are unavailable — // is ready before ANY other startup step reaches getDbInstance(). This diff --git a/src/lib/quota/saturationSignals.ts b/src/lib/quota/saturationSignals.ts index 50b6b467da..93d3d42690 100644 --- a/src/lib/quota/saturationSignals.ts +++ b/src/lib/quota/saturationSignals.ts @@ -498,6 +498,17 @@ async function fetchGenericSaturation( const result = await fetcher(connectionId, provider); if (result && typeof result === "object") { const obj = result as Record; + + // Prefer the normalized quota shape (handles nested `quotas` map for + // Antigravity / Claude / etc.). Fall back to legacy top-level fields. + const { convertUsageToQuotaInfo } = await import( + "@omniroute/open-sse/services/genericQuotaFetcher" + ); + const quota = convertUsageToQuotaInfo(result); + if (quota && Number.isFinite(quota.percentUsed)) { + return Math.min(1, Math.max(0, quota.percentUsed)); + } + const pct = typeof obj.percentUsed === "number" ? obj.percentUsed diff --git a/tests/unit/universal-quota-aware-routing.test.ts b/tests/unit/universal-quota-aware-routing.test.ts new file mode 100644 index 0000000000..76e5e18f76 --- /dev/null +++ b/tests/unit/universal-quota-aware-routing.test.ts @@ -0,0 +1,229 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import { randomUUID } from "node:crypto"; + +const genericModule = await import("../../open-sse/services/genericQuotaFetcher.ts"); +const preflightModule = await import("../../open-sse/services/quotaPreflight.ts"); +const scoringModule = await import("../../open-sse/services/combo/quotaScoring.ts"); +const strategiesModule = await import("../../open-sse/services/combo/quotaStrategies.ts"); + +const { convertUsageToQuotaInfo, registerGenericQuotaFetchers } = genericModule; +const { getQuotaFetcher, registerQuotaFetcher } = preflightModule; +const { scoreResetAwareQuota, resolveResetAwareConfig } = scoringModule; +const { orderTargetsByResetAwareQuota } = strategiesModule; + +test("registerGenericQuotaFetchers wires antigravity and claude fetchers", () => { + registerGenericQuotaFetchers(); + assert.ok(getQuotaFetcher("antigravity"), "antigravity fetcher should be registered"); + assert.ok(getQuotaFetcher("claude"), "claude fetcher should be registered"); + assert.ok(getQuotaFetcher("agy"), "agy alias fetcher should be registered"); +}); + +test("convertUsageToQuotaInfo normalizes Claude session/weekly into window5h/window7d", () => { + const result = convertUsageToQuotaInfo({ + quotas: { + "session (5h)": { + used: 20, + total: 100, + remainingPercentage: 80, + resetAt: "2026-05-14T20:00:00Z", + }, + "weekly (7d)": { + used: 50, + total: 100, + remainingPercentage: 50, + resetAt: "2026-05-21T00:00:00Z", + }, + }, + }); + assert.ok(result, "should produce a QuotaInfo"); + assert.equal(result!.percentUsed, 0.5, "worst-case percentUsed"); + assert.equal(result!.resetAt, "2026-05-21T00:00:00Z", "resetAt tracks worst window"); + assert.deepEqual(result!.windows["session (5h)"], { + percentUsed: 0.2, + resetAt: "2026-05-14T20:00:00Z", + }); + assert.deepEqual(result!.windows["weekly (7d)"], { + percentUsed: 0.5, + resetAt: "2026-05-21T00:00:00Z", + }); + assert.deepEqual(result!.windows["window5h"], { + percentUsed: 0.2, + resetAt: "2026-05-14T20:00:00Z", + }); + assert.deepEqual(result!.windows["window7d"], { + percentUsed: 0.5, + resetAt: "2026-05-21T00:00:00Z", + }); + assert.equal((result as Record).window5h?.percentUsed, 0.2); + assert.equal((result as Record).window7d?.percentUsed, 0.5); +}); + +test("convertUsageToQuotaInfo normalizes Antigravity model quotas into window5h/window7d", () => { + const result = convertUsageToQuotaInfo({ + quotas: { + "gemini-2-flash": { + used: 200, + total: 1000, + remainingPercentage: 80, + resetAt: "2026-05-14T20:00:00Z", + }, + "gemini-2-pro": { + used: 100, + total: 1000, + remainingPercentage: 90, + resetAt: "2026-05-14T22:00:00Z", + }, + "gemini_weekly": { + used: 500, + total: 1000, + remainingPercentage: 50, + resetAt: "2026-05-21T00:00:00Z", + }, + "claude_and_gpt_weekly": { + used: 100, + total: 1000, + remainingPercentage: 90, + resetAt: "2026-05-21T12:00:00Z", + }, + }, + }); + assert.ok(result, "should produce a QuotaInfo"); + assert.equal(result!.percentUsed, 0.5, "worst-case percentUsed across all windows"); + // 5h window picks the worst (most used) per-model quota. + assert.equal(result!.windows["window5h"].percentUsed, 0.2, "window5h is worst 5h model"); + assert.equal(result!.windows["window5h"].resetAt, "2026-05-14T20:00:00Z"); + // 7d window picks the worst weekly quota. + assert.equal(result!.windows["window7d"].percentUsed, 0.5, "window7d is worst weekly"); + assert.equal(result!.windows["window7d"].resetAt, "2026-05-21T00:00:00Z"); + // Native model keys are preserved. + assert.equal(result!.windows["gemini-2-flash"].percentUsed, 0.2); + assert.equal(result!.windows["gemini_weekly"].percentUsed, 0.5); +}); + +test("scoreResetAwareQuota ranks lower-used Antigravity quota higher and avoids 0.5 dead score", () => { + const resetAt5h = new Date(Date.now() + 24 * 3600 * 1000).toISOString(); + const resetAt7d = new Date(Date.now() + 5 * 24 * 3600 * 1000).toISOString(); + const lowUsage = convertUsageToQuotaInfo({ + quotas: { + "gemini-2-flash": { + used: 800, + total: 1000, + remainingPercentage: 20, + resetAt: resetAt5h, + }, + "gemini_weekly": { + used: 800, + total: 1000, + remainingPercentage: 20, + resetAt: resetAt7d, + }, + }, + }); + const highUsage = convertUsageToQuotaInfo({ + quotas: { + "gemini-2-flash": { + used: 200, + total: 1000, + remainingPercentage: 80, + resetAt: resetAt5h, + }, + "gemini_weekly": { + used: 200, + total: 1000, + remainingPercentage: 80, + resetAt: resetAt7d, + }, + }, + }); + const config = resolveResetAwareConfig({}); + const lowScore = scoreResetAwareQuota(lowUsage, config).score; + const highScore = scoreResetAwareQuota(highUsage, config).score; + assert.ok(highScore > lowScore, "more-remaining quota must score higher"); + assert.ok(lowScore !== 0.5, "low-usage score must not be the dead 0.5 fallback"); + assert.ok(highScore !== 0.5, "high-usage score must not be the dead 0.5 fallback"); +}); + +test("orderTargetsByResetAwareQuota prefers Antigravity connection with more remaining quota", async () => { + registerGenericQuotaFetchers(); + const low = `low-${randomUUID()}`; + const high = `high-${randomUUID()}`; + const resetAt5h = new Date(Date.now() + 24 * 3600 * 1000).toISOString(); + const resetAt7d = new Date(Date.now() + 5 * 24 * 3600 * 1000).toISOString(); + + registerQuotaFetcher("antigravity", async (connectionId) => { + if (connectionId === low) { + return convertUsageToQuotaInfo({ + quotas: { + "gemini-2-flash": { + used: 800, + total: 1000, + remainingPercentage: 20, + resetAt: resetAt5h, + }, + "gemini_weekly": { + used: 800, + total: 1000, + remainingPercentage: 20, + resetAt: resetAt7d, + }, + }, + }); + } + if (connectionId === high) { + return convertUsageToQuotaInfo({ + quotas: { + "gemini-2-flash": { + used: 200, + total: 1000, + remainingPercentage: 80, + resetAt: resetAt5h, + }, + "gemini_weekly": { + used: 200, + total: 1000, + remainingPercentage: 80, + resetAt: resetAt7d, + }, + }, + }); + } + return null; + }); + + const targets = [low, high].map((connectionId, index) => ({ + providerId: "antigravity", + provider: "antigravity", + model: "gemini-2-flash", + modelStr: "antigravity/gemini-2-flash", + connectionId, + executionKey: `test-${index}`, + index, + })); + + const ordered = await orderTargetsByResetAwareQuota(targets, "test", {}, { warn: () => {} }); + assert.equal(ordered[0].connectionId, high, "connection with more remaining quota must be first"); +}); + +test("getSaturation reads saturation from nested quotas map for generic providers", async () => { + const saturationModule = await import("../../src/lib/quota/saturationSignals.ts"); + saturationModule._clearSaturationCache(); + saturationModule.__setGenericUsageFetcherForTests(async () => ({ + quotas: { + "gemini-2-flash": { + used: 800, + total: 1000, + remainingPercentage: 20, + resetAt: null, + }, + }, + })); + + const saturation = await saturationModule.getSaturation("conn-ag", "antigravity", { + unit: "percent", + window: "5h", + }); + assert.equal(saturation, 0.8, "saturation must reflect the worst nested quota"); + + saturationModule.__setGenericUsageFetcherForTests(null); +});