diff --git a/src/lib/db/quotaConsumption.ts b/src/lib/db/quotaConsumption.ts index adb66e9f58..355ba76242 100644 --- a/src/lib/db/quotaConsumption.ts +++ b/src/lib/db/quotaConsumption.ts @@ -200,6 +200,58 @@ export function listConsumptionForPool(poolId: string, limit: number): Consumpti }); } +// --------------------------------------------------------------------------- +// Pool-wide aggregate +// --------------------------------------------------------------------------- + +interface BucketPairRow { + api_key_id: string; + curr: number; + prev: number; +} + +/** + * Sum the consumed values for a given (dimensionKey, currentBucketIndex) and + * (dimensionKey, currentBucketIndex - 1) across ALL api_key_id values. + * + * Returns { currTotal, prevTotal } so the caller can apply the sliding-window + * formula once with the pool-wide totals instead of summing per-key. + * + * @param dimensionKey "::" string — same format as consume/peek. + * @param currentBucket floor(nowMs / windowMs) — caller must pass the same value. + */ +export function sumPoolDimension( + dimensionKey: string, + currentBucket: number +): { currTotal: number; prevTotal: number } { + const prevBucket = currentBucket - 1; + + interface SumRow { + total: number; + } + + const currRow = getDb() + .prepare( + `SELECT COALESCE(SUM(consumed), 0) AS total + FROM quota_consumption + WHERE dimension_key = ? AND bucket_index = ?` + ) + .get(dimensionKey, currentBucket); + + const prevRow = getDb() + .prepare( + `SELECT COALESCE(SUM(consumed), 0) AS total + FROM quota_consumption + WHERE dimension_key = ? AND bucket_index = ?` + ) + .get(dimensionKey, prevBucket); + + return { + currTotal: currRow?.total ?? 0, + prevTotal: prevRow?.total ?? 0, + }; +} + // --------------------------------------------------------------------------- // GC // --------------------------------------------------------------------------- diff --git a/src/lib/localDb.ts b/src/lib/localDb.ts index dbf27aee7a..21b6ee1776 100755 --- a/src/lib/localDb.ts +++ b/src/lib/localDb.ts @@ -549,6 +549,7 @@ export { getBucket, incrementBucket, getPair, + sumPoolDimension, gcOlderThan as gcQuotaConsumption, } from "./db/quotaConsumption"; export { diff --git a/src/lib/quota/redisQuotaStore.ts b/src/lib/quota/redisQuotaStore.ts index 2fa75faa6b..e2497022c6 100644 --- a/src/lib/quota/redisQuotaStore.ts +++ b/src/lib/quota/redisQuotaStore.ts @@ -165,6 +165,51 @@ export class RedisQuotaStore implements QuotaStore { return slidingWindowEffective(curr, prev, nowMs, windowMs); } + /** + * Return the real pool-wide consumption for a dimension in the current + * sliding window, summed across ALL apiKeyIds in the pool's allocations. + * + * Strategy: pool/allocation metadata lives in SQLite (F2), so we fetch the + * allocation list for the pool, then issue a single MGET for all + * (apiKeyId, curr-bucket) and (apiKeyId, prev-bucket) keys. The Redis keys + * are per-key counters — there is no pool-level aggregate key — so we read + * each key's two buckets and sum the raw values before applying the + * sliding-window formula once on the totals. + * + * If the pool does not exist or has no allocations, returns 0. + * Keys absent in Redis are treated as 0 (MGET returns null). + */ + async poolConsumedTotal(poolId: string, dim: DimensionKey): Promise { + const pool = getPool(poolId); + if (!pool || pool.allocations.length === 0) return 0; + + const nowMs = Date.now(); + const dimKey = dimensionKeyToString(dim); + const windowMs = WINDOW_MS[dim.window]; + const currentBucket = Math.floor(nowMs / windowMs); + const prevBucket = currentBucket - 1; + + const client = await this.client(); + + // Build [currKey0, prevKey0, currKey1, prevKey1, ...] for all allocated keys + const redisKeys: string[] = []; + for (const alloc of pool.allocations) { + redisKeys.push(bucketKey(alloc.apiKeyId, dimKey, currentBucket)); + redisKeys.push(bucketKey(alloc.apiKeyId, dimKey, prevBucket)); + } + + const values = await client.mget(...redisKeys); + + let currTotal = 0; + let prevTotal = 0; + for (let i = 0; i < values.length; i += 2) { + currTotal += parseFloat(values[i] ?? "0") || 0; + prevTotal += parseFloat(values[i + 1] ?? "0") || 0; + } + + return slidingWindowEffective(currTotal, prevTotal, nowMs, windowMs); + } + /** * Aggregate pool usage. Pool and allocation metadata come from SQLite (F2); * rolling counters come from Redis. diff --git a/src/lib/quota/sqliteQuotaStore.ts b/src/lib/quota/sqliteQuotaStore.ts index 38ae17b0d5..a72526294d 100644 --- a/src/lib/quota/sqliteQuotaStore.ts +++ b/src/lib/quota/sqliteQuotaStore.ts @@ -21,6 +21,7 @@ import { getBucket, incrementBucket, getPair, + sumPoolDimension, } from "@/lib/localDb"; import { WINDOW_MS, dimensionKeyToString } from "./dimensions"; import type { DimensionKey } from "./dimensions"; @@ -112,6 +113,24 @@ export class SqliteQuotaStore implements QuotaStore { return slidingWindowEffective(curr, prev, nowMs, windowMs); } + /** + * Return the real pool-wide consumption for a dimension in the current + * sliding window, summed across ALL apiKeyIds that share the same + * dimensionKey (i.e. same poolId + unit + window). + * + * Uses the same 2-bucket sliding-window formula as peek(), applied once + * to the pool totals so the result is consistent with per-key semantics. + */ + async poolConsumedTotal(poolId: string, dim: DimensionKey): Promise { + const nowMs = Date.now(); + const dimKey = dimensionKeyToString(dim); + const windowMs = WINDOW_MS[dim.window]; + const currentBucket = Math.floor(nowMs / windowMs); + + const { currTotal, prevTotal } = sumPoolDimension(dimKey, currentBucket); + return slidingWindowEffective(currTotal, prevTotal, nowMs, windowMs); + } + /** * Return a PoolUsageSnapshot for the given pool, aggregating per-key * consumption across all dimensions and computing fairShare / deficit / diff --git a/src/lib/quota/types.ts b/src/lib/quota/types.ts index 53bd0daf2a..63b19c2b39 100644 --- a/src/lib/quota/types.ts +++ b/src/lib/quota/types.ts @@ -34,6 +34,16 @@ export interface ConsumeResult { export interface QuotaStore { consume(apiKeyId: string, dim: DimensionKey, cost: number): Promise; peek(apiKeyId: string, dim: DimensionKey): Promise; + /** + * Return the real pool-wide consumption for a dimension in the current + * sliding window — i.e. the sum of each key's effective consumption across + * ALL apiKeyIds that have contributed to (poolId, unit, window). + * + * Unlike the per-key saturation signal (which can be 0 for countable units + * whose hard-cap has never been set), this reflects actual spent units so + * the enforce path can block when the pool total hits the plan limit. + */ + poolConsumedTotal(poolId: string, dim: DimensionKey): Promise; poolUsage(poolId: string): Promise; clear(apiKeyId: string, dim: DimensionKey): Promise; } diff --git a/tests/unit/quota-store-pool-total.test.ts b/tests/unit/quota-store-pool-total.test.ts new file mode 100644 index 0000000000..a6d44d4611 --- /dev/null +++ b/tests/unit/quota-store-pool-total.test.ts @@ -0,0 +1,119 @@ +/** + * tests/unit/quota-store-pool-total.test.ts + * + * Coverage for QuotaStore.poolConsumedTotal(): + * - Two keys in same pool → sum equals individual contributions. + * - Different pool or different dim → 0 (isolation). + * - peek() for a single key still returns only that key's own value (no regression). + */ + +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-pool-total-")); +process.env.DATA_DIR = TEST_DATA_DIR; + +const core = await import("../../src/lib/db/core.ts"); + +async function resetStorage() { + core.resetDbInstance(); + for (let attempt = 0; attempt < 10; attempt++) { + try { + if (fs.existsSync(TEST_DATA_DIR)) { + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + } + break; + } catch (err: unknown) { + const e = err as { code?: string }; + if ((e?.code === "EBUSY" || e?.code === "EPERM") && attempt < 9) { + await new Promise((resolve) => setTimeout(resolve, 50 * (attempt + 1))); + } else { + throw err; + } + } + } + fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); +} + +test.beforeEach(async () => { + await resetStorage(); +}); + +test.after(async () => { + core.resetDbInstance(); + if (fs.existsSync(TEST_DATA_DIR)) { + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + } +}); + +// ─── Main scenario: two keys in same pool ──────────────────────────────────── + +test("poolConsumedTotal: keyA(3) + keyB(2) in same pool → total === 5", async () => { + const { SqliteQuotaStore } = await import("../../src/lib/quota/sqliteQuotaStore.ts"); + const store = new SqliteQuotaStore(); + + const dim = { poolId: "p1", unit: "requests" as const, window: "5h" as const }; + + await store.consume("keyA", dim, 3); + await store.consume("keyB", dim, 2); + + const total = await store.poolConsumedTotal("p1", dim); + + // Both consumes happen in the same bucket (milliseconds apart) so + // prev=0, elapsed≈0 → total ≈ 5. Allow small delta for timing. + assert.ok(total > 4.9, `Expected >4.9, got ${total}`); + assert.ok(total <= 5, `Expected <=5, got ${total}`); +}); + +// ─── Isolation: different pool ─────────────────────────────────────────────── + +test("poolConsumedTotal: different poolId → returns 0 (isolation)", async () => { + const { SqliteQuotaStore } = await import("../../src/lib/quota/sqliteQuotaStore.ts"); + const store = new SqliteQuotaStore(); + + const dimP1 = { poolId: "p1", unit: "requests" as const, window: "5h" as const }; + const dimP2 = { poolId: "p2", unit: "requests" as const, window: "5h" as const }; + + await store.consume("keyA", dimP1, 10); + + // Pool p2 has no consumption → should be 0 + const total = await store.poolConsumedTotal("p2", dimP2); + assert.equal(total, 0, `Expected 0 for different pool, got ${total}`); +}); + +// ─── Isolation: different dimension (unit) ─────────────────────────────────── + +test("poolConsumedTotal: different unit dim → returns 0 (isolation)", async () => { + const { SqliteQuotaStore } = await import("../../src/lib/quota/sqliteQuotaStore.ts"); + const store = new SqliteQuotaStore(); + + const dimRequests = { poolId: "p1", unit: "requests" as const, window: "5h" as const }; + const dimTokens = { poolId: "p1", unit: "tokens" as const, window: "5h" as const }; + + await store.consume("keyA", dimRequests, 7); + + // tokens dim has no consumption → should be 0 + const total = await store.poolConsumedTotal("p1", dimTokens); + assert.equal(total, 0, `Expected 0 for different unit, got ${total}`); +}); + +// ─── No regression: peek still returns per-key value ───────────────────────── + +test("poolConsumedTotal: peek(keyA) still returns only keyA's own consumption", async () => { + const { SqliteQuotaStore } = await import("../../src/lib/quota/sqliteQuotaStore.ts"); + const store = new SqliteQuotaStore(); + + const dim = { poolId: "p1", unit: "requests" as const, window: "5h" as const }; + + await store.consume("keyA", dim, 3); + await store.consume("keyB", dim, 2); + + const peekA = await store.peek("keyA", dim); + + // keyA's own peek should be ≈3, not 5 + assert.ok(peekA > 2.9, `Expected >2.9, got ${peekA}`); + assert.ok(peekA <= 3, `Expected <=3, got ${peekA}`); +});