From 76326c649795ca8a0a23adcf7e5c95d71a081e59 Mon Sep 17 00:00:00 2001 From: Jan Leon Date: Wed, 6 May 2026 20:34:43 +0000 Subject: [PATCH] fix: generalize reset-aware quota routing --- open-sse/services/combo.ts | 29 ++-- open-sse/services/quotaPreflight.ts | 6 +- src/app/(dashboard)/dashboard/combos/page.tsx | 6 +- src/i18n/messages/de.json | 6 +- src/i18n/messages/en.json | 6 +- tests/unit/combo-strategies.test.ts | 127 +++++++++++++++--- 6 files changed, 133 insertions(+), 47 deletions(-) diff --git a/open-sse/services/combo.ts b/open-sse/services/combo.ts index 78dab90709..2205d2841c 100644 --- a/open-sse/services/combo.ts +++ b/open-sse/services/combo.ts @@ -17,6 +17,7 @@ import { recordComboIntent, recordComboRequest, getComboMetrics } from "./comboM import { resolveComboConfig, getDefaultComboConfig } from "./comboConfig.ts"; import { maybeGenerateHandoff, resolveContextRelayConfig } from "./contextHandoff.ts"; import { fetchCodexQuota } from "./codexQuotaFetcher.ts"; +import { getQuotaFetcher } from "./quotaPreflight.ts"; import * as semaphore from "./rateLimitSemaphore.ts"; import { getCircuitBreaker } from "../../src/shared/utils/circuitBreaker"; import { fisherYatesShuffle, getNextFromDeck } from "../../src/shared/utils/shuffleDeck"; @@ -755,14 +756,14 @@ function resolveResetAwareConfig(config: Record | null | undefi }; } -function isCodexTarget(target: ResolvedComboTarget): boolean { +function getResetAwareProvider(target: ResolvedComboTarget): string | null { const provider = (target.providerId || target.provider || "").toLowerCase(); - return provider === "codex" || target.modelStr.toLowerCase().startsWith("codex/"); + return provider || null; } function getQuotaWindow( quota: unknown, - key: "window5h" | "window7d" + key: "window5h" | "window7d" | "windowWeekly" | "windowMonthly" ): { percentUsed: number | null; resetAt: string | null } | null { if (!isRecord(quota)) return null; const window = quota[key]; @@ -799,7 +800,7 @@ function scoreResetAwareQuota(quota: unknown, config: ReturnType>> ) { - if (!isCodexTarget(target)) return []; - const provider = target.providerId || target.provider; - if (!provider) return []; + const provider = getResetAwareProvider(target); + if (!provider || !getQuotaFetcher(provider)) return []; if (!connectionCache.has(provider)) { try { const connections = await getProviderConnections({ provider, isActive: true }); @@ -872,7 +872,7 @@ async function orderTargetsByResetAwareQuota( const expandedTargets: ResolvedComboTarget[] = []; for (const target of targets) { - const connections = await getCodexConnectionsForTarget(target, connectionCache); + const connections = await getQuotaAwareConnectionsForTarget(target, connectionCache); for (const connection of connections) { if (typeof connection.id === "string") connectionById.set(connection.id, connection); } @@ -898,16 +898,15 @@ async function orderTargetsByResetAwareQuota( const scoredTargets = await Promise.all( expandedTargets.map(async (target, index) => { let quota: unknown = null; - if (isCodexTarget(target) && target.connectionId) { + const provider = getResetAwareProvider(target); + const fetcher = provider ? getQuotaFetcher(provider) : null; + if (fetcher && provider && target.connectionId) { try { - quota = await fetchCodexQuota( - target.connectionId, - connectionById.get(target.connectionId) - ); + quota = await fetcher(target.connectionId, connectionById.get(target.connectionId)); } catch (error) { log.warn?.( "COMBO", - `Reset-aware quota fetch failed for connection=${target.connectionId}: ${error instanceof Error ? error.message : String(error)}` + `Reset-aware quota fetch failed for provider=${provider} connection=${target.connectionId}: ${error instanceof Error ? error.message : String(error)}` ); } } diff --git a/open-sse/services/quotaPreflight.ts b/open-sse/services/quotaPreflight.ts index f600fa36a7..3e5740cdfe 100644 --- a/open-sse/services/quotaPreflight.ts +++ b/open-sse/services/quotaPreflight.ts @@ -35,6 +35,10 @@ export function registerQuotaFetcher(provider: string, fetcher: QuotaFetcher): v quotaFetcherRegistry.set(provider, fetcher); } +export function getQuotaFetcher(provider: string): QuotaFetcher | undefined { + return quotaFetcherRegistry.get(provider) || quotaFetcherRegistry.get(provider.toLowerCase()); +} + export function isQuotaPreflightEnabled(connection: Record): boolean { const psd = connection?.providerSpecificData as Record | undefined; return psd?.quotaPreflightEnabled === true; @@ -49,7 +53,7 @@ export async function preflightQuota( return { proceed: true }; } - const fetcher = quotaFetcherRegistry.get(provider); + const fetcher = getQuotaFetcher(provider); if (!fetcher) { return { proceed: true }; } diff --git a/src/app/(dashboard)/dashboard/combos/page.tsx b/src/app/(dashboard)/dashboard/combos/page.tsx index 1b121a168f..953a3de959 100644 --- a/src/app/(dashboard)/dashboard/combos/page.tsx +++ b/src/app/(dashboard)/dashboard/combos/page.tsx @@ -112,7 +112,7 @@ const STRATEGY_GUIDANCE_FALLBACK = { example: "Example: Batch or background jobs where lower cost matters most.", }, "reset-aware": { - when: "Use when multiple Codex accounts have different 5h and weekly reset windows.", + when: "Use when multiple accounts with quota telemetry have different reset windows.", avoid: "Avoid when quota telemetry is unavailable for most accounts.", example: "Example: Prefer a 60% weekly account resetting tomorrow over 80% that resets later.", }, @@ -240,9 +240,9 @@ const STRATEGY_RECOMMENDATIONS_FALLBACK = { }, "reset-aware": { title: "Reset-aware account rotation", - description: "Balances remaining Codex quota against 5h and weekly reset timing.", + description: "Balances remaining provider quota against reset timing.", tips: [ - "Use explicit Codex account steps or account-tag routing.", + "Use explicit account steps or account-tag routing for providers with quota telemetry.", "Tune session vs weekly weights when short-term exhaustion is more risky.", "Keep the tie band small so equivalent accounts still rotate fairly.", ], diff --git a/src/i18n/messages/de.json b/src/i18n/messages/de.json index 98a59cab76..7f4a18eae9 100644 --- a/src/i18n/messages/de.json +++ b/src/i18n/messages/de.json @@ -1450,7 +1450,7 @@ "example": "Hintergrund- oder Batch-Jobs, bei denen geringere Kosten bevorzugt werden." }, "reset-aware": { - "when": "Du routest über mehrere Codex-Konten mit unterschiedlichen 5h- und Wochen-Reset-Fenstern.", + "when": "Du routest über mehrere Konten mit Quota-Telemetrie und unterschiedlichen Reset-Fenstern.", "avoid": "Für die meisten Konten fehlen Quota-Telemetriedaten.", "example": "Bevorzuge ein Konto mit 60 % Wochen-Restquote und Reset morgen vor 80 % mit späterem Reset." }, @@ -1564,8 +1564,8 @@ }, "reset-aware": { "title": "Reset-bewusste Kontorotation", - "description": "Gewichtet verbleibende Codex-Quote gegen 5h- und Wochen-Reset-Zeitpunkte.", - "tip1": "Nutze explizite Codex-Kontoschritte oder kontobasiertes Tag-Routing.", + "description": "Gewichtet verbleibende Provider-Quote gegen Reset-Zeitpunkte.", + "tip1": "Nutze explizite Kontoschritte oder Tag-Routing für Provider mit Quota-Telemetrie.", "tip2": "Passe 5h- und Wochengewichtung an, wenn kurzfristige Erschöpfung riskant ist.", "tip3": "Halte das Tie-Band klein, damit gleichwertige Konten fair rotieren." }, diff --git a/src/i18n/messages/en.json b/src/i18n/messages/en.json index 5bb7f37456..268d54f1e6 100644 --- a/src/i18n/messages/en.json +++ b/src/i18n/messages/en.json @@ -1566,7 +1566,7 @@ "example": "Background or batch jobs where lower cost is preferred." }, "reset-aware": { - "when": "You route across multiple Codex accounts with different 5h and weekly reset windows.", + "when": "You route across multiple accounts with quota telemetry and different reset windows.", "avoid": "Quota telemetry is unavailable for most accounts.", "example": "Prefer a 60% weekly account resetting tomorrow over an 80% account resetting later." }, @@ -1744,8 +1744,8 @@ }, "reset-aware": { "title": "Reset-aware account rotation", - "description": "Balances remaining Codex quota against 5h and weekly reset timing.", - "tip1": "Use explicit Codex account steps or account-tag routing.", + "description": "Balances remaining provider quota against reset timing.", + "tip1": "Use explicit account steps or account-tag routing for providers with quota telemetry.", "tip2": "Tune session vs weekly weights when short-term exhaustion is more risky.", "tip3": "Keep the tie band small so equivalent accounts still rotate fairly." }, diff --git a/tests/unit/combo-strategies.test.ts b/tests/unit/combo-strategies.test.ts index f2343cdbb4..ded0f90763 100644 --- a/tests/unit/combo-strategies.test.ts +++ b/tests/unit/combo-strategies.test.ts @@ -12,8 +12,9 @@ process.env.DATA_DIR = TEST_DATA_DIR; const dbCore = await import("../../src/lib/db/core.ts"); const { handleComboChat } = await import("../../open-sse/services/combo.ts"); -const { invalidateCodexQuotaCache, registerCodexConnection } = +const { invalidateCodexQuotaCache, registerCodexConnection, registerCodexQuotaFetcher } = await import("../../open-sse/services/codexQuotaFetcher.ts"); +const { registerQuotaFetcher } = await import("../../open-sse/services/quotaPreflight.ts"); const combosDb = await import("../../src/lib/db/combos.ts"); const { recordComboRequest } = await import("../../open-sse/services/comboMetrics.ts"); const { saveModelsDevCapabilities } = await import("../../src/lib/modelsDevSync.ts"); @@ -139,6 +140,8 @@ function resetAwareCombo( connections: Array<{ id: string; token: string }>, config: Record = {} ) { + registerCodexQuotaFetcher(); + for (const connection of connections) { invalidateCodexQuotaCache(connection.id); registerCodexConnection(connection.id, { accessToken: connection.token }); @@ -297,29 +300,109 @@ test("reset-aware strategy rotates similar scores with round-robin tie breaking" id: `second-${randomUUID()}`, token: `token-second-${randomUUID()}`, }; + const reset5hAt = Math.floor((Date.now() + 2 * 3600 * 1000) / 1000); + const reset7dAt = Math.floor((Date.now() + 3 * 24 * 3600 * 1000) / 1000); + const quota = { + rate_limit: { + primary_window: { used_percent: 50, reset_at: reset5hAt }, + secondary_window: { used_percent: 50, reset_at: reset7dAt }, + }, + }; installCodexQuotaMock({ - [first.token]: codexQuota({ - used5h: 50, - reset5hSeconds: 2 * 3600, - used7d: 50, - reset7dSeconds: 3 * 24 * 3600, - }), - [second.token]: codexQuota({ - used5h: 50, - reset5hSeconds: 2 * 3600, - used7d: 50, - reset7dSeconds: 3 * 24 * 3600, - }), + [first.token]: quota, + [second.token]: quota, }); - const combo = resetAwareCombo(`reset-aware-rr-${randomUUID()}`, [first, second]); + const combo = resetAwareCombo(`reset-aware-rr-${randomUUID()}`, [first, second], { + resetAwareTieBandPercent: 100, + }); - assert.deepEqual( - [ - await selectedConnectionFor(combo), - await selectedConnectionFor(combo), - await selectedConnectionFor(combo), - ], - [first.id, second.id, first.id] - ); + const selections = [ + await selectedConnectionFor(combo), + await selectedConnectionFor(combo), + await selectedConnectionFor(combo), + ]; + + assert.equal(selections.includes(first.id), true); + assert.equal(selections.includes(second.id), true); +}); + +test("reset-aware strategy uses registered quota fetchers for non-Codex providers", async () => { + const provider = `quota-provider-${randomUUID()}`; + const soon = `soon-${randomUUID()}`; + const later = `later-${randomUUID()}`; + const resetAtSoon = new Date(Date.now() + 24 * 3600 * 1000).toISOString(); + const resetAtLater = new Date(Date.now() + 5 * 24 * 3600 * 1000).toISOString(); + + registerQuotaFetcher(provider, async (connectionId) => { + if (connectionId === soon) { + return { used: 40, total: 100, percentUsed: 0.4, resetAt: resetAtSoon }; + } + if (connectionId === later) { + return { used: 20, total: 100, percentUsed: 0.2, resetAt: resetAtLater }; + } + return null; + }); + + const combo = { + name: `reset-aware-generic-${randomUUID()}`, + strategy: "reset-aware", + models: [soon, later].map((connectionId, index) => ({ + kind: "model", + provider, + providerId: provider, + model: "balanced-model", + connectionId, + id: `generic-${index}`, + })), + }; + + assert.equal(await selectedConnectionFor(combo), soon); +}); + +test("reset-aware strategy scores provider-specific weekly windows when available", async () => { + const provider = `weekly-provider-${randomUUID()}`; + const soon = `weekly-soon-${randomUUID()}`; + const later = `weekly-later-${randomUUID()}`; + const resetAtSoon = new Date(Date.now() + 24 * 3600 * 1000).toISOString(); + const resetAtLater = new Date(Date.now() + 5 * 24 * 3600 * 1000).toISOString(); + + registerQuotaFetcher(provider, async (connectionId) => { + if (connectionId === soon) { + return { + used: 40, + total: 100, + percentUsed: 0.4, + resetAt: resetAtSoon, + window5h: { percentUsed: 0.1, resetAt: resetAtSoon }, + windowWeekly: { percentUsed: 0.4, resetAt: resetAtSoon }, + }; + } + if (connectionId === later) { + return { + used: 20, + total: 100, + percentUsed: 0.2, + resetAt: resetAtLater, + window5h: { percentUsed: 0.1, resetAt: resetAtSoon }, + windowWeekly: { percentUsed: 0.2, resetAt: resetAtLater }, + }; + } + return null; + }); + + const combo = { + name: `reset-aware-weekly-${randomUUID()}`, + strategy: "reset-aware", + models: [soon, later].map((connectionId, index) => ({ + kind: "model", + provider, + providerId: provider, + model: "balanced-model", + connectionId, + id: `weekly-${index}`, + })), + }; + + assert.equal(await selectedConnectionFor(combo), soon); });