diff --git a/open-sse/handlers/chatCore.ts b/open-sse/handlers/chatCore.ts index dc0286ff50..b8832a8b32 100644 --- a/open-sse/handlers/chatCore.ts +++ b/open-sse/handlers/chatCore.ts @@ -4307,6 +4307,14 @@ export async function handleChatCore({ providerResponse.status, model ); + + // Store rate-limit headers for quota saturation signals + try { + const { storeRateLimitHeaders } = await import("@/lib/quota/saturationSignals"); + storeRateLimitHeaders(connectionId, provider, providerResponse.headers as Record); + } catch { + // fail-open: saturation signal is best-effort + } } catch (error) { trackPendingRequest(model, provider, connectionId, false); if (isSemaphoreCapacityError(error)) { diff --git a/open-sse/handlers/embeddings.ts b/open-sse/handlers/embeddings.ts index cc2eb494c9..cee2d5b930 100644 --- a/open-sse/handlers/embeddings.ts +++ b/open-sse/handlers/embeddings.ts @@ -169,6 +169,27 @@ export async function handleEmbedding({ } try { + // Quota share enforcement (fail-open: errors allow the request through) + if (apiKeyId && connectionId && provider) { + try { + const { enforceQuotaShare } = await import("@/lib/quota/enforce"); + const quotaDecision = await enforceQuotaShare({ + apiKeyId, + connectionId, + provider, + }); + if (quotaDecision.kind === "block") { + return { + success: false, + status: quotaDecision.httpStatus ?? 429, + error: quotaDecision.reason || "Quota share limit reached", + }; + } + } catch { + // fail-open per B16 + } + } + // Log provider request reqLogger.logTargetRequest(providerConfig.baseUrl, headers, upstreamBody); @@ -263,7 +284,24 @@ export async function handleEmbedding({ connectionId, }).catch(() => {}); - // Normalize response to OpenAI format + // Record quota consumption (fire-and-forget, never blocks) + if (apiKeyId && connectionId && provider) { + try { + const { scheduleRecordConsumption } = await import("@/lib/quota/spendRecorder"); + scheduleRecordConsumption({ + apiKeyId, + connectionId, + provider, + cost: { + tokens: data.usage?.prompt_tokens || data.usage?.total_tokens || 0, + requests: 1, + }, + }); + } catch { + // fail-open per B29 + } + } + return { success: true, data: normalizedResponse, diff --git a/src/app/api/quota/pools/[id]/usage/route.ts b/src/app/api/quota/pools/[id]/usage/route.ts index 727746f48a..b3983c2c51 100644 --- a/src/app/api/quota/pools/[id]/usage/route.ts +++ b/src/app/api/quota/pools/[id]/usage/route.ts @@ -2,19 +2,7 @@ * GET /api/quota/pools/[id]/usage — pool consumption snapshot with dimensions * * Resolves the pool's provider plan to get dimensions, then calls - * poolUsageWithDimensions on the concrete store implementation. - * - * Note on poolUsageWithDimensions availability: - * This method is defined on SqliteQuotaStore (and RedisQuotaStore) but is NOT - * part of the QuotaStore interface (keeping the interface minimal). F8 accesses - * it via dynamic type-narrowing: - * - * const storeExt = store as { poolUsageWithDimensions?: (...) => Promise<...> }; - * if (typeof storeExt.poolUsageWithDimensions === "function") { ... } - * else { fallback to store.poolUsage(id) } - * - * This avoids modifying the QuotaStore interface (F6 responsibility) while - * still using the richer method when available. + * poolUsageWithDimensions on the QuotaStore interface. * * Auth: requireManagementAuth * Sanitization: all error responses via buildErrorBody (Hard Rule #12, B25) @@ -51,24 +39,14 @@ export async function GET(request: Request, { params }: RouteParams): Promise - ) => Promise; - }; - - if ( - typeof storeExt.poolUsageWithDimensions === "function" && - plan.dimensions.length > 0 - ) { - snapshot = await storeExt.poolUsageWithDimensions(id, plan.dimensions); + if (plan.dimensions.length > 0) { + snapshot = await store.poolUsageWithDimensions(id, plan.dimensions); } else { - // Fallback: use the interface-standard poolUsage (dimensions come from stored data only) + // Fallback: no plan dimensions configured — return minimal snapshot snapshot = await store.poolUsage(id); } diff --git a/src/lib/db/quotaPools.ts b/src/lib/db/quotaPools.ts index d71ac79813..14b9032e3a 100644 --- a/src/lib/db/quotaPools.ts +++ b/src/lib/db/quotaPools.ts @@ -451,6 +451,14 @@ export function deletePool(id: string): boolean { export function upsertAllocations(poolId: string, allocations: PoolAllocation[]): void { const database = getDb(); + // Normalize: when all weights are 0, distribute equally so the pool is usable + // without requiring a manual re-save. Persists the normalized weights. + const totalWeight = allocations.reduce((s, a) => s + (Number.isFinite(a.weight) ? a.weight : 0), 0); + const normalizedAllocations = + totalWeight === 0 && allocations.length > 0 + ? allocations.map((a) => ({ ...a, weight: 100 / allocations.length })) + : allocations; + // Resolve the target pool's group so we can propagate to siblings. // Defensive: fall back to [poolId] (single-pool semantics) if pool not found. const targetPool = database @@ -475,9 +483,8 @@ export function upsertAllocations(poolId: string, allocations: PoolAllocation[]) VALUES (?, ?, ?, ?, ?, ?)` ); for (const pid of poolIdsInGroup) { - // Replace allocations for this pool. database.prepare("DELETE FROM quota_allocations WHERE pool_id = ?").run(pid); - for (const alloc of allocations) { + for (const alloc of normalizedAllocations) { insert.run( pid, alloc.apiKeyId, diff --git a/src/lib/quota/burnRate.ts b/src/lib/quota/burnRate.ts index 3e6cfd9626..9675e8218d 100644 --- a/src/lib/quota/burnRate.ts +++ b/src/lib/quota/burnRate.ts @@ -24,6 +24,43 @@ export interface BurnRateResult { timeToExhaustionMs: number | null; } +/** + * Compute burn rate from a single snapshot using the sliding window context. + * + * When only one sample is available (the common case for on-demand pool usage + * queries), we derive the rate from the consumption within the current window: + * rate = consumedTotal / elapsedInWindow + * + * This assumes consumption is roughly uniform within the window — a reasonable + * approximation for token/request budgets over hourly/daily/weekly periods. + * + * @param consumedTotal Cumulative consumption in the current sliding window. + * @param windowMs The window duration in milliseconds (e.g. 5h = 18_000_000). + * @param remaining Optional remaining quota (same unit as consumedTotal). + */ +export function computeBurnRateFromWindow( + consumedTotal: number, + windowMs: number, + remaining?: number +): BurnRateResult { + if (consumedTotal <= 0 || windowMs <= 0) { + return { tokensPerSecond: 0, timeToExhaustionMs: null }; + } + + const nowMs = Date.now(); + const currentBucketIndex = Math.floor(nowMs / windowMs); + const windowStartMs = currentBucketIndex * windowMs; + const elapsedMs = Math.max(1, nowMs - windowStartMs); // avoid division by zero + + const safeRate = consumedTotal / (elapsedMs / 1000); // per second + const timeToExhaustionMs = + safeRate > 0 && remaining !== undefined && remaining >= 0 + ? (remaining / safeRate) * 1000 + : null; + + return { tokensPerSecond: safeRate, timeToExhaustionMs }; +} + /** * Compute the current burn rate from a series of samples. * diff --git a/src/lib/quota/enforce.ts b/src/lib/quota/enforce.ts index ffa385be3a..404538e60d 100644 --- a/src/lib/quota/enforce.ts +++ b/src/lib/quota/enforce.ts @@ -187,6 +187,19 @@ export async function enforceQuotaShare(input: EnforceInput): Promise sum + a.weight, 0); - const burnSamples: Array<{ ts: number; consumed: number }> = []; const dimensionSnapshots: PoolUsageSnapshot["dimensions"] = []; for (const planDim of planDimensions) { @@ -289,7 +288,6 @@ export class RedisQuotaStore implements QuotaStore { }); } - burnSamples.push({ ts: nowMs, consumed: consumedTotal }); dimensionSnapshots.push({ unit: planDim.unit as PoolUsageSnapshot["dimensions"][number]["unit"], window: planDim.window as PoolUsageSnapshot["dimensions"][number]["window"], @@ -301,9 +299,10 @@ export class RedisQuotaStore implements QuotaStore { const tokenDim = dimensionSnapshots.find((d) => d.unit === "tokens"); let burnRate: PoolUsageSnapshot["burnRate"]; - if (tokenDim && burnSamples.length >= 1) { + if (tokenDim && tokenDim.consumedTotal > 0) { + const windowMs = WINDOW_MS[tokenDim.window as keyof typeof WINDOW_MS]; const remaining = tokenDim.limit - tokenDim.consumedTotal; - const rateResult = computeBurnRate(burnSamples, remaining); + const rateResult = computeBurnRateFromWindow(tokenDim.consumedTotal, windowMs, remaining); burnRate = { tokensPerSecond: rateResult.tokensPerSecond, timeToExhaustionMs: rateResult.timeToExhaustionMs, diff --git a/src/lib/quota/saturationSignals.ts b/src/lib/quota/saturationSignals.ts index ce11b01e03..c045288707 100644 --- a/src/lib/quota/saturationSignals.ts +++ b/src/lib/quota/saturationSignals.ts @@ -41,6 +41,51 @@ const CACHE_TTL_MS = 30_000; // 30 seconds const _cache = new Map(); +// --------------------------------------------------------------------------- +// Rate-limit header cache (populated by response handlers) +// --------------------------------------------------------------------------- + +interface RateLimitHeaderEntry { + limit: number; + remaining: number; + ts: number; +} + +const _rateLimitHeaders = new Map(); +const RL_HEADER_TTL_MS = 5 * 60 * 1000; // 5 minutes + +/** + * Store rate-limit headers from an upstream response for saturation signal use. + * Called by the response handler after a successful request. + */ +export function storeRateLimitHeaders( + connectionId: string, + provider: string, + headers: Record +): void { + // Anthropic: anthropic-ratelimit-requests-limit / anthropic-ratelimit-requests-remaining + const limitStr = + headers["anthropic-ratelimit-requests-limit"] ?? + headers["x-ratelimit-limit-requests"] ?? + headers["x-ratelimit-limit"]; + const remainingStr = + headers["anthropic-ratelimit-requests-remaining"] ?? + headers["x-ratelimit-remaining-requests"] ?? + headers["x-ratelimit-remaining"]; + + if (limitStr && remainingStr) { + const limit = Number(limitStr); + const remaining = Number(remainingStr); + if (Number.isFinite(limit) && limit > 0 && Number.isFinite(remaining)) { + _rateLimitHeaders.set(`${provider}:${connectionId}`, { + limit, + remaining, + ts: Date.now(), + }); + } + } +} + function cacheKey(connectionId: string, provider: string, dim: DimensionSpec): string { return `${provider}:${connectionId}:${dim.unit}:${dim.window}`; } @@ -95,40 +140,55 @@ async function fetchBailianSaturation( const quota = await mod.fetchBailianQuota(connectionId); if (!quota) return 0; - // Select the window matching the dimension + const q = quota as unknown as Record; let pct = 0; switch (dim.window) { case "5h": - pct = quota.window5h?.percentUsed ?? 0; + pct = (q.window5h as Record)?.percentUsed as number ?? 0; break; case "weekly": - pct = quota.windowWeekly?.percentUsed ?? 0; + pct = (q.windowWeekly as Record)?.percentUsed as number ?? 0; break; case "monthly": - pct = quota.windowMonthly?.percentUsed ?? 0; + pct = (q.windowMonthly as Record)?.percentUsed as number ?? 0; break; default: - pct = quota.percentUsed ?? 0; + pct = (q.percentUsed as number) ?? 0; } return Math.min(1, Math.max(0, pct)); } +async function fetchAnthropicSaturation( + connectionId: string, + dim: DimensionSpec +): Promise { + const entry = _rateLimitHeaders.get(`anthropic:${connectionId}`); + if (!entry || Date.now() - entry.ts > RL_HEADER_TTL_MS) return 0; + + const used = entry.limit - entry.remaining; + return Math.min(1, Math.max(0, used / entry.limit)); +} + async function fetchGenericSaturation( connectionId: string, provider: string ): Promise { - const mod = await import("@omniroute/open-sse/services/usage"); - // getUsageForProvider returns an object with percentUsed or similar - const result = await mod.getUsageForProvider(provider, connectionId); - if (!result || typeof result !== "object") return 0; - const obj = result as Record; - const pct = - typeof obj.percentUsed === "number" - ? obj.percentUsed - : typeof obj.used_percent === "number" - ? obj.used_percent - : 0; - return Math.min(1, Math.max(0, pct)); + try { + const mod = await import("@omniroute/open-sse/services/usage"); + const conn = { id: connectionId, provider } as Parameters[0]; + const result = await mod.getUsageForProvider(conn); + if (!result || typeof result !== "object") return 0; + const obj = result as Record; + const pct = + typeof obj.percentUsed === "number" + ? obj.percentUsed + : typeof obj.used_percent === "number" + ? obj.used_percent + : 0; + return Math.min(1, Math.max(0, pct)); + } catch { + return 0; + } } // --------------------------------------------------------------------------- @@ -163,6 +223,10 @@ export async function getSaturation( case "bailian": value = await fetchBailianSaturation(connectionId, dim); break; + case "anthropic": + case "claude": + value = await fetchAnthropicSaturation(connectionId, dim); + break; default: value = await fetchGenericSaturation(connectionId, provider); break; diff --git a/src/lib/quota/sqliteQuotaStore.ts b/src/lib/quota/sqliteQuotaStore.ts index a72526294d..51417df3bd 100644 --- a/src/lib/quota/sqliteQuotaStore.ts +++ b/src/lib/quota/sqliteQuotaStore.ts @@ -26,7 +26,7 @@ import { import { WINDOW_MS, dimensionKeyToString } from "./dimensions"; import type { DimensionKey } from "./dimensions"; import type { QuotaStore, PoolUsageSnapshot } from "./types"; -import { computeBurnRate } from "./burnRate"; +import { computeBurnRateFromWindow } from "./burnRate"; // --------------------------------------------------------------------------- // In-memory mutex (anti-thundering-herd, same pattern as auth.ts) @@ -261,9 +261,6 @@ export class SqliteQuotaStore implements QuotaStore { const { allocations } = pool; const totalWeight = allocations.reduce((sum, a) => sum + a.weight, 0); - // Burn rate samples: collect peek values at nowMs and nowMs - 60s - const burnSamples: Array<{ ts: number; consumed: number }> = []; - const dimensionSnapshots: PoolUsageSnapshot["dimensions"] = []; for (const planDim of planDimensions) { @@ -285,7 +282,6 @@ export class SqliteQuotaStore implements QuotaStore { const effectiveWeight = totalWeight > 0 ? alloc.weight : 0; const fairShare = (effectiveWeight / 100) * planDim.limit; const deficit = consumed - fairShare; - // borrowing = key consumed more than its fair share const borrowing = consumed > fairShare; perKey.push({ @@ -297,8 +293,6 @@ export class SqliteQuotaStore implements QuotaStore { }); } - burnSamples.push({ ts: nowMs, consumed: consumedTotal }); - dimensionSnapshots.push({ unit: planDim.unit as PoolUsageSnapshot["dimensions"][number]["unit"], window: planDim.window as PoolUsageSnapshot["dimensions"][number]["window"], @@ -308,12 +302,13 @@ export class SqliteQuotaStore implements QuotaStore { }); } - // Compute burn rate from token-like dimensions + // Burn rate: derive from the sliding window (single-snapshot, no history needed). const tokenDim = dimensionSnapshots.find((d) => d.unit === "tokens"); let burnRate: PoolUsageSnapshot["burnRate"]; - if (tokenDim && burnSamples.length >= 1) { + if (tokenDim && tokenDim.consumedTotal > 0) { + const windowMs = WINDOW_MS[tokenDim.window as keyof typeof WINDOW_MS]; const remaining = tokenDim.limit - tokenDim.consumedTotal; - const rateResult = computeBurnRate(burnSamples, remaining); + const rateResult = computeBurnRateFromWindow(tokenDim.consumedTotal, windowMs, remaining); burnRate = { tokensPerSecond: rateResult.tokensPerSecond, timeToExhaustionMs: rateResult.timeToExhaustionMs, diff --git a/src/lib/quota/types.ts b/src/lib/quota/types.ts index 63b19c2b39..98c119e2ee 100644 --- a/src/lib/quota/types.ts +++ b/src/lib/quota/types.ts @@ -45,6 +45,20 @@ export interface QuotaStore { */ poolConsumedTotal(poolId: string, dim: DimensionKey): Promise; poolUsage(poolId: string): Promise; + /** + * Build a PoolUsageSnapshot with explicit plan dimensions. This is the + * primary method for dashboard / REST usage — it resolves per-key + * consumption, fair-share, deficit, borrowing, and burn-rate from the + * plan's dimension list. + * + * The parameterless `poolUsage()` is kept for backward compatibility but + * returns minimal data (no plan context). Prefer this method when plan + * dimensions are available. + */ + poolUsageWithDimensions( + poolId: string, + planDimensions: Array<{ unit: string; window: string; limit: number }> + ): Promise; clear(apiKeyId: string, dim: DimensionKey): Promise; } diff --git a/tests/unit/quota-sharing-fixes.test.ts b/tests/unit/quota-sharing-fixes.test.ts new file mode 100644 index 0000000000..83743fea83 --- /dev/null +++ b/tests/unit/quota-sharing-fixes.test.ts @@ -0,0 +1,260 @@ +/** + * tests/unit/quota-sharing-fixes.test.ts + * + * Verifies the quota sharing improvements from fix/quota-sharing-improvements: + * 1. computeBurnRateFromWindow() produces correct non-zero output + * 2. storeRateLimitHeaders() + fetchAnthropicSaturation() round-trip + * 3. upsertAllocations() normalizes zero weights to equal distribution + * 4. poolUsageWithDimensions() returns real burn rate data + * 5. QuotaStore interface includes poolUsageWithDimensions() + */ + +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-fixes-")); +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"); + +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 }); + } +}); + +// ─── Fix 2: computeBurnRateFromWindow ───────────────────────────────────────── + +test("computeBurnRateFromWindow: returns non-zero rate for non-zero consumption", async () => { + const { computeBurnRateFromWindow } = await import("../../src/lib/quota/burnRate.ts"); + + // Simulate 500 tokens consumed in a 1-hour window + const result = computeBurnRateFromWindow(500, 60 * 60 * 1000, 500); + + assert.ok(result.tokensPerSecond > 0, `rate should be >0, got ${result.tokensPerSecond}`); + assert.ok( + result.timeToExhaustionMs !== null && result.timeToExhaustionMs > 0, + `timeToExhaustion should be >0, got ${result.timeToExhaustionMs}` + ); +}); + +test("computeBurnRateFromWindow: returns zero for zero consumption", async () => { + const { computeBurnRateFromWindow } = await import("../../src/lib/quota/burnRate.ts"); + + const result = computeBurnRateFromWindow(0, 60 * 60 * 1000); + assert.equal(result.tokensPerSecond, 0); + assert.equal(result.timeToExhaustionMs, null); +}); + +test("computeBurnRateFromWindow: rate is consumption / elapsed (not / full window)", async () => { + const { computeBurnRateFromWindow } = await import("../../src/lib/quota/burnRate.ts"); + + // 1000 tokens consumed. The rate should be based on elapsed time within the + // current bucket, not the full window duration. We can't control time, but + // we can verify the rate is at least as fast as consumed/windowMs (lower bound). + const windowMs = 60 * 60 * 1000; // 1h + const consumed = 1000; + const result = computeBurnRateFromWindow(consumed, windowMs, 1000); + + // Lower bound: if we're at the very end of the window, rate = consumed/window + const lowerBound = consumed / (windowMs / 1000); + assert.ok( + result.tokensPerSecond >= lowerBound * 0.9, // allow 10% margin + `rate ${result.tokensPerSecond} should be >= lower bound ${lowerBound * 0.9}` + ); +}); + +// ─── Fix 3: Weight normalization ────────────────────────────────────────────── + +test("upsertAllocations: normalizes zero weights to equal distribution", async () => { + const pool = poolsDb.createPool({ + connectionId: "conn-weight-test", + name: "Weight Test Pool", + allocations: [ + { apiKeyId: "key-x", weight: 0, policy: "hard" }, + { apiKeyId: "key-y", weight: 0, policy: "hard" }, + { apiKeyId: "key-z", weight: 0, policy: "hard" }, + ], + }); + + // Call upsertAllocations with all-zero weights + poolsDb.upsertAllocations(pool.id, [ + { apiKeyId: "key-x", weight: 0, policy: "hard" }, + { apiKeyId: "key-y", weight: 0, policy: "hard" }, + { apiKeyId: "key-z", weight: 0, policy: "hard" }, + ]); + + // Read back — weights should be normalized to ~33.33 each + const updated = poolsDb.getPool(pool.id); + assert.ok(updated, "pool should exist"); + assert.equal(updated!.allocations.length, 3); + + for (const alloc of updated!.allocations) { + assert.ok( + Math.abs(alloc.weight - 100 / 3) < 0.01, + `weight should be ~33.33, got ${alloc.weight}` + ); + } +}); + +test("upsertAllocations: preserves non-zero weights", async () => { + const pool = poolsDb.createPool({ + connectionId: "conn-weight-preserve", + name: "Weight Preserve Pool", + }); + + poolsDb.upsertAllocations(pool.id, [ + { apiKeyId: "key-a", weight: 70, policy: "hard" }, + { apiKeyId: "key-b", weight: 30, policy: "soft" }, + ]); + + const updated = poolsDb.getPool(pool.id); + assert.ok(updated); + assert.equal(updated!.allocations[0].weight, 70); + assert.equal(updated!.allocations[1].weight, 30); +}); + +// ─── Fix 4: storeRateLimitHeaders + Anthropic saturation ───────────────────── + +test("storeRateLimitHeaders: stores headers and getSaturation reads them", async () => { + const { storeRateLimitHeaders, _clearSaturationCache } = await import( + "../../src/lib/quota/saturationSignals.ts" + ); + + _clearSaturationCache(); + + // Store headers as if Anthropic returned them (800/1000 remaining = 20% used) + storeRateLimitHeaders("conn-anthropic-1", "anthropic", { + "anthropic-ratelimit-requests-limit": "1000", + "anthropic-ratelimit-requests-remaining": "800", + }); + + // Now read via getSaturation — should return ~0.2 (20% used) + const { getSaturation } = await import("../../src/lib/quota/saturationSignals.ts"); + const value = await getSaturation("conn-anthropic-1", "anthropic", { + unit: "requests", + window: "hourly", + }); + + assert.ok(value > 0.15, `saturation should be >0.15, got ${value}`); + assert.ok(value < 0.25, `saturation should be <0.25, got ${value}`); +}); + +test("storeRateLimitHeaders: ignores non-Anthropic headers gracefully", async () => { + const { storeRateLimitHeaders, _clearSaturationCache, getSaturation } = await import( + "../../src/lib/quota/saturationSignals.ts" + ); + + _clearSaturationCache(); + + // Store headers with no rate-limit info + storeRateLimitHeaders("conn-other", "openai", { + "content-type": "application/json", + }); + + const value = await getSaturation("conn-other", "openai", { + unit: "tokens", + window: "hourly", + }); + + // Should return 0 (no data → generous mode) + assert.equal(value, 0); +}); + +// ─── Fix 1: poolUsageWithDimensions returns burn rate ───────────────────────── + +test("poolUsageWithDimensions: returns non-null burn rate for token dimensions", async () => { + const { SqliteQuotaStore } = await import("../../src/lib/quota/sqliteQuotaStore.ts"); + const store = new SqliteQuotaStore(); + + const pool = poolsDb.createPool({ + connectionId: "conn-burn-rate", + name: "Burn Rate Pool", + allocations: [ + { apiKeyId: "key-br-1", weight: 100, policy: "hard" }, + ], + }); + + const dim = { poolId: pool.id, unit: "tokens" as const, window: "hourly" as const }; + await store.consume("key-br-1", dim, 5000); + + const snapshot = await store.poolUsageWithDimensions(pool.id, [ + { unit: "tokens", window: "hourly", limit: 50000 }, + ]); + + assert.ok(snapshot.burnRate, "burnRate should be present"); + assert.ok( + snapshot.burnRate!.tokensPerSecond > 0, + `tokensPerSecond should be >0, got ${snapshot.burnRate!.tokensPerSecond}` + ); + assert.ok( + snapshot.burnRate!.timeToExhaustionMs !== null && snapshot.burnRate!.timeToExhaustionMs > 0, + `timeToExhaustionMs should be >0, got ${snapshot.burnRate!.timeToExhaustionMs}` + ); +}); + +test("poolUsageWithDimensions: no burn rate when consumedTotal is 0", async () => { + const { SqliteQuotaStore } = await import("../../src/lib/quota/sqliteQuotaStore.ts"); + const store = new SqliteQuotaStore(); + + const pool = poolsDb.createPool({ + connectionId: "conn-no-burn", + name: "No Burn Pool", + allocations: [ + { apiKeyId: "key-nb-1", weight: 100, policy: "hard" }, + ], + }); + + const snapshot = await store.poolUsageWithDimensions(pool.id, [ + { unit: "tokens", window: "hourly", limit: 50000 }, + ]); + + assert.equal(snapshot.burnRate, undefined, "burnRate should be undefined when nothing consumed"); +}); + +// ─── Fix 1: QuotaStore interface includes poolUsageWithDimensions ──────────── + +test("QuotaStore interface: poolUsageWithDimensions is on the interface", async () => { + // Type-level check: if this compiles, the method is on the interface. + // We verify at runtime that both implementations have it. + const { SqliteQuotaStore } = await import("../../src/lib/quota/sqliteQuotaStore.ts"); + const sqlite = new SqliteQuotaStore(); + assert.equal(typeof sqlite.poolUsageWithDimensions, "function", "SqliteQuotaStore must have poolUsageWithDimensions"); + + // Redis store (just check the prototype) + const { RedisQuotaStore } = await import("../../src/lib/quota/redisQuotaStore.ts"); + assert.equal( + typeof RedisQuotaStore.prototype.poolUsageWithDimensions, + "function", + "RedisQuotaStore must have poolUsageWithDimensions" + ); +});