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-token-limits-")); process.env.DATA_DIR = TEST_DATA_DIR; const core = await import("../../src/lib/db/core.ts"); const tokenLimits = await import("../../src/lib/db/tokenLimits.ts"); const counter = await import("../../open-sse/services/tokenLimitCounter.ts"); const flush = () => new Promise((r) => setImmediate(r)); // Window pivots (UTC). const NOW_JAN = Date.UTC(2026, 0, 15, 12, 0, 0); const NOW_FEB = Date.UTC(2026, 1, 15, 12, 0, 0); const D1 = Date.UTC(2026, 0, 10, 12); const D2 = Date.UTC(2026, 0, 11, 12); const W1 = Date.UTC(2026, 0, 6, 12); // Tue const W2 = Date.UTC(2026, 0, 13, 12); // next Tue const M1 = Date.UTC(2026, 0, 8, 6); // same daily window as D1? no -> Jan 8; use for "same window" daily checks separately 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 (error: any) { if ((error?.code === "EBUSY" || error?.code === "EPERM") && attempt < 9) { await new Promise((resolve) => setTimeout(resolve, 50 * (attempt + 1))); } else { throw error; } } } fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); } function insertUsage( apiKeyId: string, provider: string, model: string, tokensInput: number, tokensOutput: number, ts: string, extra: { cacheRead?: number; cacheCreation?: number; reasoning?: number } = {} ) { const db = core.getDbInstance(); db.prepare( `INSERT INTO usage_history (provider, model, api_key_id, tokens_input, tokens_output, tokens_cache_read, tokens_cache_creation, tokens_reasoning, success, timestamp) VALUES (?, ?, ?, ?, ?, ?, ?, ?, 1, ?)` ).run( provider, model, apiKeyId, tokensInput, tokensOutput, extra.cacheRead ?? 0, extra.cacheCreation ?? 0, extra.reasoning ?? 0, ts ); } test.beforeEach(async () => { await resetStorage(); counter.clearTokenLimitCache(); }); test.after(async () => { core.resetDbInstance(); fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); }); test("window rollover: daily/weekly/monthly produce distinct windowStart", async () => { const daily = tokenLimits.upsertTokenLimit({ apiKeyId: "k-daily", scopeType: "global", tokenLimit: 1000, resetInterval: "daily", }); const weekly = tokenLimits.upsertTokenLimit({ apiKeyId: "k-week", scopeType: "global", tokenLimit: 1000, resetInterval: "weekly", }); const monthly = tokenLimits.upsertTokenLimit({ apiKeyId: "k-month", scopeType: "global", tokenLimit: 1000, resetInterval: "monthly", }); // Distinct windows across a boundary. assert.notEqual( tokenLimits.resetWindowIfElapsed(daily, D1).windowStart, tokenLimits.resetWindowIfElapsed(daily, D2).windowStart ); assert.notEqual( tokenLimits.resetWindowIfElapsed(weekly, W1).windowStart, tokenLimits.resetWindowIfElapsed(weekly, W2).windowStart ); assert.notEqual( tokenLimits.resetWindowIfElapsed(monthly, NOW_JAN).windowStart, tokenLimits.resetWindowIfElapsed(monthly, NOW_FEB).windowStart ); // Same window: two distinct `now`s inside the same period give same windowStart. const dailySameA = Date.UTC(2026, 0, 10, 1); const dailySameB = Date.UTC(2026, 0, 10, 23); assert.equal( tokenLimits.resetWindowIfElapsed(daily, dailySameA).windowStart, tokenLimits.resetWindowIfElapsed(daily, dailySameB).windowStart ); // weekly same window: Mon and Sun of same week. const weekMon = Date.UTC(2026, 0, 5, 3); // Monday const weekSun = Date.UTC(2026, 0, 11, 20); // Sunday same week assert.equal( tokenLimits.resetWindowIfElapsed(weekly, weekMon).windowStart, tokenLimits.resetWindowIfElapsed(weekly, weekSun).windowStart ); // monthly same window: two days in January. assert.equal( tokenLimits.resetWindowIfElapsed(monthly, NOW_JAN).windowStart, tokenLimits.resetWindowIfElapsed(monthly, Date.UTC(2026, 0, 28, 4)).windowStart ); }); test("seed-on-miss equals usage_history SUM for the active window", async () => { const limit = tokenLimits.upsertTokenLimit({ apiKeyId: "k2", scopeType: "model", scopeValue: "gpt-4o", tokenLimit: 100000, resetInterval: "monthly", }); // In-window, same-model rows (counted). insertUsage("k2", "openai", "gpt-4o", 100, 50, new Date(Date.UTC(2026, 0, 12)).toISOString()); insertUsage("k2", "openai", "gpt-4o", 30, 20, new Date(Date.UTC(2026, 0, 14)).toISOString()); // Different month (excluded). insertUsage("k2", "openai", "gpt-4o", 999, 999, new Date(Date.UTC(2025, 11, 31)).toISOString()); // Different model (excluded). insertUsage("k2", "openai", "gpt-4o-mini", 777, 777, new Date(Date.UTC(2026, 0, 13)).toISOString()); const expected = 100 + 50 + 30 + 20; assert.equal(counter.seedWindowUsageFromHistory(limit, NOW_JAN), expected); }); test("getCurrentWindowUsage seeds from history and PERSISTS the seed (FIX 3)", async () => { const limit = tokenLimits.upsertTokenLimit({ apiKeyId: "k2b", scopeType: "model", scopeValue: "gpt-4o", tokenLimit: 100000, resetInterval: "monthly", }); insertUsage("k2b", "openai", "gpt-4o", 200, 100, new Date(Date.UTC(2026, 0, 12)).toISOString()); const seeded = 300; assert.equal(counter.seedWindowUsageFromHistory(limit, NOW_JAN), seeded); // FIX 3: a force-fresh read on a cold window now PERSISTS the seed to the // counter row (previously the seed was read-only and DB usage stayed 0). assert.equal(counter.getCurrentWindowUsage(limit, NOW_JAN, true), seeded); assert.equal(tokenLimits.getWindowUsage(limit, NOW_JAN), seeded); // A subsequent increment accumulates ON TOP of the persisted seed — the prior // historical usage is NOT forgotten. const { windowStart } = tokenLimits.resetWindowIfElapsed(limit, NOW_JAN); tokenLimits.incrementWindowTokens(limit.id, windowStart, 25); assert.equal(tokenLimits.getWindowUsage(limit, NOW_JAN), seeded + 25); assert.equal(counter.getCurrentWindowUsage(limit, NOW_JAN, true), seeded + 25); }); test("seed total excludes cache tokens (no double-count) (FIX 2)", async () => { const limit = tokenLimits.upsertTokenLimit({ apiKeyId: "k2c", scopeType: "model", scopeValue: "claude-sonnet", tokenLimit: 100000, resetInterval: "monthly", }); // tokens_input ALREADY INCLUDES cache_read + cache_creation (these columns are a // breakdown, per migration 012). Billable = input + output + reasoning ONLY. insertUsage("k2c", "anthropic", "claude-sonnet", 500, 200, new Date(Date.UTC(2026, 0, 12)).toISOString(), { cacheRead: 300, cacheCreation: 100, reasoning: 40, }); // 500 + 200 + 40 = 740. Must NOT add cacheRead/cacheCreation again (would be 1140). assert.equal(counter.seedWindowUsageFromHistory(limit, NOW_JAN), 740); }); test("cold-window recordTokenUsage seeds from history before increment (FIX 4)", async () => { // recordTokenUsage uses Date.now() internally; insert history in the current window. const limit = tokenLimits.upsertTokenLimit({ apiKeyId: "k2d", scopeType: "model", scopeValue: "gpt-4o", tokenLimit: 1000000, resetInterval: "monthly", }); // Prior historical usage in this window, but NO counter row yet (cold window). insertUsage("k2d", "openai", "gpt-4o", 400, 100, new Date().toISOString()); assert.equal(tokenLimits.getWindowUsage(limit, Date.now()), 0); // no counter row yet // First record on the cold window must seed (500) then add the delta (50) = 550, // NOT restart from 0 (which would yield 50). counter.recordTokenUsage("k2d", "openai", "gpt-4o", 50); await flush(); await flush(); assert.equal(tokenLimits.getWindowUsage(limit, Date.now()), 550); }); test("most-restrictive breach wins when model and provider both match", async () => { // Case A: both breach; provider has smaller limitValue (tie on remaining=0). const modelLimit = tokenLimits.upsertTokenLimit({ apiKeyId: "k3", scopeType: "model", scopeValue: "gpt-4o", tokenLimit: 100, resetInterval: "monthly", }); const providerLimit = tokenLimits.upsertTokenLimit({ apiKeyId: "k3", scopeType: "provider", scopeValue: "openai", tokenLimit: 50, resetInterval: "monthly", }); const mWs = tokenLimits.resetWindowIfElapsed(modelLimit, NOW_JAN).windowStart; const pWs = tokenLimits.resetWindowIfElapsed(providerLimit, NOW_JAN).windowStart; tokenLimits.incrementWindowTokens(modelLimit.id, mWs, 100); // remaining 0 tokenLimits.incrementWindowTokens(providerLimit.id, pWs, 55); // remaining 0, smaller limitValue const breachA = counter.checkTokenLimits("k3", "openai", "gpt-4o", NOW_JAN); assert.ok(breachA); assert.equal(breachA!.scopeType, "provider"); assert.equal(breachA!.limitValue, 50); // Case B: only the model limit breaches → it is returned. const modelLimit2 = tokenLimits.upsertTokenLimit({ apiKeyId: "k3b", scopeType: "model", scopeValue: "gpt-4o", tokenLimit: 100, resetInterval: "monthly", }); const providerLimit2 = tokenLimits.upsertTokenLimit({ apiKeyId: "k3b", scopeType: "provider", scopeValue: "openai", tokenLimit: 200, resetInterval: "monthly", }); const mWs2 = tokenLimits.resetWindowIfElapsed(modelLimit2, NOW_JAN).windowStart; const pWs2 = tokenLimits.resetWindowIfElapsed(providerLimit2, NOW_JAN).windowStart; tokenLimits.incrementWindowTokens(modelLimit2.id, mWs2, 100); // breach (>=100) tokenLimits.incrementWindowTokens(providerLimit2.id, pWs2, 150); // 150 < 200 → no breach const breachB = counter.checkTokenLimits("k3b", "openai", "gpt-4o", NOW_JAN); assert.ok(breachB); assert.equal(breachB!.scopeType, "model"); assert.equal(breachB!.limitValue, 100); }); test("disabled limit is ignored by checkTokenLimits", async () => { const limit = tokenLimits.upsertTokenLimit({ apiKeyId: "k4", scopeType: "model", scopeValue: "gpt-4o", tokenLimit: 10, resetInterval: "monthly", enabled: false, }); const ws = tokenLimits.resetWindowIfElapsed(limit, NOW_JAN).windowStart; tokenLimits.incrementWindowTokens(limit.id, ws, 999); assert.equal(counter.checkTokenLimits("k4", "openai", "gpt-4o", NOW_JAN), null); }); test("global fallback applies when no model/provider limit", async () => { const limit = tokenLimits.upsertTokenLimit({ apiKeyId: "k5", scopeType: "global", tokenLimit: 10, resetInterval: "monthly", }); const ws = tokenLimits.resetWindowIfElapsed(limit, NOW_JAN).windowStart; tokenLimits.incrementWindowTokens(limit.id, ws, 20); const breach = counter.checkTokenLimits("k5", "openai", "gpt-4o", NOW_JAN); assert.ok(breach); assert.equal(breach!.scopeType, "global"); assert.equal(breach!.limitValue, 10); }); test("atomic increment under repeated calls has no lost updates", async () => { const limit = tokenLimits.upsertTokenLimit({ apiKeyId: "k6", scopeType: "global", tokenLimit: 1000000, resetInterval: "monthly", }); const { windowStart } = tokenLimits.resetWindowIfElapsed(limit, NOW_JAN); for (let i = 0; i < 100; i++) { tokenLimits.incrementWindowTokens(limit.id, windowStart, 7); } assert.equal(tokenLimits.getWindowUsage(limit, NOW_JAN), 700); }); test("getCurrentWindowUsage cache hit / miss / forceFresh", async () => { const limit = tokenLimits.upsertTokenLimit({ apiKeyId: "k7", scopeType: "global", tokenLimit: 1000, resetInterval: "monthly", }); // Prime cache via write-through. counter.addWindowTokens(limit, 50, NOW_JAN); // cache + DB = 50 const { windowStart: ws } = tokenLimits.resetWindowIfElapsed(limit, NOW_JAN); // Mutate DB directly behind the cache → DB=80, cache stale at 50. tokenLimits.incrementWindowTokens(limit.id, ws, 30); // Cache HIT (within TTL) returns stale 50. assert.equal(counter.getCurrentWindowUsage(limit, NOW_JAN, false), 50); // forceFresh returns authoritative 80 and refreshes cache. assert.equal(counter.getCurrentWindowUsage(limit, NOW_JAN, true), 80); // Subsequent normal read returns refreshed 80. assert.equal(counter.getCurrentWindowUsage(limit, NOW_JAN, false), 80); }); test("recordTokenUsage is fire-and-forget and records after microtask flush", async () => { // recordTokenUsage uses Date.now() internally — create + read with default now. const modelLimit = tokenLimits.upsertTokenLimit({ apiKeyId: "k8", scopeType: "model", scopeValue: "gpt-4o", tokenLimit: 100000, resetInterval: "monthly", }); // Returns synchronously (void) and does not throw before the microtask runs. assert.equal(counter.recordTokenUsage("k8", "openai", "gpt-4o", 42), undefined); await flush(); await flush(); assert.ok(tokenLimits.getWindowUsage(modelLimit, Date.now()) >= 42); // No-op cases: tokens <= 0 and empty apiKey. const before = tokenLimits.getWindowUsage(modelLimit, Date.now()); counter.recordTokenUsage("k8", "openai", "gpt-4o", 0); counter.recordTokenUsage("k8", "openai", "gpt-4o", -5); counter.recordTokenUsage("", "openai", "gpt-4o", 99); await flush(); await flush(); assert.equal(tokenLimits.getWindowUsage(modelLimit, Date.now()), before); });