diff --git a/changelog.d/fixes/10686-combo-quota-token-limit-await.md b/changelog.d/fixes/10686-combo-quota-token-limit-await.md new file mode 100644 index 0000000000..a9e7b910e3 --- /dev/null +++ b/changelog.d/fixes/10686-combo-quota-token-limit-await.md @@ -0,0 +1 @@ +- **Combo routing:** await each connection's token limit before reserving quota. The old lookup treated the `Promise` as a connection and dropped `rateLimitOverrides.tpm` ([#10686](https://github.com/diegosouzapw/OmniRoute/pull/10686)). diff --git a/open-sse/services/combo.ts b/open-sse/services/combo.ts index 71d1f47384..b6ae5d534a 100644 --- a/open-sse/services/combo.ts +++ b/open-sse/services/combo.ts @@ -87,7 +87,7 @@ import { selectQuotaShareTarget } from "./combo/quotaShareStrategy.ts"; import { makeConnectionConcurrencyResolver, lookupPositiveCap } from "./combo/concurrencyCaps.ts"; import { acquireQuotaShareConcurrencySlot } from "./combo/quotaShareConcurrency.ts"; import { canAffordRequest } from "../../src/lib/quota/quotaScheduler.ts"; -import { getCachedProviderConnectionById } from "../../src/lib/localDb.ts"; +import { getCachedProviderConnectionById } from "../../src/lib/db/readCache.ts"; import { orderTargetsByEvalScores } from "./evalRouting.ts"; /** @@ -96,11 +96,13 @@ import { orderTargetsByEvalScores } from "./evalRouting.ts"; * keeps the previously recorded limit (or 0 for a fresh row, meaning "no * budget enforced"). */ -function resolveTargetTokenLimit(target: { connectionId?: string | null }): number | undefined { +async function resolveTargetTokenLimit(target: { + connectionId?: string | null; +}): Promise { const connectionId = target?.connectionId; if (!connectionId) return undefined; try { - const connection = getCachedProviderConnectionById(connectionId); + const connection = await getCachedProviderConnectionById(connectionId); const overrides = (connection as { rateLimitOverrides?: Record | null } | null) ?.rateLimitOverrides; const tpm = overrides?.tpm; @@ -3082,7 +3084,7 @@ async function handleRoundRobinCombo({ try { const { reserveQuota } = await import("../../src/lib/quota/quotaScheduler.ts"); reserveQuota(target.connectionId, modelStr, attemptBody as Record, { - tokenLimit: resolveTargetTokenLimit(target), + tokenLimit: await resolveTargetTokenLimit(target), }); } catch { // best-effort only diff --git a/tests/unit/combo-quota-token-limit.test.ts b/tests/unit/combo-quota-token-limit.test.ts new file mode 100644 index 0000000000..53ecdb91c1 --- /dev/null +++ b/tests/unit/combo-quota-token-limit.test.ts @@ -0,0 +1,76 @@ +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-combo-quota-token-limit-")); +const ORIGINAL_DATA_DIR = process.env.DATA_DIR; +const ORIGINAL_QUOTA_ROUTING = process.env.OMNIROUTE_QUOTA_AWARE_ROUTING; +process.env.DATA_DIR = TEST_DATA_DIR; +process.env.OMNIROUTE_QUOTA_AWARE_ROUTING = "1"; + +const { handleComboChat } = await import("../../open-sse/services/combo.ts"); +const { getProviderQuota } = await import("../../src/lib/quota/providerQuotaState.ts"); +const providersDb = await import("../../src/lib/db/providers.ts"); +const dbCore = await import("../../src/lib/db/core.ts"); + +const log = { info() {}, warn() {}, debug() {}, error() {} }; + +test.after(() => { + dbCore.resetDbInstance(); + if (ORIGINAL_DATA_DIR === undefined) delete process.env.DATA_DIR; + else process.env.DATA_DIR = ORIGINAL_DATA_DIR; + if (ORIGINAL_QUOTA_ROUTING === undefined) delete process.env.OMNIROUTE_QUOTA_AWARE_ROUTING; + else process.env.OMNIROUTE_QUOTA_AWARE_ROUTING = ORIGINAL_QUOTA_ROUTING; + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +test("round-robin quota reservation keeps the connection token limit", async () => { + const connection = await providersDb.createProviderConnection({ + provider: "openai", + authType: "apikey", + name: "quota token limit test", + apiKey: "sk-quota-token-limit-test", + rateLimitOverrides: { tpm: 5000 }, + }); + assert.ok(connection?.id); + + const model = "openai/gpt-4o"; + const combo = { + name: "quota-token-limit-test", + strategy: "round-robin", + config: { maxRetries: 0, disableSessionStickiness: true }, + models: [ + { + kind: "model", + provider: "openai", + providerId: "openai", + model: "gpt-4o", + connectionId: connection.id, + id: "quota-token-limit-test-target", + }, + ], + }; + + const result = await handleComboChat({ + body: { + model, + messages: [{ role: "user", content: "reserve this request" }], + max_tokens: 8, + stream: false, + }, + combo, + allCombos: [combo], + isModelAvailable: async () => true, + settings: {}, + log, + handleSingleModel: async () => Response.json({ choices: [{ message: { content: "ok" } }] }), + }); + + assert.equal(result.ok, true); + const snapshot = getProviderQuota(connection.id, model); + assert.equal(snapshot?.known, true); + assert.equal(snapshot?.tokenLimit, 5000); + assert.ok((snapshot?.tokensUsed ?? 0) > 0); +});