From 0d31fd3d244c3391fbe938367b86372da9a81eaf Mon Sep 17 00:00:00 2001 From: Diego Rodrigues de Sa e Souza Date: Fri, 18 Sep 2026 11:57:17 -0300 Subject: [PATCH] fix(quota): resolve plan from pool's primary connection (#13876) (#14042) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Multi-connection Quota Sharing pools resolve a DIFFERENT provider plan depending on which member connection actually served a request (write path, enforceQuotaShare/recordConsumption) vs. the pool's primary connection (dashboard read path, /api/quota/pools/[id]/usage). The wizard's "Limite" step PUTs a manual plan override only to the primary connection, so any other pool member fell back to a different (catalog/empty) plan shape. Since the quota_consumption dimension key is poolId:unit:window, a different unit/window meant recordConsumption wrote to a bucket the dashboard never read, so real traffic served via a non-primary connection never appeared as "consumed". Fix: resolve the plan from the pool's canonical primary connection (pool.connectionId) in both enforceQuotaShare and recordConsumption, matching the dashboard's read path. recordConsumption now keeps the matched pool object (not just its id) so it can reach connectionId. getSaturation(input.connectionId, ...) is untouched — that signal is legitimately per-connection. --- ...876-quota-share-multiconn-plan-mismatch.md | 1 + src/lib/quota/enforce.ts | 25 +++- .../quota-share-multiconn-enforce-cap.test.ts | 84 +++++++++++++ ...uota-share-multiconn-plan-mismatch.test.ts | 114 ++++++++++++++++++ 4 files changed, 218 insertions(+), 6 deletions(-) create mode 100644 changelog.d/fixes/13876-quota-share-multiconn-plan-mismatch.md create mode 100644 tests/unit/quota-share-multiconn-enforce-cap.test.ts create mode 100644 tests/unit/quota-share-multiconn-plan-mismatch.test.ts diff --git a/changelog.d/fixes/13876-quota-share-multiconn-plan-mismatch.md b/changelog.d/fixes/13876-quota-share-multiconn-plan-mismatch.md new file mode 100644 index 0000000000..0c9392a979 --- /dev/null +++ b/changelog.d/fixes/13876-quota-share-multiconn-plan-mismatch.md @@ -0,0 +1 @@ +- fix(quota): resolve Quota Sharing plan/limits from the pool's canonical primary connection in `enforceQuotaShare` and `recordConsumption`, not the serving connection — a multi-connection pool whose "Limite" wizard override only ever lands on the primary connection was writing/checking consumption under a different dimension key when a request was actually served via a non-primary pool member, so the dashboard's "consumed" amount never reflected real traffic (#13876) diff --git a/src/lib/quota/enforce.ts b/src/lib/quota/enforce.ts index 3d8d844f5a..d46c0d3c65 100644 --- a/src/lib/quota/enforce.ts +++ b/src/lib/quota/enforce.ts @@ -112,7 +112,14 @@ export async function enforceQuotaShare(input: EnforceInput): Promise { + core.resetDbInstance(); + if (fs.existsSync(TEST_DATA_DIR)) { + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true, maxRetries: 5, retryDelay: 100 }); + } +}); + +test("multi-connection pool: a cap set on the primary connection's plan is enforced for a request served by a non-primary member", async () => { + const { enforceQuotaShare, recordConsumption } = await import("../../src/lib/quota/enforce.ts"); + + const PROVIDER = "quota-13876-enforce-test-provider"; // not in the built-in catalog on purpose + + poolsDb.createPool({ + connectionId: "conn-primary-2", + connectionIds: ["conn-primary-2", "conn-secondary-2"], + name: "Repro Pool #13876 (enforce)", + allocations: [{ apiKeyId: "key-shared-2", weight: 100, policy: "hard" }], + }); + + // The operator sets a small daily token cap via the wizard — only the + // primary connection gets the manual override. The pool has 2 member + // connections, so the effective (global) limit is limit × accountCount — + // set the per-account limit small enough that a single request still + // saturates the whole pool. + plansDb.upsertPlan( + "conn-primary-2", + PROVIDER, + [{ unit: "tokens", window: "daily", limit: 100 }], + "manual" + ); + + // A request served via the secondary connection consumes past the + // effective (accountCount-multiplied) cap of 200 tokens. + await recordConsumption({ + apiKeyId: "key-shared-2", + connectionId: "conn-secondary-2", + provider: PROVIDER, + cost: { tokens: 1500, usd: 0, requests: 1 }, + }); + + // The NEXT request, also served via the secondary connection, must now be + // blocked by the cap that was configured only on the primary connection's + // plan — matching what the dashboard (which reads via pool.connectionId) + // shows as consumed. + const decision = await enforceQuotaShare({ + apiKeyId: "key-shared-2", + connectionId: "conn-secondary-2", + provider: PROVIDER, + }); + + assert.equal( + decision.kind, + "block", + `expected the primary connection's cap to block a request served via the secondary connection, got ${JSON.stringify(decision)}` + ); +}); diff --git a/tests/unit/quota-share-multiconn-plan-mismatch.test.ts b/tests/unit/quota-share-multiconn-plan-mismatch.test.ts new file mode 100644 index 0000000000..307e40e5ec --- /dev/null +++ b/tests/unit/quota-share-multiconn-plan-mismatch.test.ts @@ -0,0 +1,114 @@ +/** + * Repro for issue #13876 — Quota Sharing dashboard "consumed" does not update. + * + * Root cause: the "Limite" wizard step writes a manual plan override (PUT + * /api/quota/plans/[connectionId]) ONLY for the pool's primary connection + * (connectionIds[0]) — see PoolWizard.tsx. But a multi-connection pool + * actually serves traffic through ANY member connection (fallback / + * round-robin), and: + * - the POST-hook (enforce.ts recordConsumption) resolved the plan with + * resolvePlan(input.connectionId, provider) — the connection that + * actually served THIS request. + * - the dashboard read path (GET /api/quota/pools/[id]/usage) resolves the + * plan with resolvePlan(pool.connectionId, provider) — always the + * primary connection. + * + * When the primary connection has a manual override and a non-primary member + * connection does not (falls through to catalog / empty dimensions), the two + * paths used DIFFERENT dimension shapes. Consumption got written under one + * dimensionKey (poolId:unit:window) while the dashboard read a different + * one, so real traffic never appeared as "consumed" even though it happened. + * + * Fix: recordConsumption / enforceQuotaShare now resolve the plan from the + * pool's canonical (primary) connection, matching the dashboard read path. + */ + +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-quota-13876-")); +process.env.DATA_DIR = TEST_DATA_DIR; + +const core = await import("../../src/lib/db/core.ts"); +const poolsDb = await import("../../src/lib/db/quotaPools.ts"); +const plansDb = await import("../../src/lib/db/providerPlans.ts"); + +test.after(async () => { + core.resetDbInstance(); + if (fs.existsSync(TEST_DATA_DIR)) { + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true, maxRetries: 5, retryDelay: 100 }); + } +}); + +test("multi-connection pool: consumption served via a non-primary connection is visible on the usage dashboard", async () => { + const { recordConsumption } = await import("../../src/lib/quota/enforce.ts"); + const { resolvePlan } = await import("../../src/lib/quota/planResolver.ts"); + const { getSqliteQuotaStore } = await import("../../src/lib/quota/sqliteQuotaStore.ts"); + + const PROVIDER = "quota-13876-test-provider"; // not in the built-in catalog on purpose + + // Pool with 2 member connections sharing quota for one API key. + // conn-primary is connectionIds[0] — the one the wizard PUTs the plan to. + // conn-secondary is a normal pool member that can also serve requests + // (fallback/round-robin across accounts in the same pool). + const pool = poolsDb.createPool({ + connectionId: "conn-primary", + connectionIds: ["conn-primary", "conn-secondary"], + name: "Repro Pool #13876", + allocations: [{ apiKeyId: "key-shared", weight: 100, policy: "hard" }], + }); + + // The operator opens the wizard's "Limite" step and sets a custom plan — + // this PUTs only to /api/quota/plans/conn-primary (PoolWizard.tsx primaryConnectionId). + plansDb.upsertPlan( + "conn-primary", + PROVIDER, + [{ unit: "tokens", window: "daily", limit: 100000 }], + "manual" + ); + + // Sanity: conn-secondary has no override and the provider has no catalog + // entry, so resolvePlan for it falls through to the empty plan. + const secondaryPlan = resolvePlan("conn-secondary", PROVIDER); + assert.equal( + secondaryPlan.dimensions.length, + 0, + "precondition: conn-secondary must resolve to a DIFFERENT (empty) plan than conn-primary" + ); + + // A real request gets served through conn-secondary (the pool's fallback + // member) and completes successfully — the POST-hook fires. + await recordConsumption({ + apiKeyId: "key-shared", + connectionId: "conn-secondary", + provider: PROVIDER, + cost: { tokens: 5000, usd: 0, requests: 1 }, + }); + + // The Quota Sharing dashboard reads usage the same way the REST route does: + // resolve the plan from pool.connectionId (always the primary) and read + // poolUsageWithDimensions with those dimensions. + const primaryPlan = resolvePlan(pool.connectionId, PROVIDER); + assert.equal( + primaryPlan.dimensions.length, + 1, + "primary connection plan should have the override" + ); + + const store = getSqliteQuotaStore(); + const snapshot = await store.poolUsageWithDimensions(pool.id, primaryPlan.dimensions); + + const tokenDim = snapshot.dimensions.find((d) => d.unit === "tokens"); + assert.ok(tokenDim, "tokens dimension should be present in the snapshot"); + + // Consumption served via conn-secondary must be visible under the pool's + // canonical dimension key, exactly as the dashboard reads it. + assert.equal( + tokenDim!.consumedTotal, + 5000, + `dashboard should show the 5000 tokens actually consumed via conn-secondary, got ${tokenDim!.consumedTotal}` + ); +});