From 1461a78d0755cca347af1a4429c8fcaa6b53abd7 Mon Sep 17 00:00:00 2001 From: Muhammad Mugni Hadi Date: Fri, 29 May 2026 15:29:10 +0700 Subject: [PATCH] feat(usage): per-API-key token limits scoped to model/provider/global (#2888) Integrated into release/v3.8.6 --- docs/reference/API_REFERENCE.md | 40 +- open-sse/handlers/chatCore.ts | 68 +++- open-sse/services/combo.ts | 51 ++- open-sse/services/tokenLimitCounter.ts | 342 ++++++++++++++++ src/app/api/usage/token-limits/route.ts | 98 +++++ .../migrations/073_per_model_token_limits.sql | 61 +++ src/lib/db/tokenLimits.ts | 295 ++++++++++++++ src/lib/localDb.ts | 19 + src/shared/utils/apiKeyPolicy.ts | 34 ++ src/shared/validation/schemas.ts | 28 ++ tests/unit/token-limits.test.ts | 377 ++++++++++++++++++ 11 files changed, 1400 insertions(+), 13 deletions(-) create mode 100644 open-sse/services/tokenLimitCounter.ts create mode 100644 src/app/api/usage/token-limits/route.ts create mode 100644 src/lib/db/migrations/073_per_model_token_limits.sql create mode 100644 src/lib/db/tokenLimits.ts create mode 100644 tests/unit/token-limits.test.ts diff --git a/docs/reference/API_REFERENCE.md b/docs/reference/API_REFERENCE.md index b888832345..34803a7fe5 100644 --- a/docs/reference/API_REFERENCE.md +++ b/docs/reference/API_REFERENCE.md @@ -325,12 +325,13 @@ Response example: ### Usage & Analytics -| Endpoint | Method | Description | -| --------------------------- | ------ | -------------------- | -| `/api/usage/history` | GET | Usage history | -| `/api/usage/logs` | GET | Usage logs | -| `/api/usage/request-logs` | GET | Request-level logs | -| `/api/usage/[connectionId]` | GET | Per-connection usage | +| Endpoint | Method | Description | +| --------------------------- | ----------------- | --------------------------------- | +| `/api/usage/history` | GET | Usage history | +| `/api/usage/logs` | GET | Usage logs | +| `/api/usage/request-logs` | GET | Request-level logs | +| `/api/usage/[connectionId]` | GET | Per-connection usage | +| `/api/usage/token-limits` | GET/POST/DELETE | Per-API-key token-limit budgets | ### Settings @@ -590,6 +591,33 @@ Content-Type: application/json > **Schema notes** (`setBudgetSchema`): `apiKeyId` is required; at least one of `dailyLimitUsd`, `weeklyLimitUsd`, or `monthlyLimitUsd` must be greater than zero. Optional fields: `warningThreshold` (0–1), `resetInterval` (`daily` | `weekly` | `monthly`), `resetTime` (`HH:MM`). The legacy `{keyId, limit, period}` shape returns `400 Bad Request`. +## Token Limits + +Per-API-key **token** budgets (distinct from the USD-based Budget above). Enforced inline on the request path: when a key's current window usage reaches its limit, requests are rejected with `429 Too Many Requests`. Limits can be scoped to a specific `model`, a `provider`, or applied `global`ly across the key; when several limits match a request, the most restrictive one wins. + +```bash +# List a key's token limits (includes live window usage) +GET /api/usage/token-limits?apiKeyId=key-123 + +# Create or update a token limit +POST /api/usage/token-limits +Content-Type: application/json + +{ + "apiKeyId": "key-123", + "scopeType": "model", + "scopeValue": "openai/gpt-4o", + "tokenLimit": 1000000, + "resetInterval": "monthly", + "enabled": true +} + +# Delete a token limit by id +DELETE /api/usage/token-limits?id=tl-abc +``` + +> **Schema notes** (`setTokenLimitSchema`): `apiKeyId` and `scopeType` (`model` | `provider` | `global`) are required. `scopeValue` is required unless `scopeType` is `global` (e.g. a model id for `model` scope, a provider id for `provider` scope). `tokenLimit` must be a positive integer (coerced from string). Optional: `id` (omit to create, supply to update), `resetInterval` (`daily` | `weekly` | `monthly`, default `monthly`), `resetTime` (`HH:MM`), `enabled` (default `true`). `GET` responses enrich each limit with `tokensUsed`, `remaining`, `windowStart`, `periodStartAt`, and `nextResetAt`. This is a management-class endpoint (auth enforced centrally by the authz pipeline). + ## Request Processing 1. Client sends request to `/v1/*` diff --git a/open-sse/handlers/chatCore.ts b/open-sse/handlers/chatCore.ts index 58996ee722..99107250f4 100644 --- a/open-sse/handlers/chatCore.ts +++ b/open-sse/handlers/chatCore.ts @@ -38,6 +38,10 @@ import { formatProviderError, sanitizeErrorMessage, } from "../utils/error.ts"; +import { + checkTokenLimits, + recordTokenUsage, +} from "@omniroute/open-sse/services/tokenLimitCounter.ts"; import { COOLDOWN_MS, HTTP_STATUS, @@ -84,7 +88,12 @@ import { appendRequestLog, saveCallLog, } from "@/lib/usageDb"; -import { formatUsageLog } from "@/lib/usage/tokenAccounting"; +import { + formatUsageLog, + getLoggedInputTokens, + getLoggedOutputTokens, + getReasoningTokens, +} from "@/lib/usage/tokenAccounting"; import { recordCost } from "@/domain/costRules"; import { calculateCost } from "@/lib/usage/costCalculator"; import { buildOmniRouteResponseMetaHeaders } from "@/domain/omnirouteResponseMeta"; @@ -827,6 +836,16 @@ function createAbortError(signal: AbortSignal): Error { return err; } +/** Billable token total — mirrors the columns persisted by saveRequestUsage so the + * live token-limit counter stays consistent with usage_history seed-on-miss. */ +function computeBillableTokens(usage: unknown): number { + // Cache read/creation tokens are a BREAKDOWN already contained inside + // getLoggedInputTokens (prompt_tokens / input_tokens). Adding them here would + // double-count. Canonical billable total = input + output + reasoning, matching + // the columns persisted by saveRequestUsage and seedWindowUsageFromHistory. + return getLoggedInputTokens(usage) + getLoggedOutputTokens(usage) + getReasoningTokens(usage); +} + function getExecutorTimeoutMs(executor: unknown): number { const getTimeoutMs = (executor as { getTimeoutMs?: () => unknown } | null)?.getTimeoutMs; if (typeof getTimeoutMs !== "function") return FETCH_TIMEOUT_MS; @@ -3887,6 +3906,35 @@ export async function handleChatCore({ 0; log?.debug?.("REQUEST", `${provider.toUpperCase()} | ${model} | ${msgCount} msgs`); + // ── Tier 2: Authoritative per-model/provider token-limit check (provider now resolved) ── + if (apiKeyInfo?.id) { + try { + const tokenBreach = checkTokenLimits(apiKeyInfo.id, provider || undefined, model || undefined); + if (tokenBreach) { + const scopeLabel = + tokenBreach.scopeType === "global" + ? "account" + : `${tokenBreach.scopeType} "${tokenBreach.scopeValue}"`; + // FIX 6: clear the pending request marker before the early return so we do + // not leak a phantom pending request (start was tracked at line ~1847). + trackPendingRequest(model, provider, connectionId, false); + // FIX 5: tag this as a per-API-key token-limit breach (errorCode + // TOKEN_LIMIT_EXCEEDED) so the combo loop can distinguish it from an + // upstream 429 and NOT cool shared accounts / retry it transiently. + return createErrorResult( + HTTP_STATUS.RATE_LIMITED, + `Token limit exceeded for ${scopeLabel}: ${tokenBreach.tokensUsed}/${tokenBreach.limitValue} tokens used in the current window. Please try again later.`, + null, + "TOKEN_LIMIT_EXCEEDED" + ); + } + } catch (err) { + // Fail-open at Tier 2: Tier 1 already enforced the model/global limit pre-dispatch. + // A transient counter read error here must not break an otherwise-valid request. + log?.warn?.("TOKEN_LIMIT", "Tier 2 token-limit check failed; allowing request", { err }); + } + } + // Execute request using executor (handles URL building, headers, fallback, transform) let providerResponse; let providerUrl; @@ -4821,6 +4869,15 @@ export async function handleChatCore({ }).catch((err) => { console.error("Failed to save usage stats:", err.message); }); + + if (apiKeyInfo?.id) { + try { + const billable = computeBillableTokens(usage); + if (billable > 0) recordTokenUsage(apiKeyInfo.id, provider || "unknown", model || "unknown", billable); + } catch { + // never block the response on counter recording + } + } } // Translate response to client's expected format (usually OpenAI) @@ -5201,6 +5258,15 @@ export async function handleChatCore({ }).catch((err) => { console.error("Failed to save usage stats:", err.message); }); + + if (apiKeyInfo?.id && streamStatus === 200) { + try { + const billable = computeBillableTokens(streamUsage); + if (billable > 0) recordTokenUsage(apiKeyInfo.id, provider || "unknown", model || "unknown", billable); + } catch { + // never block the stream on counter recording + } + } } persistAttemptLogs({ diff --git a/open-sse/services/combo.ts b/open-sse/services/combo.ts index 63bb9f7dd2..07857619b5 100644 --- a/open-sse/services/combo.ts +++ b/open-sse/services/combo.ts @@ -448,6 +448,20 @@ function isStreamReadinessFailureErrorBody(errorBody: unknown): boolean { return code === "STREAM_READINESS_TIMEOUT" || code === "STREAM_EARLY_EOF"; } +/** + * A local per-API-key token-limit breach surfaces as a 429 tagged with + * errorCode "TOKEN_LIMIT_EXCEEDED" (see chatCore.ts Tier 2 early return). This + * is NOT an upstream rate limit, so the combo loop must not cool the shared + * account/provider, must not add it to transientRateLimitedProviders, and must + * not retry it transiently — it propagates to the client as a terminal 429. + */ +function isTokenLimitBreachErrorBody(errorBody: unknown): boolean { + if (!errorBody || typeof errorBody !== "object") return false; + const error = (errorBody as Record).error; + if (!error || typeof error !== "object") return false; + return (error as Record).code === "TOKEN_LIMIT_EXCEEDED"; +} + function toRecordedTarget(target: ResolvedComboTarget) { return { executionKey: target.executionKey, @@ -3457,6 +3471,10 @@ export async function handleComboChat({ (result.status === 502 || result.status === 504) && isStreamReadinessFailureErrorBody(errorBody); + // FIX 5: a local per-API-key token-limit 429 must not cool shared accounts. + const isTokenLimitBreach = + result.status === 429 && isTokenLimitBreachErrorBody(errorBody); + // Fix #1681: Status 499 means client disconnected — stop combo loop immediately. // There is no point trying fallback models when nobody is listening. if (result.status === 499) { @@ -3510,7 +3528,12 @@ export async function handleComboChat({ "COMBO", `Provider ${provider} quota exhausted — marking for skip on remaining targets (#1731)` ); - } else if (result.status === 429 && provider && provider !== "unknown") { + } else if ( + result.status === 429 && + !isTokenLimitBreach && + provider && + provider !== "unknown" + ) { transientRateLimitedProviders.add(provider); } @@ -3532,9 +3555,12 @@ export async function handleComboChat({ recordProviderFailure(provider, log, target.connectionId, profile); } - // Check if this is a transient error worth retrying on same model + // Check if this is a transient error worth retrying on same model. + // A token-limit 429 is terminal for the client — never retry it. const isTransient = - !isStreamReadinessFailure && [408, 429, 500, 502, 503, 504].includes(result.status); + !isStreamReadinessFailure && + !isTokenLimitBreach && + [408, 429, 500, 502, 503, 504].includes(result.status); if (retry < maxRetries && isTransient && !providerExhausted) { continue; // Retry same model } @@ -3907,6 +3933,10 @@ async function handleRoundRobinCombo({ (result.status === 502 || result.status === 504) && isStreamReadinessFailureErrorBody(errorBody); + // FIX 5: a local per-API-key token-limit 429 must not cool shared accounts. + const isTokenLimitBreach = + result.status === 429 && isTokenLimitBreachErrorBody(errorBody); + // Round-robin uses the same target-level fallback rule as other combo // strategies: non-ok target responses fall through to the next target. // Classification stays here only to support cooldown/semaphore pacing, @@ -3949,13 +3979,19 @@ async function handleRoundRobinCombo({ if (providerExhausted) { exhaustedProviders.add(provider); log.info("COMBO-RR", `Provider ${provider} quota exhausted — marking for skip (#1731)`); - } else if (result.status === 429 && provider && provider !== "unknown") { + } else if ( + result.status === 429 && + !isTokenLimitBreach && + provider && + provider !== "unknown" + ) { transientRateLimitedProviders.add(provider); } // Transient errors → mark in semaphore so round-robin stops stampeding this target. if ( !isStreamReadinessFailure && + !isTokenLimitBreach && TRANSIENT_FOR_SEMAPHORE.includes(result.status) && cooldownMs > 0 ) { @@ -3970,9 +4006,12 @@ async function handleRoundRobinCombo({ ); } - // Transient error → retry same model + // Transient error → retry same model. + // A token-limit 429 is terminal for the client — never retry it. const isTransient = - !isStreamReadinessFailure && [408, 429, 500, 502, 503, 504].includes(result.status); + !isStreamReadinessFailure && + !isTokenLimitBreach && + [408, 429, 500, 502, 503, 504].includes(result.status); if (retry < maxRetries && isTransient && !providerExhausted) { continue; } diff --git a/open-sse/services/tokenLimitCounter.ts b/open-sse/services/tokenLimitCounter.ts new file mode 100644 index 0000000000..ad6fa38e8c --- /dev/null +++ b/open-sse/services/tokenLimitCounter.ts @@ -0,0 +1,342 @@ +/** + * Token Limit Counter — in-memory write-through accelerator. + * + * The DB (api_key_token_counters) is the source of truth for enforcement. + * This module keeps a short-lived in-memory mirror of the current window's + * usage per limit to avoid a DB round-trip on every hot-path read, and seeds + * a cold window by SUMming usage_history for the active window when no counter + * row exists yet. + * + * Steps 006/007 (checkTokenLimits / recordTokenUsage) build on top of this. + * + * @module services/tokenLimitCounter + */ + +import { getDbInstance } from "../../src/lib/db/core.ts"; +import { + resetWindowIfElapsed, + getWindowUsage, + incrementWindowTokens, + getTokenLimitsForRequest, + logTokenLimitReset, + type TokenLimit, +} from "@/lib/localDb"; + +interface CacheEntry { + windowStart: string; + tokensUsed: number; + /** epoch-ms when this cache entry was last synced from / written to the DB */ + syncedAt: number; +} + +/** key = `${limitId}` ; value tracks the *current* window only */ +const cache = new Map(); + +/** Cache entries older than this are re-validated against the DB on read. */ +const CACHE_TTL_MS = 5_000; + +/** + * Sum billable tokens recorded in usage_history for this limit's API key within + * the active window, filtered by scope (model / provider / global). Used to seed + * a cold counter so enforcement is correct even before the first write-through. + * + * Returns the windowed total (>= 0). DB-authoritative point-in-time read. + */ +export function seedWindowUsageFromHistory(limit: TokenLimit, now = Date.now()): number { + const { periodStartAt } = resetWindowIfElapsed(limit, now); + const lowerBound = new Date(periodStartAt).toISOString(); + const db = getDbInstance(); + + // Canonical billable total = input + output + reasoning. tokens_cache_read and + // tokens_cache_creation are a BREAKDOWN already inside tokens_input (see migration + // 012_fix_token_input_cache_tokens.sql) — summing them again would double-count. + // This must mirror computeBillableTokens() in chatCore.ts. + const tokenSum = `COALESCE(SUM( + COALESCE(tokens_input, 0) + COALESCE(tokens_output, 0) + + COALESCE(tokens_reasoning, 0) + ), 0) AS total`; + + let row: unknown; + if (limit.scopeType === "model") { + row = db + .prepare( + `SELECT ${tokenSum} FROM usage_history + WHERE api_key_id = ? AND model = ? AND timestamp >= ?` + ) + .get(limit.apiKeyId, limit.scopeValue, lowerBound); + } else if (limit.scopeType === "provider") { + row = db + .prepare( + `SELECT ${tokenSum} FROM usage_history + WHERE api_key_id = ? AND provider = ? AND timestamp >= ?` + ) + .get(limit.apiKeyId, limit.scopeValue, lowerBound); + } else { + row = db + .prepare( + `SELECT ${tokenSum} FROM usage_history + WHERE api_key_id = ? AND timestamp >= ?` + ) + .get(limit.apiKeyId, lowerBound); + } + + const total = row && typeof row === "object" ? (row as { total?: unknown }).total : 0; + const n = typeof total === "number" ? total : Number(total); + return Number.isFinite(n) && n > 0 ? Math.floor(n) : 0; +} + +/** + * Current-window usage for a limit. DB is authoritative; the in-memory cache is + * a read accelerator only. On a fresh window with no DB counter row, seeds the + * value from usage_history so enforcement is correct immediately. + * + * `forceFresh` bypasses the cache (used by the authoritative enforcement read). + */ +export function getCurrentWindowUsage( + limit: TokenLimit, + now = Date.now(), + forceFresh = false +): number { + const { windowStart } = resetWindowIfElapsed(limit, now); + const cached = cache.get(limit.id); + + if ( + !forceFresh && + cached && + cached.windowStart === windowStart && + now - cached.syncedAt < CACHE_TTL_MS + ) { + return cached.tokensUsed; + } + + // DB point-read (authoritative for the window). + let dbUsage = getWindowUsage(limit, now); + + // Cold window: no counter row yet → seed from usage_history and PERSIST the + // seed to the counter row so a subsequent recordTokenUsage increment does not + // forget the existing historical usage in this window (force-fresh read then + // first record must accumulate on top of history, not restart from 0). + if (dbUsage === 0 && (!cached || cached.windowStart !== windowStart)) { + const seeded = seedWindowUsageFromHistory(limit, now); + if (seeded > 0) { + // UPSERT creates the row at `seeded`; safe because there is no row yet + // (dbUsage === 0). Returns the new authoritative total. + dbUsage = incrementWindowTokens(limit.id, windowStart, seeded); + } + } + + cache.set(limit.id, { windowStart, tokensUsed: dbUsage, syncedAt: now }); + return dbUsage; +} + +/** + * Atomically add `tokens` to the DB counter for the limit's current window and + * update the in-memory cache to the new authoritative total. Returns the total. + * + * NOTE: callers that need transactional reset-detection (step 007) should call + * the DB increment inside their own transaction; this helper is the simple + * write-through used outside a transaction. + */ +export function addWindowTokens(limit: TokenLimit, tokens: number, now = Date.now()): number { + const { windowStart } = resetWindowIfElapsed(limit, now); + const delta = tokens > 0 ? Math.floor(tokens) : 0; + const newTotal = incrementWindowTokens(limit.id, windowStart, delta); + cache.set(limit.id, { windowStart, tokensUsed: newTotal, syncedAt: now }); + return newTotal; +} + +/** Overwrite the cache entry for a limit (e.g. after a transactional increment in step 007). */ +export function syncCache(limitId: string, windowStart: string, tokensUsed: number): void { + cache.set(limitId, { windowStart, tokensUsed, syncedAt: Date.now() }); +} + +/** Drop a single limit's cache entry (e.g. on limit delete/update). */ +export function invalidateLimit(limitId: string): void { + cache.delete(limitId); +} + +/** Clear the whole cache (tests / config reload). */ +export function clearTokenLimitCache(): void { + cache.clear(); +} + +/** + * Breach detail returned when an applicable token limit is exceeded. `remaining` + * is `limitValue - tokensUsed` (clamped at 0). The breach with the SMALLEST + * remaining is the most-restrictive and is the one returned by checkTokenLimits. + */ +export interface TokenLimitBreach { + limitId: string; + scopeType: TokenLimit["scopeType"]; + scopeValue: string; + limitValue: number; + tokensUsed: number; + remaining: number; + windowStart: string; + nextResetAt: number; +} + +/** + * Enforcement check. Loads every enabled limit applicable to (apiKeyId, provider, + * model) — model-scoped, provider-scoped, and global — reads DB-authoritative + * window usage for each (forceFresh, bypassing the read cache), and returns the + * most-restrictive breach (smallest remaining tokens) or null if none breach. + * + * A limit is breached when window usage is at or above the configured limit + * (tokensUsed >= limitValue), i.e. there is no remaining budget for more tokens. + * + * @param apiKeyId the API key id (required) + * @param provider resolved upstream provider id (optional; "" matches no provider scope) + * @param model resolved model id (optional; "" matches no model scope) + */ +export function checkTokenLimits( + apiKeyId: string, + provider = "", + model = "", + now = Date.now() +): TokenLimitBreach | null { + if (!apiKeyId) return null; + + const limits = getTokenLimitsForRequest(apiKeyId, provider, model); + if (!limits || limits.length === 0) return null; + + let worst: TokenLimitBreach | null = null; + + for (const limit of limits) { + if (limit.enabled === false) continue; + const limitValue = limit.tokenLimit; + if (!Number.isFinite(limitValue) || limitValue <= 0) continue; + + // Authoritative read (bypass the in-memory accelerator). + const tokensUsed = getCurrentWindowUsage(limit, now, true); + if (tokensUsed < limitValue) continue; // within budget + + const { windowStart, nextResetAt } = resetWindowIfElapsed(limit, now); + const remaining = Math.max(0, limitValue - tokensUsed); + const breach: TokenLimitBreach = { + limitId: limit.id, + scopeType: limit.scopeType, + scopeValue: limit.scopeValue, + limitValue, + tokensUsed, + remaining, + windowStart, + nextResetAt, + }; + + // Most-restrictive wins: smallest remaining. Tie-break: smaller limitValue. + if ( + worst === null || + breach.remaining < worst.remaining || + (breach.remaining === worst.remaining && breach.limitValue < worst.limitValue) + ) { + worst = breach; + } + } + + return worst; +} + +/** + * Record token consumption against every applicable token limit for a request. + * + * FIRE-AND-FORGET: scheduled on a microtask so it NEVER blocks the SSE stream. + * The DB work runs inside a synchronous better-sqlite3 transaction so the + * reset-detection + reset-log + atomic increment for all matching limits commit + * atomically. The in-memory cache is updated to the new authoritative totals. + * + * A rollover (window reset) is detected when the current window has no counter + * row yet but a prior window row for the same limit still holds usage; in that + * case a reset-log row is written before incrementing the new window. + * + * @param apiKeyId API key id (no-op if falsy) + * @param provider resolved upstream provider id + * @param model resolved model id + * @param tokens billable tokens consumed by this request (no-op if <= 0) + */ +export function recordTokenUsage( + apiKeyId: string, + provider: string, + model: string, + tokens: number +): void { + if (!apiKeyId) return; + const delta = Number.isFinite(tokens) && tokens > 0 ? Math.floor(tokens) : 0; + if (delta <= 0) return; + + // Schedule off the hot path; never await, never block the stream. + Promise.resolve() + .then(() => { + const now = Date.now(); + const limits = getTokenLimitsForRequest(apiKeyId, provider || "", model || ""); + if (!limits || limits.length === 0) return; + + const db = getDbInstance(); + const applied: Array<{ limitId: string; windowStart: string; total: number }> = []; + + const tx = db.transaction(() => { + for (const limit of limits) { + if (limit.enabled === false) continue; + + const { windowStart } = resetWindowIfElapsed(limit, now); + + const currentRow = db + .prepare( + "SELECT tokens_used FROM api_key_token_counters WHERE limit_id = ? AND window_start = ?" + ) + .get(limit.id, windowStart) as { tokens_used?: number } | undefined; + + // First write to a new window? Detect rollover from the prior window. + if (!currentRow) { + const priorRow = db + .prepare( + `SELECT window_start, tokens_used FROM api_key_token_counters + WHERE limit_id = ? AND window_start < ? + ORDER BY window_start DESC LIMIT 1` + ) + .get(limit.id, windowStart) as + | { window_start?: string; tokens_used?: number } + | undefined; + const prevTokens = + priorRow && typeof priorRow.tokens_used === "number" ? priorRow.tokens_used : 0; + if (prevTokens > 0) { + logTokenLimitReset(limit.id, prevTokens, windowStart); + } + + // Cold window with no counter row: seed from usage_history so the + // running total reflects prior usage already recorded in this window + // before applying the new delta. Synchronous (better-sqlite3) — safe + // inside this transaction. Mirrors getCurrentWindowUsage seed-on-miss. + const seeded = seedWindowUsageFromHistory(limit, now); + if (seeded > 0) { + incrementWindowTokens(limit.id, windowStart, seeded); + } + } + + const total = incrementWindowTokens(limit.id, windowStart, delta); + applied.push({ limitId: limit.id, windowStart, total }); + } + }); + + try { + tx(); + // Update the read accelerator to the new authoritative totals. + for (const a of applied) { + syncCache(a.limitId, a.windowStart, a.total); + } + } catch (err) { + // better-sqlite3 auto-rolls-back on throw; verify we are not stuck mid-txn. + if (db.inTransaction) { + try { + db.exec("ROLLBACK"); + } catch { + // already rolled back + } + } + // Swallow — usage recording must never surface to the request path. + } + }) + .catch(() => { + // Microtask scheduling/setup failure — non-fatal, never blocks the stream. + }); +} diff --git a/src/app/api/usage/token-limits/route.ts b/src/app/api/usage/token-limits/route.ts new file mode 100644 index 0000000000..18d5ab664e --- /dev/null +++ b/src/app/api/usage/token-limits/route.ts @@ -0,0 +1,98 @@ +/** + * Per-API-Key Token Limits — CRUD Route + * + * Management-class endpoint for listing, creating/updating, and deleting + * token-limit budgets attached to an API key (model / provider / global scope). + * Auth and CORS are enforced centrally by the global authz pipeline + * (src/proxy.ts → runAuthzPipeline); this file intentionally adds neither. + * + * @route /api/usage/token-limits + */ + +import { NextResponse } from "next/server"; +import { setTokenLimitSchema } from "@/shared/validation/schemas"; +import { isValidationFailure, validateBody } from "@/shared/validation/helpers"; +import { buildErrorBody, sanitizeErrorMessage } from "@omniroute/open-sse/utils/error"; +import { + listTokenLimits, + upsertTokenLimit, + deleteTokenLimit, + getWindowUsage, + resetWindowIfElapsed, +} from "@/lib/localDb"; +import type { TokenLimit } from "@/lib/db/tokenLimits"; + +export async function GET(request: Request) { + try { + const { searchParams } = new URL(request.url); + const apiKeyId = searchParams.get("apiKeyId"); + if (!apiKeyId) { + return NextResponse.json(buildErrorBody(400, "apiKeyId query param is required"), { + status: 400, + }); + } + const limits = listTokenLimits(apiKeyId).map((limit: TokenLimit) => { + const usage = getWindowUsage(limit); + const window = resetWindowIfElapsed(limit); + return { + ...limit, + tokensUsed: usage, + windowStart: window.windowStart, + periodStartAt: window.periodStartAt, + nextResetAt: window.nextResetAt, + remaining: Math.max(0, limit.tokenLimit - usage), + }; + }); + return NextResponse.json({ apiKeyId, limits }); + } catch (error) { + console.error("Error listing token limits:", error); + return NextResponse.json(buildErrorBody(500, sanitizeErrorMessage(error)), { status: 500 }); + } +} + +export async function POST(request: Request) { + let rawBody: unknown; + try { + rawBody = await request.json(); + } catch { + return NextResponse.json(buildErrorBody(400, "Invalid JSON body"), { status: 400 }); + } + + try { + const validation = validateBody(setTokenLimitSchema, rawBody); + if (isValidationFailure(validation)) { + return NextResponse.json({ error: validation.error }, { status: 400 }); + } + const { id, apiKeyId, scopeType, scopeValue, tokenLimit, resetInterval, resetTime, enabled } = + validation.data; + const limit = upsertTokenLimit({ + id, + apiKeyId, + scopeType, + scopeValue, + tokenLimit, + resetInterval, + resetTime, + enabled, + }); + return NextResponse.json({ success: true, limit }); + } catch (error) { + console.error("Error setting token limit:", error); + return NextResponse.json(buildErrorBody(500, sanitizeErrorMessage(error)), { status: 500 }); + } +} + +export async function DELETE(request: Request) { + const { searchParams } = new URL(request.url); + const id = searchParams.get("id"); + if (!id) { + return NextResponse.json(buildErrorBody(400, "id query param is required"), { status: 400 }); + } + try { + const deleted = deleteTokenLimit(id); + return NextResponse.json({ success: deleted }, { status: deleted ? 200 : 404 }); + } catch (error) { + console.error("Error deleting token limit:", error); + return NextResponse.json(buildErrorBody(500, sanitizeErrorMessage(error)), { status: 500 }); + } +} diff --git a/src/lib/db/migrations/073_per_model_token_limits.sql b/src/lib/db/migrations/073_per_model_token_limits.sql new file mode 100644 index 0000000000..e64360f4aa --- /dev/null +++ b/src/lib/db/migrations/073_per_model_token_limits.sql @@ -0,0 +1,61 @@ +-- Migration 073: per-API-key token limits (scoped to model and/or provider) +-- +-- Adds enforcement-grade token budgets attachable to an API key and scoped to a +-- specific model, a specific provider, or globally (all traffic for the key). +-- When a key exceeds its configured token budget within the active reset window, +-- requests are rejected with HTTP 429. This complements (does not replace) the +-- USD cost budgets in `domain_budgets` / src/domain/costRules.ts. Token usage is +-- already captured per request in usage_history; these tables provide the +-- authoritative rolling-window counters used for pre-flight enforcement. +-- +-- Three tables: +-- api_key_token_limits - the limit definitions (one per scope per key) +-- api_key_token_counters - rolling-window usage counters (implicit rollover) +-- api_key_token_limit_reset_logs - audit trail of window resets +-- +-- Plus a composite index on usage_history to make the cold-start seed-on-miss +-- SUM (api_key_id, provider, model, timestamp) an index scan rather than a table +-- scan. The build does NOT run with PRAGMA foreign_keys=ON, so the FK clause is +-- declarative only (matches existing repo convention; application code is the +-- single writer and prunes counters via ON DELETE CASCADE semantics best-effort). + +CREATE TABLE IF NOT EXISTS api_key_token_limits ( + id TEXT PRIMARY KEY, + api_key_id TEXT NOT NULL, + scope_type TEXT NOT NULL CHECK (scope_type IN ('model', 'provider', 'global')), + scope_value TEXT NOT NULL DEFAULT '', + token_limit INTEGER NOT NULL CHECK (token_limit > 0), + reset_interval TEXT NOT NULL DEFAULT 'monthly' CHECK (reset_interval IN ('daily', 'weekly', 'monthly')), + reset_time TEXT, + enabled INTEGER NOT NULL DEFAULT 1, + created_at TEXT NOT NULL DEFAULT (datetime('now')), + updated_at TEXT NOT NULL DEFAULT (datetime('now')), + UNIQUE (api_key_id, scope_type, scope_value) +); + +CREATE INDEX IF NOT EXISTS idx_aktl_api_key_id + ON api_key_token_limits (api_key_id); + +CREATE TABLE IF NOT EXISTS api_key_token_counters ( + limit_id TEXT NOT NULL, + window_start TEXT NOT NULL, + tokens_used INTEGER NOT NULL DEFAULT 0, + updated_at TEXT NOT NULL DEFAULT (datetime('now')), + PRIMARY KEY (limit_id, window_start), + FOREIGN KEY (limit_id) REFERENCES api_key_token_limits (id) ON DELETE CASCADE +); + +CREATE TABLE IF NOT EXISTS api_key_token_limit_reset_logs ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + limit_id TEXT NOT NULL, + reset_at TEXT NOT NULL DEFAULT (datetime('now')), + prev_tokens INTEGER NOT NULL DEFAULT 0, + window_start TEXT NOT NULL, + FOREIGN KEY (limit_id) REFERENCES api_key_token_limits (id) ON DELETE CASCADE +); + +CREATE INDEX IF NOT EXISTS idx_aktlrl_limit_id + ON api_key_token_limit_reset_logs (limit_id); + +CREATE INDEX IF NOT EXISTS idx_uh_key_provider_model_ts + ON usage_history (api_key_id, provider, model, timestamp); diff --git a/src/lib/db/tokenLimits.ts b/src/lib/db/tokenLimits.ts new file mode 100644 index 0000000000..820b6c6920 --- /dev/null +++ b/src/lib/db/tokenLimits.ts @@ -0,0 +1,295 @@ +/** + * Per-API-Key Token Limits — Persistence Layer + * + * Enforcement-grade token budgets attachable to an API key, scoped to a + * specific model, a specific provider, or globally. Complements the USD cost + * budgets in domainState.ts / src/domain/costRules.ts. + * + * Tables (migration 073): api_key_token_limits, api_key_token_counters, + * api_key_token_limit_reset_logs. + * + * @module lib/db/tokenLimits + */ + +import { randomUUID } from "crypto"; +import { getDbInstance } from "./core"; +import { getBudgetWindow, type BudgetResetInterval } from "@/domain/costRules"; + +export type TokenLimitScopeType = "model" | "provider" | "global"; + +export interface TokenLimit { + id: string; + apiKeyId: string; + scopeType: TokenLimitScopeType; + scopeValue: string; + tokenLimit: number; + resetInterval: BudgetResetInterval; + resetTime: string; + enabled: boolean; + createdAt: string; + updatedAt: string; +} + +export interface UpsertTokenLimitInput { + id?: string; + apiKeyId: string; + scopeType: TokenLimitScopeType; + scopeValue?: string; + tokenLimit: number; + resetInterval?: BudgetResetInterval; + resetTime?: string; + enabled?: boolean; +} + +export interface TokenWindowState { + windowStart: string; + didReset: boolean; + periodStartAt: number; + nextResetAt: number; +} + +type JsonRecord = Record; + +let _schemaChecked = false; + +function asRecord(value: unknown): JsonRecord { + return value && typeof value === "object" && !Array.isArray(value) ? (value as JsonRecord) : {}; +} + +function toNumber(value: unknown, fallback = 0): number { + if (typeof value === "number" && Number.isFinite(value)) return value; + if (typeof value === "string" && value.trim().length > 0) { + const parsed = Number(value); + return Number.isFinite(parsed) ? parsed : fallback; + } + return fallback; +} + +function normalizeScopeType(value: unknown): TokenLimitScopeType { + if (value === "model" || value === "provider" || value === "global") return value; + return "global"; +} + +function normalizeResetInterval(value: unknown): BudgetResetInterval { + if (value === "daily" || value === "weekly" || value === "monthly") return value; + return "monthly"; +} + +function ensureSchema() { + if (_schemaChecked) return; + const db = getDbInstance(); + db.exec(` + CREATE TABLE IF NOT EXISTS api_key_token_limits ( + id TEXT PRIMARY KEY, + api_key_id TEXT NOT NULL, + scope_type TEXT NOT NULL CHECK (scope_type IN ('model', 'provider', 'global')), + scope_value TEXT NOT NULL DEFAULT '', + token_limit INTEGER NOT NULL CHECK (token_limit > 0), + reset_interval TEXT NOT NULL DEFAULT 'monthly' CHECK (reset_interval IN ('daily', 'weekly', 'monthly')), + reset_time TEXT, + enabled INTEGER NOT NULL DEFAULT 1, + created_at TEXT NOT NULL DEFAULT (datetime('now')), + updated_at TEXT NOT NULL DEFAULT (datetime('now')), + UNIQUE (api_key_id, scope_type, scope_value) + ); + CREATE INDEX IF NOT EXISTS idx_aktl_api_key_id ON api_key_token_limits (api_key_id); + CREATE TABLE IF NOT EXISTS api_key_token_counters ( + limit_id TEXT NOT NULL, + window_start TEXT NOT NULL, + tokens_used INTEGER NOT NULL DEFAULT 0, + updated_at TEXT NOT NULL DEFAULT (datetime('now')), + PRIMARY KEY (limit_id, window_start), + FOREIGN KEY (limit_id) REFERENCES api_key_token_limits (id) ON DELETE CASCADE + ); + CREATE TABLE IF NOT EXISTS api_key_token_limit_reset_logs ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + limit_id TEXT NOT NULL, + reset_at TEXT NOT NULL DEFAULT (datetime('now')), + prev_tokens INTEGER NOT NULL DEFAULT 0, + window_start TEXT NOT NULL, + FOREIGN KEY (limit_id) REFERENCES api_key_token_limits (id) ON DELETE CASCADE + ); + CREATE INDEX IF NOT EXISTS idx_aktlrl_limit_id ON api_key_token_limit_reset_logs (limit_id); + `); + _schemaChecked = true; +} + +function rowToTokenLimit(row: unknown): TokenLimit { + const r = asRecord(row); + return { + id: typeof r.id === "string" ? r.id : "", + apiKeyId: typeof r.api_key_id === "string" ? r.api_key_id : "", + scopeType: normalizeScopeType(r.scope_type), + scopeValue: typeof r.scope_value === "string" ? r.scope_value : "", + tokenLimit: toNumber(r.token_limit), + resetInterval: normalizeResetInterval(r.reset_interval), + resetTime: typeof r.reset_time === "string" && r.reset_time ? r.reset_time : "00:00", + enabled: toNumber(r.enabled, 1) !== 0, + createdAt: typeof r.created_at === "string" ? r.created_at : "", + updatedAt: typeof r.updated_at === "string" ? r.updated_at : "", + }; +} + +// ──────────────── CRUD ──────────────── + +/** + * Insert or update a token limit. Upsert key is (api_key_id, scope_type, scope_value). + * Returns the persisted row. + */ +export function upsertTokenLimit(input: UpsertTokenLimitInput): TokenLimit { + ensureSchema(); + const db = getDbInstance(); + const scopeType = normalizeScopeType(input.scopeType); + const scopeValue = scopeType === "global" ? "" : (input.scopeValue ?? "").trim(); + const resetInterval = normalizeResetInterval(input.resetInterval); + const resetTime = + typeof input.resetTime === "string" && input.resetTime ? input.resetTime : "00:00"; + const enabled = input.enabled === false ? 0 : 1; + const tokenLimit = Math.floor(toNumber(input.tokenLimit)); + const id = input.id && input.id.trim() ? input.id.trim() : randomUUID(); + + db.prepare( + `INSERT INTO api_key_token_limits + (id, api_key_id, scope_type, scope_value, token_limit, reset_interval, reset_time, enabled, created_at, updated_at) + VALUES (@id, @apiKeyId, @scopeType, @scopeValue, @tokenLimit, @resetInterval, @resetTime, @enabled, datetime('now'), datetime('now')) + ON CONFLICT(api_key_id, scope_type, scope_value) + DO UPDATE SET token_limit = excluded.token_limit, + reset_interval = excluded.reset_interval, + reset_time = excluded.reset_time, + enabled = excluded.enabled, + updated_at = datetime('now')` + ).run({ id, apiKeyId: input.apiKeyId, scopeType, scopeValue, tokenLimit, resetInterval, resetTime, enabled }); + + const row = db + .prepare( + "SELECT * FROM api_key_token_limits WHERE api_key_id = ? AND scope_type = ? AND scope_value = ?" + ) + .get(input.apiKeyId, scopeType, scopeValue); + return rowToTokenLimit(row); +} + +/** List all token limits for an API key (ordered most-specific first: model, provider, global). */ +export function listTokenLimits(apiKeyId: string): TokenLimit[] { + ensureSchema(); + const db = getDbInstance(); + return db + .prepare( + `SELECT * FROM api_key_token_limits + WHERE api_key_id = ? + ORDER BY CASE scope_type WHEN 'model' THEN 0 WHEN 'provider' THEN 1 ELSE 2 END, scope_value` + ) + .all(apiKeyId) + .map(rowToTokenLimit); +} + +/** + * Return the enabled limits that apply to a given request: the model-scoped row + * (scope_value === model), the provider-scoped row (scope_value === provider), + * and the global row. Used by the enforcement read. + */ +export function getTokenLimitsForRequest( + apiKeyId: string, + provider: string, + model: string +): TokenLimit[] { + ensureSchema(); + const db = getDbInstance(); + return db + .prepare( + `SELECT * FROM api_key_token_limits + WHERE api_key_id = @apiKeyId + AND enabled = 1 + AND ( + (scope_type = 'global') + OR (scope_type = 'model' AND scope_value = @model) + OR (scope_type = 'provider' AND scope_value = @provider) + )` + ) + .all({ apiKeyId, model: model || "", provider: provider || "" } as JsonRecord) + .map(rowToTokenLimit); +} + +/** Delete a token limit by id (counters + reset logs cascade in app code below). */ +export function deleteTokenLimit(id: string): boolean { + ensureSchema(); + const db = getDbInstance(); + // FK pragma is OFF in this build; delete dependents explicitly. + db.prepare("DELETE FROM api_key_token_counters WHERE limit_id = ?").run(id); + db.prepare("DELETE FROM api_key_token_limit_reset_logs WHERE limit_id = ?").run(id); + const info = db.prepare("DELETE FROM api_key_token_limits WHERE id = ?").run(id); + return info.changes > 0; +} + +// ──────────────── Counters (concurrency-safe under WAL) ──────────────── + +/** + * Pure boundary calculator. Resolves the active window for a limit at `now` + * and reports whether the window has rolled since the limit's last-known window. + * No DB writes. Window math reused from costRules.getBudgetWindow. + */ +export function resetWindowIfElapsed(limit: TokenLimit, now = Date.now()): TokenWindowState { + const window = getBudgetWindow(limit.resetInterval, limit.resetTime, now); + const windowStart = String(window.periodStartAt); + return { + windowStart, + didReset: false, // caller compares against the stored counter row to decide rollover + periodStartAt: window.periodStartAt, + nextResetAt: window.nextResetAt, + }; +} + +/** + * Read-only point-read of the current window's usage for a limit. + * Returns 0 if no counter row exists yet (cold window). DB-authoritative. + */ +export function getWindowUsage(limit: TokenLimit, now = Date.now()): number { + ensureSchema(); + const db = getDbInstance(); + const { windowStart } = resetWindowIfElapsed(limit, now); + const row = db + .prepare( + "SELECT tokens_used FROM api_key_token_counters WHERE limit_id = ? AND window_start = ?" + ) + .get(limit.id, windowStart); + return toNumber(asRecord(row).tokens_used); +} + +/** + * Atomically add `tokens` to the counter for (limitId, windowStart) and return + * the new running total. Uses UPSERT (no read-then-write) so concurrent + * increments under WAL cannot lose updates. + */ +export function incrementWindowTokens( + limitId: string, + windowStart: string, + tokens: number +): number { + ensureSchema(); + const db = getDbInstance(); + const delta = Math.max(0, Math.floor(toNumber(tokens))); + const row = db + .prepare( + `INSERT INTO api_key_token_counters (limit_id, window_start, tokens_used, updated_at) + VALUES (@limitId, @windowStart, @tokens, datetime('now')) + ON CONFLICT(limit_id, window_start) + DO UPDATE SET tokens_used = tokens_used + excluded.tokens_used, + updated_at = datetime('now') + RETURNING tokens_used` + ) + .get({ limitId, windowStart, tokens: delta }); + return toNumber(asRecord(row).tokens_used); +} + +/** Append a window-reset audit log row. */ +export function logTokenLimitReset( + limitId: string, + prevTokens: number, + windowStart: string +): void { + ensureSchema(); + const db = getDbInstance(); + db.prepare( + `INSERT INTO api_key_token_limit_reset_logs (limit_id, reset_at, prev_tokens, window_start) + VALUES (?, datetime('now'), ?, ?)` + ).run(limitId, Math.max(0, Math.floor(toNumber(prevTokens))), windowStart); +} diff --git a/src/lib/localDb.ts b/src/lib/localDb.ts index 648c03f658..3ac0ca7847 100755 --- a/src/lib/localDb.ts +++ b/src/lib/localDb.ts @@ -508,3 +508,22 @@ export { } from "./db/freeProxies"; export type { FreeProxyRecord, FreeProxyStats } from "./db/freeProxies"; + +export { + // Per-API-Key Token Limits (migration 073) + upsertTokenLimit, + listTokenLimits, + getTokenLimitsForRequest, + deleteTokenLimit, + getWindowUsage, + incrementWindowTokens, + resetWindowIfElapsed, + logTokenLimitReset, +} from "./db/tokenLimits"; + +export type { + TokenLimit, + TokenLimitScopeType, + UpsertTokenLimitInput, + TokenWindowState, +} from "./db/tokenLimits"; diff --git a/src/shared/utils/apiKeyPolicy.ts b/src/shared/utils/apiKeyPolicy.ts index 5ec03199f6..17530ced45 100644 --- a/src/shared/utils/apiKeyPolicy.ts +++ b/src/shared/utils/apiKeyPolicy.ts @@ -12,6 +12,7 @@ import { extractApiKey } from "@/sse/services/auth"; import { getApiKeyMetadata, getComboByName, isModelAllowedForKey } from "@/lib/localDb"; import { resolveComboForModel } from "@/lib/db/modelComboMappings"; import { checkBudget } from "@/domain/costRules"; +import { checkTokenLimits } from "@omniroute/open-sse/services/tokenLimitCounter.ts"; import { errorResponse } from "@omniroute/open-sse/utils/error.ts"; import { HTTP_STATUS } from "@omniroute/open-sse/config/constants.ts"; import * as log from "@/sse/utils/logger"; @@ -394,6 +395,39 @@ export async function enforceApiKeyPolicy( } } + // ── Check 4.5: Per-model / per-provider token limits (Tier 1) ── + if (apiKeyInfo.id) { + try { + const breach = checkTokenLimits(apiKeyInfo.id, undefined, modelStr ?? undefined); + if (breach) { + const scopeLabel = + breach.scopeType === "global" + ? "account" + : `${breach.scopeType} "${breach.scopeValue}"`; + return { + apiKey, + apiKeyInfo, + rejection: errorResponse( + HTTP_STATUS.RATE_LIMITED, + `Token limit exceeded for ${scopeLabel}: ${breach.tokensUsed}/${breach.limitValue} tokens used in the current window. Please try again later.` + ), + }; + } + } catch (error) { + // Fail-closed: token-limit backend error should block the request, + // consistent with the budget check above. + log.error("API_POLICY", "Token limit check failed. Request blocked.", { error }); + return { + apiKey, + apiKeyInfo, + rejection: errorResponse( + HTTP_STATUS.SERVICE_UNAVAILABLE, + "Token limit policy unavailable" + ), + }; + } + } + // ── Check 5: Generic Multi-Window Rate Limits ── if (apiKeyInfo.id) { const hasCustomRateLimits = Boolean(apiKeyInfo.rateLimits && apiKeyInfo.rateLimits.length > 0); diff --git a/src/shared/validation/schemas.ts b/src/shared/validation/schemas.ts index 8316329a94..53bb9e8bae 100644 --- a/src/shared/validation/schemas.ts +++ b/src/shared/validation/schemas.ts @@ -853,6 +853,34 @@ export const setBudgetSchema = z } }); +export const setTokenLimitSchema = z + .object({ + id: z.string().trim().min(1).optional(), + apiKeyId: z.string().trim().min(1, "apiKeyId is required"), + scopeType: z.enum(["model", "provider", "global"]), + scopeValue: z.string().trim().default(""), + tokenLimit: z.coerce + .number() + .int("tokenLimit must be an integer") + .positive("tokenLimit must be greater than zero"), + resetInterval: z.enum(["daily", "weekly", "monthly"]).default("monthly"), + resetTime: z + .string() + .trim() + .regex(/^\d{2}:\d{2}$/, "resetTime must be in HH:MM format") + .optional(), + enabled: z.boolean().default(true), + }) + .superRefine((value, ctx) => { + if (value.scopeType !== "global" && (!value.scopeValue || value.scopeValue.length === 0)) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: "scopeValue is required unless scopeType is 'global'", + path: ["scopeValue"], + }); + } + }); + export const policyActionSchema = z .object({ action: z.enum(["unlock"]), diff --git a/tests/unit/token-limits.test.ts b/tests/unit/token-limits.test.ts new file mode 100644 index 0000000000..aeb802626e --- /dev/null +++ b/tests/unit/token-limits.test.ts @@ -0,0 +1,377 @@ +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); +});