pjandro/OmniRoute
1
0
Fork 0
mirror of https://github.com/diegosouzapw/OmniRoute.git synced 2026-07-26 09:52:11 +03:00
Code Issues Packages Projects Releases Wiki Activity

feat(usage): per-API-key token limits scoped to model/provider/global (#2888)

Browse Source 
Integrated into release/v3.8.6
This commit is contained in:
Muhammad Mugni Hadi
2026-05-29 15:29:10 +07:00
committed by GitHub
parent d96eeef9c1
commit 1461a78d07
11 changed files with 1400 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
docs/reference/API_REFERENCE.md
View File

@@ -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/*`

68
open-sse/handlers/chatCore.ts
View File

@@ -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({

51
open-sse/services/combo.ts
View File

@@ -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<string, unknown>).error;
if (!error || typeof error !== "object") return false;
return (error as Record<string, unknown>).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;
}

342
open-sse/services/tokenLimitCounter.ts Normal file
View File

@@ -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<string, CacheEntry>();
/** 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.
});
}

98
src/app/api/usage/token-limits/route.ts Normal file
View File

@@ -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 });
}
}

61
src/lib/db/migrations/073_per_model_token_limits.sql Normal file
View File

@@ -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);

295
src/lib/db/tokenLimits.ts Normal file
View File

@@ -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<string, unknown>;
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);
}

19
src/lib/localDb.ts
View File

@@ -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";

34
src/shared/utils/apiKeyPolicy.ts
View File

@@ -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);

28
src/shared/validation/schemas.ts
View File

@@ -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"]),

377
tests/unit/token-limits.test.ts Normal file
View File

@@ -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);
});