diff --git a/open-sse/handlers/chatCore.ts b/open-sse/handlers/chatCore.ts index 397dae6eaa3..c4dbf8ef3cd 100644 --- a/open-sse/handlers/chatCore.ts +++ b/open-sse/handlers/chatCore.ts @@ -275,7 +275,7 @@ import { ensureEngineBreakdown } from "../services/compression/engineBreakdown.t import { handleBypassRequest } from "../utils/bypassHandler.ts"; import { saveRequestUsage, trackPendingRequest, appendRequestLog } from "@/lib/usageDb"; import { finalizePendingScope, updatePendingScope } from "@/lib/usage/pendingRequestScope"; -import { recordCost } from "@/domain/costRules"; +import { recordCost, recordChatCallCost, buildCostCtx } from "@/domain/costRules"; import { calculateCost } from "@/lib/usage/costCalculator"; import { buildClaudePassthroughToolNameMap, @@ -309,7 +309,7 @@ import { } from "./chatCore/pluginOnResponse.ts"; import { scheduleStreamingQuotaShareConsumption } from "./chatCore/streamingQuotaShare.ts"; import { recordStreamingUsageStats } from "./chatCore/streamingUsageStats.ts"; -import { recordStreamingCost } from "./chatCore/streamingCost.ts"; +import { recordStreamingCost, buildStreamLedgerDetails } from "./chatCore/streamingCost.ts"; import { isJsonRecord } from "./chatCore/nonStreamingResponseParse.ts"; import { recordNonStreamingUsageStats } from "./chatCore/nonStreamingUsageStats.ts"; import { @@ -5447,6 +5447,7 @@ export async function handleChatCore({ const estimatedCost = costUsage ? await calculateCost(provider, model, costUsage, { serviceTier: effectiveServiceTier }) : 0; + const chatCostCtx = buildCostCtx(provider, model, usage, effectiveServiceTier, traceId); if (postCallGuardrails.blocked) { const guardrailMessage = postCallGuardrails.message || "Response blocked by guardrail"; @@ -5467,9 +5468,7 @@ export async function handleChatCore({ claudeCacheUsageMeta: cacheUsageLogMeta, cacheSource: "upstream", }); - if (apiKeyInfo?.id && estimatedCost > 0) { - recordCost(apiKeyInfo.id, estimatedCost); - } + recordChatCallCost(apiKeyInfo, estimatedCost, chatCostCtx, false); log?.warn?.( "GUARDRAIL", `Response blocked by ${postCallGuardrails.guardrail || "guardrail"}: ${guardrailMessage}` @@ -5607,9 +5606,7 @@ export async function handleChatCore({ claudeCacheUsageMeta: cacheUsageLogMeta, cacheSource: "upstream", }); - if (apiKeyInfo?.id && estimatedCost > 0) { - recordCost(apiKeyInfo.id, estimatedCost); - } + recordChatCallCost(apiKeyInfo, estimatedCost, chatCostCtx, true); // === Quota Share POST-hook (B/F7) — fire-and-forget, fail-open === await scheduleQuotaShareConsumption({ @@ -6041,6 +6038,7 @@ export async function handleChatCore({ serviceTier: effectiveServiceTier, calculateCost, recordCost, + ledger: buildStreamLedgerDetails(effectiveServiceTier, normalizedStreamStatus < 400, traceId), }); // === Quota Share POST-hook streaming (B/F7) — fire-and-forget, fail-open === diff --git a/open-sse/handlers/chatCore/nonStreamingUsageStats.ts b/open-sse/handlers/chatCore/nonStreamingUsageStats.ts index 13e4a9ae912..b4631a32b55 100644 --- a/open-sse/handlers/chatCore/nonStreamingUsageStats.ts +++ b/open-sse/handlers/chatCore/nonStreamingUsageStats.ts @@ -11,6 +11,7 @@ */ import { saveRequestUsage } from "@/lib/usageDb"; +import { recordKeyQuotaUsage } from "@/domain/keyQuota"; import { formatUsageLog } from "@/lib/usage/tokenAccounting"; import { COLORS } from "../../utils/stream.ts"; import { recordTokenUsage } from "../../services/tokenLimitCounter.ts"; @@ -71,6 +72,10 @@ function recordBillableTokens( if (!apiKeyInfo?.id) return; try { const billable = computeBillableTokens(usage); + // Key-quota tpm/rpm counters advance for every successful call regardless + // of token count (rpm always +1), so the gate in apiKeyPolicy sees current + // usage. Mirrors recordTokenUsage's swallow-and-continue. + recordKeyQuotaUsage(apiKeyInfo.id, billable); if (billable > 0) recordTokenUsage(apiKeyInfo.id, provider || "unknown", model || "unknown", billable); } catch { diff --git a/open-sse/handlers/chatCore/streamingCost.ts b/open-sse/handlers/chatCore/streamingCost.ts index 634c1ba9735..285a4122f60 100644 --- a/open-sse/handlers/chatCore/streamingCost.ts +++ b/open-sse/handlers/chatCore/streamingCost.ts @@ -16,6 +16,23 @@ type CostResolver = ( options: { serviceTier?: string } ) => Promise; +/** Extra per-request ledger context threaded from the handler. */ +type LedgerDetails = { + serviceTier?: string | null; + success?: boolean; + timestamp?: string; + requestId?: string | null; +}; + +/** Build a {@link LedgerDetails} for a completed streaming response. */ +export function buildStreamLedgerDetails( + serviceTier: string | null | undefined, + success: boolean, + requestId: string | null | undefined +): LedgerDetails { + return { serviceTier, success, timestamp: new Date().toISOString(), requestId }; +} + export function recordStreamingCost(args: { apiKeyId: string | null | undefined; provider: string | null | undefined; @@ -23,7 +40,20 @@ export function recordStreamingCost(args: { streamUsage: Record | null | undefined; serviceTier?: string; calculateCost: CostResolver; - recordCost: (apiKeyId: string, cost: number) => void; + recordCost: ( + apiKeyId: string, + cost: number, + details?: { + provider?: string | null; + model?: string | null; + tokens?: unknown; + serviceTier?: string | null; + success?: boolean; + timestamp?: string; + requestId?: string | null; + } + ) => void; + ledger?: LedgerDetails; }): void { if (!args.apiKeyId || !args.streamUsage) return; @@ -31,7 +61,17 @@ export function recordStreamingCost(args: { args .calculateCost(args.provider, args.model, args.streamUsage, { serviceTier: args.serviceTier }) .then((estimatedCost) => { - if (estimatedCost > 0) args.recordCost(apiKeyId, estimatedCost); + if (estimatedCost > 0) { + args.recordCost(apiKeyId, estimatedCost, { + provider: args.provider, + model: args.model, + tokens: args.streamUsage, + serviceTier: args.serviceTier ?? args.ledger?.serviceTier, + success: args.ledger?.success ?? true, + timestamp: args.ledger?.timestamp, + requestId: args.ledger?.requestId, + }); + } }) .catch(() => {}); } diff --git a/open-sse/handlers/chatCore/streamingUsageStats.ts b/open-sse/handlers/chatCore/streamingUsageStats.ts index 1a48b8d8572..d827a86de3c 100644 --- a/open-sse/handlers/chatCore/streamingUsageStats.ts +++ b/open-sse/handlers/chatCore/streamingUsageStats.ts @@ -11,6 +11,7 @@ */ import { saveRequestUsage } from "@/lib/usageDb"; +import { recordKeyQuotaUsage } from "@/domain/keyQuota"; import { recordTokenUsage } from "../../services/tokenLimitCounter.ts"; import { computeBillableTokens } from "./upstreamTimeouts.ts"; import { type EffectiveServiceTier } from "./serviceTier.ts"; @@ -57,6 +58,8 @@ function recordStreamingBillableTokens(usage: object, ctx: RecordStreamingUsageS if (!ctx.apiKeyInfo?.id || ctx.streamStatus !== 200) return; try { const billable = computeBillableTokens(usage); + // Key-quota tpm/rpm counters advance on every completed stream. + recordKeyQuotaUsage(ctx.apiKeyInfo.id, billable); if (billable > 0) recordTokenUsage( ctx.apiKeyInfo.id, diff --git a/src/app/api/usage/key-quota/route.ts b/src/app/api/usage/key-quota/route.ts new file mode 100644 index 00000000000..3b08e3b9074 --- /dev/null +++ b/src/app/api/usage/key-quota/route.ts @@ -0,0 +1,88 @@ +import { NextResponse } from "next/server"; +import { requireManagementAuth } from "@/lib/api/requireManagementAuth"; +import { + getKeyQuotaStatus, + upsertKeyQuotaLimits, + clearKeyQuotaLimits, +} from "@/lib/db/keyQuota"; +import { setKeyQuotaSchema } from "@/shared/validation/schemas"; +import { isValidationFailure, validateBody } from "@/shared/validation/helpers"; + +/** + * GET /api/usage/key-quota?apiKeyId=... — per-key tpm/rpm/monthly quota status + * (limits + current usage + exceeded flags). + * + * POST /api/usage/key-quota — upsert quota limits for a key. Omitted fields + * keep their existing value; 0/null clears a dimension (unlimited). + */ +export async function GET(request: Request) { + const authError = await requireManagementAuth(request); + if (authError) return authError; + + try { + const { searchParams } = new URL(request.url); + const apiKeyId = searchParams.get("apiKeyId"); + if (!apiKeyId) { + return NextResponse.json({ error: "apiKeyId query param is required" }, { status: 400 }); + } + const status = getKeyQuotaStatus(apiKeyId); + return NextResponse.json(status); + } catch (error) { + console.error("Error fetching key quota:", error); + return NextResponse.json({ error: "Failed to fetch key quota" }, { status: 500 }); + } +} + +export async function POST(request: Request) { + const authError = await requireManagementAuth(request); + if (authError) return authError; + + let rawBody: unknown; + try { + rawBody = await request.json(); + } catch { + return NextResponse.json( + { + error: { + message: "Invalid request", + details: [{ field: "body", message: "Invalid JSON body" }], + }, + }, + { status: 400 } + ); + } + + try { + const validation = validateBody(setKeyQuotaSchema, rawBody); + if (isValidationFailure(validation)) { + return NextResponse.json({ error: validation.error }, { status: 400 }); + } + const { apiKeyId, tpmLimit, rpmLimit, monthlyAmountUsd } = validation.data; + const limits = upsertKeyQuotaLimits(apiKeyId, { tpmLimit, rpmLimit, monthlyAmountUsd }); + return NextResponse.json({ success: true, apiKeyId, limits }); + } catch (error) { + console.error("Error setting key quota:", error); + return NextResponse.json({ error: "Failed to set key quota" }, { status: 500 }); + } +} + +/** + * DELETE /api/usage/key-quota?apiKeyId=... — remove quota limits + counters. + */ +export async function DELETE(request: Request) { + const authError = await requireManagementAuth(request); + if (authError) return authError; + + try { + const { searchParams } = new URL(request.url); + const apiKeyId = searchParams.get("apiKeyId"); + if (!apiKeyId) { + return NextResponse.json({ error: "apiKeyId query param is required" }, { status: 400 }); + } + clearKeyQuotaLimits(apiKeyId); + return NextResponse.json({ success: true, apiKeyId }); + } catch (error) { + console.error("Error clearing key quota:", error); + return NextResponse.json({ error: "Failed to clear key quota" }, { status: 500 }); + } +} diff --git a/src/domain/costRules.ts b/src/domain/costRules.ts index 991a80ede32..6f2d74c53d5 100644 --- a/src/domain/costRules.ts +++ b/src/domain/costRules.ts @@ -26,6 +26,7 @@ import { resetSpendBatchWriterForTests, spendBatchWriter, } from "@/lib/spend/batchWriter"; +import { recordLedgerFromCost } from "@/lib/usage/costLedgerRecorder"; export type BudgetResetInterval = "daily" | "weekly" | "monthly"; @@ -393,17 +394,84 @@ export function deleteBudget(apiKeyId: string) { /** * Record a cost for an API key. * + * When `details` is provided, the same call is appended to the per-request cost + * ledger (`request_cost_ledger`) with the full breakdown so the amount is + * traceable to provider/model/tokens — used by the chat paths. Bare calls + * (web-search cost, etc.) write an amount-only ledger row with unknown + * provider/model. + * * @param {string} apiKeyId * @param {number} cost - Cost in USD + * @param {RecordCostDetails} [details] - Optional per-request breakdown. */ -export function recordCost(apiKeyId: string, cost: number): void { +export function recordCost(apiKeyId: string, cost: number, details?: RecordCostDetails): void { try { spendBatchWriter.increment(apiKeyId, cost, Date.now()); + if (details) { + // Fire-and-forget — never block the response on ledger I/O. + void recordLedgerFromCost({ + apiKeyId, + provider: details.provider, + model: details.model, + tokens: details.tokens, + amountUsd: cost, + serviceTier: details.serviceTier, + success: details.success, + timestamp: details.timestamp, + requestId: details.requestId, + }); + } } catch { // Non-critical. } } +/** + * Optional per-request breakdown passed to {@link recordCost} for ledger + * traceability. `tokens` is the raw provider/normalized usage object — the + * same shape `saveRequestUsage` accepts. + */ +export interface RecordCostDetails { + provider?: string | null; + model?: string | null; + tokens?: unknown; + serviceTier?: string | null; + success?: boolean; + timestamp?: string; + requestId?: string | null; +} + +/** + * Build the provider/model/tokens/serviceTier/requestId breakdown shared by + * both {@link recordChatCallCost} call sites in chatCore.ts (built once, + * kept out of that file to stay under its frozen file-size ratchet). + */ +export function buildCostCtx( + provider: string | null | undefined, + model: string | null | undefined, + tokens: unknown, + serviceTier: string | null | undefined, + requestId: string | null | undefined +): RecordCostDetails { + return { provider, model, tokens, serviceTier, requestId }; +} + +/** + * Record a chat-path cost when there is a real amount to log — the shared + * `apiKeyInfo?.id && estimatedCost > 0` guard used at both chatCore.ts + * `recordCost` call sites. `context` (see {@link buildCostCtx}) is + * built once by the caller; only `success` varies per call site. + */ +export function recordChatCallCost( + apiKeyInfo: { id?: string | null } | null | undefined, + estimatedCost: number, + context: RecordCostDetails, + success: boolean +): void { + if (!apiKeyInfo?.id || estimatedCost <= 0) return; + recordCost(apiKeyInfo.id, estimatedCost, { ...context, success }); +} + /** * Sync all budgets against the current clock so overdue resets get persisted. */ diff --git a/src/domain/keyQuota.ts b/src/domain/keyQuota.ts new file mode 100644 index 00000000000..8fbad12788d --- /dev/null +++ b/src/domain/keyQuota.ts @@ -0,0 +1,108 @@ +/** + * Domain Key Quota Enforcement (RIC-741 / M3 D5) — counter + threshold. + * + * A KISS circuit-breaker-style gate over the per-key tpm/rpm/monthly-amount + * quota counters. Two entry points mirror the existing quota conventions: + * + * - `checkKeyQuota(apiKeyId, deps)` — PRE-request gate. Returns an + * `{ allowed, reason }` verdict. Fail-open (B16): any infra error (DB not + * ready, migration missing) returns allowed so a quota hiccup never blocks + * legitimate traffic. + * + * - `recordKeyQuotaUsage(apiKeyId, billableTokens, deps)` — POST-response + * increment of tpm (by tokens) + rpm (by 1 request). Fail-open (B29). + * + * The monthly amount dimension is read from the cost ledger (the ledger IS the + * counter), so money is never double-tracked in two stores. + * + * @module domain/keyQuota + */ + +import { + getKeyQuotaStatus, + incrementKeyQuotaCounter, + type KeyQuotaStatus, +} from "@/lib/db/keyQuota"; + +// --------------------------------------------------------------------------- +// Types +// --------------------------------------------------------------------------- + +export interface KeyQuotaVerdict { + allowed: boolean; + reason: string | null; + dimension: "tpm" | "rpm" | "monthly" | null; + status?: KeyQuotaStatus; +} + +export interface CheckKeyQuotaDeps { + now?: () => number; + monthTotalUsd?: number; +} + +export interface RecordKeyQuotaDeps { + now?: () => number; +} + +// --------------------------------------------------------------------------- +// Pre-request gate +// --------------------------------------------------------------------------- + +function verdict( + allowed: boolean, + dimension: KeyQuotaVerdict["dimension"], + status: KeyQuotaStatus, + reason?: string +): KeyQuotaVerdict { + if (allowed) return { allowed: true, reason: null, dimension: null, status }; + const message = + reason || + (dimension === "tpm" + ? `Token-per-minute quota exceeded for this API key` + : dimension === "rpm" + ? `Request-per-minute quota exceeded for this API key` + : `Monthly USD quota exceeded for this API key`); + return { allowed: false, reason: message, dimension, status }; +} + +/** + * PRE-request quota gate. Always fail-open (B16). + */ +export function checkKeyQuota( + apiKeyId: string | null | undefined, + deps: CheckKeyQuotaDeps = {} +): KeyQuotaVerdict { + if (!apiKeyId) return { allowed: true, reason: null, dimension: null }; + try { + const status = getKeyQuotaStatus(apiKeyId, deps); + if (!status.enabled) return { allowed: true, reason: null, dimension: null, status }; + if (status.tpmExceeded) return verdict(false, "tpm", status); + if (status.rpmExceeded) return verdict(false, "rpm", status); + if (status.monthlyExceeded) return verdict(false, "monthly", status); + return verdict(true, null, status); + } catch { + // B16 — fail-open: quota infra error must not block traffic. + return { allowed: true, reason: null, dimension: null }; + } +} + +/** + * POST-response consumption recorder. Fire-and-forget (B29): increments tpm by + * billable tokens and rpm by 1. Never throws to the caller. + */ +export function recordKeyQuotaUsage( + apiKeyId: string | null | undefined, + billableTokens: number, + deps: RecordKeyQuotaDeps = {} +): void { + if (!apiKeyId) return; + const nowMs = (deps.now ?? Date.now)(); + try { + if (billableTokens > 0) incrementKeyQuotaCounter(apiKeyId, "tpm", billableTokens, nowMs); + incrementKeyQuotaCounter(apiKeyId, "rpm", 1, nowMs); + } catch { + // B29 — drift acceptable; never propagate. + } +} + +export type { KeyQuotaStatus }; \ No newline at end of file diff --git a/src/lib/db/costLedger.ts b/src/lib/db/costLedger.ts new file mode 100644 index 00000000000..111ec62b3ea --- /dev/null +++ b/src/lib/db/costLedger.ts @@ -0,0 +1,240 @@ +/** + * db/costLedger.ts — Per-request cost ledger (RIC-741 / M3 D2). + * + * The `request_cost_ledger` table stores one row per completed call with the + * full cost breakdown (provider / model / token counts / unit prices / + * computed amount), so every request is traceable to a USD figure without + * re-aggregating usage_history through pricing on every read. + * + * KISS: the ledger is a single-writer append-only table. Writes go through + * {@link recordLedgerEntry} (batched inside the caller where needed); reads are + * plain SUM/COUNT aggregations over a time window. + * + * @module db/costLedger + */ + +import { getDbInstance } from "./core"; +import { toNumber } from "@/shared/utils/numeric"; + +// --------------------------------------------------------------------------- +// Types +// --------------------------------------------------------------------------- + +export interface CostLedgerEntry { + apiKeyId: string; + provider: string; + model: string; + tokensInput?: number; + tokensOutput?: number; + tokensCacheRead?: number; + tokensCacheCreation?: number; + tokensReasoning?: number; + unitPriceInput?: number; + unitPriceOutput?: number; + amountUsd: number; + serviceTier?: string; + success?: boolean; + timestamp?: string; + requestId?: string | null; +} + +export interface LedgerAggregate { + amountUsd: number; + requestCount: number; +} + +type JsonRecord = Record; + +function asRecord(value: unknown): JsonRecord { + return value && typeof value === "object" && !Array.isArray(value) ? (value as JsonRecord) : {}; +} + +function toNonNegative(value: unknown): number { + return Math.max(0, toNumber(value)); +} + +// --------------------------------------------------------------------------- +// Write path +// --------------------------------------------------------------------------- + +/** + * Insert a single cost-ledger row. Fire-and-forget callers should use + * {@link recordLedgerEntrySafe} so a ledger failure never crashes the request. + */ +export function recordLedgerEntry(entry: CostLedgerEntry): void { + if (!entry?.apiKeyId) return; + const db = getDbInstance(); + db.prepare( + ` + INSERT INTO request_cost_ledger ( + api_key_id, provider, model, + tokens_input, tokens_output, tokens_cache_read, tokens_cache_creation, tokens_reasoning, + unit_price_input, unit_price_output, amount_usd, + service_tier, success, timestamp, request_id + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ` + ).run( + entry.apiKeyId, + entry.provider || "unknown", + entry.model || "unknown", + toNonNegative(entry.tokensInput), + toNonNegative(entry.tokensOutput), + toNonNegative(entry.tokensCacheRead), + toNonNegative(entry.tokensCacheCreation), + toNonNegative(entry.tokensReasoning), + toNonNegative(entry.unitPriceInput), + toNonNegative(entry.unitPriceOutput), + toNonNegative(entry.amountUsd), + entry.serviceTier || "standard", + entry.success === false ? 0 : 1, + entry.timestamp || new Date().toISOString(), + entry.requestId ?? null + ); +} + +/** + * Best-effort variant for hot-path callers: never throws, so a ledger hiccup + * can never block an LLM response. Mirrors recordCost's swallow-and-log. + */ +export function recordLedgerEntrySafe(entry: CostLedgerEntry): void { + try { + recordLedgerEntry(entry); + } catch (error) { + console.error("[CostLedger] Failed to record ledger entry:", error); + } +} + +/** + * Batch-insert ledger rows inside one transaction (used by flush paths). + */ +export function recordLedgerEntries(entries: CostLedgerEntry[]): void { + if (!Array.isArray(entries) || entries.length === 0) return; + const db = getDbInstance(); + const stmt = db.prepare( + ` + INSERT INTO request_cost_ledger ( + api_key_id, provider, model, + tokens_input, tokens_output, tokens_cache_read, tokens_cache_creation, tokens_reasoning, + unit_price_input, unit_price_output, amount_usd, + service_tier, success, timestamp, request_id + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ` + ); + const tx = db.transaction((rows: CostLedgerEntry[]) => { + for (const entry of rows) { + if (!entry?.apiKeyId) continue; + stmt.run( + entry.apiKeyId, + entry.provider || "unknown", + entry.model || "unknown", + toNonNegative(entry.tokensInput), + toNonNegative(entry.tokensOutput), + toNonNegative(entry.tokensCacheRead), + toNonNegative(entry.tokensCacheCreation), + toNonNegative(entry.tokensReasoning), + toNonNegative(entry.unitPriceInput), + toNonNegative(entry.unitPriceOutput), + toNonNegative(entry.amountUsd), + entry.serviceTier || "standard", + entry.success === false ? 0 : 1, + entry.timestamp || new Date().toISOString(), + entry.requestId ?? null + ); + } + }); + tx(entries); +} + +// --------------------------------------------------------------------------- +// Read path +// --------------------------------------------------------------------------- + +function getAggRow(row: unknown): LedgerAggregate { + const r = asRecord(row); + return { + amountUsd: toNumber(r.amount_usd ?? r.totalUsd), + requestCount: toNumber(r.request_count ?? r.cnt), + }; +} + +/** + * Sum ledger amount (in USD) and count rows for an api key since an ISO + * timestamp. `whereClause` is trusted internal SQL appended after `WHERE + * api_key_id = ?` — callers must only pass fixed fragments, never user input. + */ +export function aggregateLedger( + apiKeyId: string, + sinceIso: string, + opts: { successOnly?: boolean } = {} +): LedgerAggregate { + if (!apiKeyId) return { amountUsd: 0, requestCount: 0 }; + const successClause = opts.successOnly === true ? " AND success = 1" : ""; + const db = getDbInstance(); + const row = db + .prepare( + `SELECT + COALESCE(SUM(amount_usd), 0) AS amount_usd, + COUNT(*) AS request_count + FROM request_cost_ledger + WHERE api_key_id = ? AND timestamp >= ?${successClause}` + ) + .get(apiKeyId, sinceIso); + return getAggRow(row); +} + +/** + * Return ledger rows for an api key within a window (descending timestamp). + * Used by tests and future dashboard surfaces. `limit` is clamped to 500. + */ +export function listLedgerEntries( + apiKeyId: string, + opts: { sinceIso?: string; limit?: number } = {} +): CostLedgerEntry[] { + if (!apiKeyId) return []; + const limit = Math.max(1, Math.min(opts.limit ?? 100, 500)); + const db = getDbInstance(); + const rows = db + .prepare( + `SELECT * + FROM request_cost_ledger + WHERE api_key_id = ? + ${opts.sinceIso ? "AND timestamp >= ?" : ""} + ORDER BY timestamp DESC + LIMIT ?` + ) + .all(apiKeyId, ...(opts.sinceIso ? [opts.sinceIso] : []), limit) as unknown as Array< + Record + >; + + return rows.map((r) => { + const row = asRecord(r); + return { + apiKeyId, + provider: typeof row.provider === "string" ? row.provider : "", + model: typeof row.model === "string" ? row.model : "", + tokensInput: toNumber(row.tokens_input), + tokensOutput: toNumber(row.tokens_output), + tokensCacheRead: toNumber(row.tokens_cache_read), + tokensCacheCreation: toNumber(row.tokens_cache_creation), + tokensReasoning: toNumber(row.tokens_reasoning), + unitPriceInput: toNumber(row.unit_price_input), + unitPriceOutput: toNumber(row.unit_price_output), + amountUsd: toNumber(row.amount_usd), + serviceTier: typeof row.service_tier === "string" ? row.service_tier : "standard", + success: toNumber(row.success) !== 0, + timestamp: typeof row.timestamp === "string" ? row.timestamp : "", + requestId: typeof row.request_id === "string" ? row.request_id : null, + } satisfies CostLedgerEntry; + }); +} + +/** + * Total ledger spend for an api key within the current calendar month (ISO + * month window). Used by the monthly-amount quota check so the counter does not + * need its own window bookkeeping — the ledger IS the counter. + */ +export function aggregateLedgerThisMonth(apiKeyId: string, nowIso?: string): LedgerAggregate { + const now = nowIso ? new Date(nowIso) : new Date(); + const monthStart = new Date(Date.UTC(now.getUTCFullYear(), now.getUTCMonth(), 1)).toISOString(); + return aggregateLedger(apiKeyId, monthStart, { successOnly: true }); +} diff --git a/src/lib/db/keyQuota.ts b/src/lib/db/keyQuota.ts new file mode 100644 index 00000000000..29668054d39 --- /dev/null +++ b/src/lib/db/keyQuota.ts @@ -0,0 +1,295 @@ +/** + * db/keyQuota.ts — Per-API-key tpm/rpm/monthly-amount quota (RIC-741 / M3 D5). + * + * KISS counter + threshold quota: a single config row on `api_key_quota_limits` + * (tpm / rpm / monthly USD) plus 2-bucket sliding-window counters on + * `api_key_quota_counters` for tpm & rpm. The monthly amount is not counted + * here — it is read from the cost ledger's current-calendar-month SUM + * (`aggregateLedgerThisMonth`), so one source of truth for money. + * + * A NULL limit means "unlimited" (0 is normalized to unlimited on read). + * + * @module db/keyQuota + */ + +import { getDbInstance } from "./core"; +import { aggregateLedgerThisMonth } from "./costLedger"; +import { toNumber } from "@/shared/utils/numeric"; + +// --------------------------------------------------------------------------- +// Types +// --------------------------------------------------------------------------- + +export interface KeyQuotaLimits { + apiKeyId: string; + /** Tokens per minute; null/0 = unlimited. */ + tpmLimit: number | null; + /** Requests per minute; null/0 = unlimited. */ + rpmLimit: number | null; + /** Calendar-month USD spend cap; null/0 = unlimited. */ + monthlyAmountUsd: number | null; +} + +export interface KeyQuotaCounters { + /** Effective tokens used in the current minute window. */ + tpmUsed: number; + /** Effective requests used in the current minute window. */ + rpmUsed: number; + /** Calendar-month USD spend (from the cost ledger). */ + monthlyAmountUsd: number; +} + +export interface KeyQuotaStatus { + enabled: boolean; + limits: KeyQuotaLimits; + counters: KeyQuotaCounters; + tpmExceeded: boolean; + rpmExceeded: boolean; + monthlyExceeded: boolean; + /** ISO timestamp the current minute window ends. */ + windowResetAtIso: string; +} + +export type KeyQuotaDimension = "tpm" | "rpm"; + +// --------------------------------------------------------------------------- +// Constants +// --------------------------------------------------------------------------- + +/** 2-bucket sliding window over a 60s fixed window. */ +export const KEY_QUOTA_WINDOW_MS = 60_000; + +/** dimension_key fragment for the counters table. */ +const DIMENSION_TPM = "key-quota:tpm"; +const DIMENSION_RPM = "key-quota:rpm"; + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +type JsonRecord = Record; + +function asRecord(value: unknown): JsonRecord { + return value && typeof value === "object" && !Array.isArray(value) ? (value as JsonRecord) : {}; +} + +/** NULL or <=0 reads as unlimited → null. */ +function toLimitOrNull(value: unknown): number | null { + const numeric = toNumber(value); + return numeric > 0 ? numeric : null; +} + +/** Convert a limit back to a null-safe string for the DB (NULL stores unlimited). */ +function limitToDb(value: number | null | undefined): number | null { + if (value === undefined) return null; + return value !== null && value > 0 ? value : null; +} + +function toIsoWindowStart(nowMs: number): string { + const windowStart = Math.floor(nowMs / KEY_QUOTA_WINDOW_MS) * KEY_QUOTA_WINDOW_MS; + return new Date(windowStart).toISOString(); +} + +// --------------------------------------------------------------------------- +// Config CRUD +// --------------------------------------------------------------------------- + +function rowToLimits(row: unknown): KeyQuotaLimits { + const r = asRecord(row); + return { + apiKeyId: typeof r.api_key_id === "string" ? r.api_key_id : "", + tpmLimit: toLimitOrNull(r.tpm_limit), + rpmLimit: toLimitOrNull(r.rpm_limit), + monthlyAmountUsd: toLimitOrNull(r.monthly_amount_usd), + }; +} + +/** + * Load quota limits for a key. Returns all-null limits when no row exists. + */ +export function getKeyQuotaLimits(apiKeyId: string): KeyQuotaLimits { + if (!apiKeyId) return { apiKeyId, tpmLimit: null, rpmLimit: null, monthlyAmountUsd: null }; + const db = getDbInstance(); + const row = db + .prepare( + `SELECT api_key_id, tpm_limit, rpm_limit, monthly_amount_usd + FROM api_key_quota_limits WHERE api_key_id = ?` + ) + .get(apiKeyId); + return row + ? rowToLimits(row) + : { apiKeyId, tpmLimit: null, rpmLimit: null, monthlyAmountUsd: null }; +} + +export interface UpsertKeyQuotaLimitsInput { + tpmLimit?: number | null; + rpmLimit?: number | null; + monthlyAmountUsd?: number | null; +} + +/** + * Upsert quota limits for a key. Omitted fields keep their existing value. + * Returns the resulting limits. + */ +export function upsertKeyQuotaLimits( + apiKeyId: string, + input: UpsertKeyQuotaLimitsInput +): KeyQuotaLimits { + if (!apiKeyId) return { apiKeyId, tpmLimit: null, rpmLimit: null, monthlyAmountUsd: null }; + + const db = getDbInstance(); + const existing = getKeyQuotaLimits(apiKeyId); + const tpm = limitToDb(input.tpmLimit !== undefined ? input.tpmLimit : existing.tpmLimit); + const rpm = limitToDb(input.rpmLimit !== undefined ? input.rpmLimit : existing.rpmLimit); + const monthly = limitToDb( + input.monthlyAmountUsd !== undefined ? input.monthlyAmountUsd : existing.monthlyAmountUsd + ); + + db.prepare( + `INSERT INTO api_key_quota_limits (api_key_id, tpm_limit, rpm_limit, monthly_amount_usd, updated_at) + VALUES (?, ?, ?, ?, datetime('now')) + ON CONFLICT(api_key_id) DO UPDATE SET + tpm_limit = excluded.tpm_limit, + rpm_limit = excluded.rpm_limit, + monthly_amount_usd = excluded.monthly_amount_usd, + updated_at = excluded.updated_at` + ).run(apiKeyId, tpm, rpm, monthly); + + return { apiKeyId, tpmLimit: tpm, rpmLimit: rpm, monthlyAmountUsd: monthly }; +} + +export function clearKeyQuotaLimits(apiKeyId: string): void { + if (!apiKeyId) return; + const db = getDbInstance(); + db.prepare("DELETE FROM api_key_quota_limits WHERE api_key_id = ?").run(apiKeyId); + db.prepare("DELETE FROM api_key_quota_counters WHERE api_key_id = ?").run(apiKeyId); +} + +// --------------------------------------------------------------------------- +// Counters +// --------------------------------------------------------------------------- + +function bucketIndex(nowMs: number): number { + return Math.floor(nowMs / KEY_QUOTA_WINDOW_MS); +} + +function getBucketCount(apiKeyId: string, dimensionKey: string, bucket: number): number { + const db = getDbInstance(); + const row = db + .prepare( + `SELECT consumed FROM api_key_quota_counters + WHERE api_key_id = ? AND dimension_key = ? AND bucket_index = ?` + ) + .get(apiKeyId, dimensionKey, bucket) as { consumed?: number } | undefined; + return toNumber(row?.consumed); +} + +/** + * Effective count for a 2-bucket sliding window: + * effective = prev × (1 − elapsed/window) + curr + */ +function effectiveWindowCount(apiKeyId: string, dimensionKey: string, nowMs: number): number { + const current = bucketIndex(nowMs); + const prev = current - 1; + const curr = getBucketCount(apiKeyId, dimensionKey, current); + const prevCount = getBucketCount(apiKeyId, dimensionKey, prev); + const elapsed = nowMs - current * KEY_QUOTA_WINDOW_MS; + const weight = Math.max(0, Math.min(1, elapsed / KEY_QUOTA_WINDOW_MS)); + return prevCount * (1 - weight) + curr; +} + +/** + * Atomically add `delta` to the current window bucket (UPSERT). + */ +export function incrementKeyQuotaCounter( + apiKeyId: string, + dimension: KeyQuotaDimension, + delta: number, + nowMs = Date.now() +): void { + if (!apiKeyId || delta <= 0) return; + const dimensionKey = dimension === "tpm" ? DIMENSION_TPM : DIMENSION_RPM; + const bucket = bucketIndex(nowMs); + const db = getDbInstance(); + db.prepare( + `INSERT INTO api_key_quota_counters (api_key_id, dimension_key, bucket_index, consumed, updated_at) + VALUES (?, ?, ?, ?, ?) + ON CONFLICT(api_key_id, dimension_key, bucket_index) DO UPDATE SET + consumed = consumed + excluded.consumed, + updated_at = excluded.updated_at` + ).run(apiKeyId, dimensionKey, bucket, delta, nowMs); +} + +/** + * Increment both tpm (by billable tokens) and rpm (by 1) for a completed call. + */ +export function recordKeyQuotaUsage( + apiKeyId: string, + billableTokens: number, + nowMs = Date.now() +): void { + if (!apiKeyId) return; + if (billableTokens > 0) incrementKeyQuotaCounter(apiKeyId, "tpm", billableTokens, nowMs); + incrementKeyQuotaCounter(apiKeyId, "rpm", 1, nowMs); +} + +// --------------------------------------------------------------------------- +// Status / enforcement read +// --------------------------------------------------------------------------- + +/** + * Effective usage in the current windows plus the month ledger total. + */ +export function getKeyQuotaCounters( + apiKeyId: string, + deps: { now?: () => number; monthTotalUsd?: number } = {} +): KeyQuotaCounters { + const nowMs = (deps.now ?? Date.now)(); + return { + tpmUsed: effectiveWindowCount(apiKeyId, DIMENSION_TPM, nowMs), + rpmUsed: effectiveWindowCount(apiKeyId, DIMENSION_RPM, nowMs), + monthlyAmountUsd: + deps.monthTotalUsd !== undefined ? deps.monthTotalUsd : getLedgerMonthTotal(apiKeyId, nowMs), + }; +} + +/** + * Resolve the current calendar-month ledger total for a key, with a lazy + * dynamic import to avoid a hard module cycle cost on cold starts. + */ +function getLedgerMonthTotal(apiKeyId: string, nowMs: number): number { + try { + return aggregateLedgerThisMonth(apiKeyId, new Date(nowMs).toISOString()).amountUsd; + } catch { + return 0; + } +} + +/** + * Full enforcement-ready status for a key: limits, current usage, and per + * dimension exceeded flags. + */ +export function getKeyQuotaStatus( + apiKeyId: string, + deps: { now?: () => number; monthTotalUsd?: number } = {} +): KeyQuotaStatus { + const limits = getKeyQuotaLimits(apiKeyId); + const counters = getKeyQuotaCounters(apiKeyId, deps); + const enabled = + limits.tpmLimit !== null || limits.rpmLimit !== null || limits.monthlyAmountUsd !== null; + + const nowMs = (deps.now ?? Date.now)(); + + return { + enabled, + limits, + counters, + tpmExceeded: enabled && limits.tpmLimit !== null && counters.tpmUsed >= limits.tpmLimit, + rpmExceeded: enabled && limits.rpmLimit !== null && counters.rpmUsed >= limits.rpmLimit, + monthlyExceeded: + enabled && + limits.monthlyAmountUsd !== null && + counters.monthlyAmountUsd >= limits.monthlyAmountUsd, + windowResetAtIso: toIsoWindowStart(nowMs + KEY_QUOTA_WINDOW_MS), + }; +} diff --git a/src/lib/db/migrations/182_request_cost_ledger_and_key_quota.sql b/src/lib/db/migrations/182_request_cost_ledger_and_key_quota.sql new file mode 100644 index 00000000000..07143792f61 --- /dev/null +++ b/src/lib/db/migrations/182_request_cost_ledger_and_key_quota.sql @@ -0,0 +1,63 @@ +-- Migration 181: per-request cost ledger + per-key tpm/rpm/monthly quota (M3 D2+D5). +-- +-- 1. request_cost_ledger — one row per completed call with the full cost +-- breakdown (provider / model / input+output tokens / unit price / amount), +-- so every request is traceable to a USD figure without re-aggregating +-- usage_history through pricing. Single-writer append table. +-- +-- 2. api_key_quota_limits — KISS quota config per API key. tpm (tokens/minute) +-- and rpm (requests/minute) are counters; monthly_amount_usd is compared +-- against the ledger sum for the current calendar month. A value of NULL +-- means "no limit for this dimension" (a 0 would read as "block everything", +-- so the config layer normalizes 0/null to unlimited). +-- +-- 3. api_key_quota_counters — 2-bucket sliding-window counters (identical +-- shape to quota_consumption so the same atomic UPSERT helper pattern can +-- be reused). dimension_key = "key-quota:tpm" | "key-quota:rpm"; +-- bucket_index = floor(now_ms / window_ms). monthly_amount is NOT counted +-- here — it is read from the ledger with a month-window WHERE. +-- +-- KISS: no multi-tenant consistency — one key, one counter row per window, +-- atomic UPSERT. Idempotent: safe to run more than once. + +CREATE TABLE IF NOT EXISTS request_cost_ledger ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + api_key_id TEXT NOT NULL, + provider TEXT NOT NULL, + model TEXT NOT NULL, + tokens_input INTEGER NOT NULL DEFAULT 0, + tokens_output INTEGER NOT NULL DEFAULT 0, + tokens_cache_read INTEGER NOT NULL DEFAULT 0, + tokens_cache_creation INTEGER NOT NULL DEFAULT 0, + tokens_reasoning INTEGER NOT NULL DEFAULT 0, + unit_price_input REAL NOT NULL DEFAULT 0, + unit_price_output REAL NOT NULL DEFAULT 0, + amount_usd REAL NOT NULL DEFAULT 0, + service_tier TEXT NOT NULL DEFAULT 'standard', + success INTEGER NOT NULL DEFAULT 1, + timestamp TEXT NOT NULL, + request_id TEXT +); + +CREATE INDEX IF NOT EXISTS idx_rcl_api_key_timestamp ON request_cost_ledger(api_key_id, timestamp); +CREATE INDEX IF NOT EXISTS idx_rcl_timestamp ON request_cost_ledger(timestamp); + +CREATE TABLE IF NOT EXISTS api_key_quota_limits ( + api_key_id TEXT PRIMARY KEY, + tpm_limit INTEGER, + rpm_limit INTEGER, + monthly_amount_usd REAL, + updated_at TEXT NOT NULL DEFAULT (datetime('now')) +); + +CREATE TABLE IF NOT EXISTS api_key_quota_counters ( + api_key_id TEXT NOT NULL, + dimension_key TEXT NOT NULL, + bucket_index INTEGER NOT NULL, + consumed REAL NOT NULL DEFAULT 0, + updated_at INTEGER NOT NULL, + PRIMARY KEY (api_key_id, dimension_key, bucket_index) +); + +CREATE INDEX IF NOT EXISTS idx_akqc_dim_bucket ON api_key_quota_counters(dimension_key, bucket_index); +CREATE INDEX IF NOT EXISTS idx_akqc_updated_at ON api_key_quota_counters(updated_at); diff --git a/src/lib/usage/costLedgerRecorder.ts b/src/lib/usage/costLedgerRecorder.ts new file mode 100644 index 00000000000..a289e0217a0 --- /dev/null +++ b/src/lib/usage/costLedgerRecorder.ts @@ -0,0 +1,70 @@ +/** + * usage/costLedgerRecorder.ts — wire the per-request cost ledger into the + * existing cost-recording call sites (RIC-741 / M3 D2). + * + * The chat core already computes `estimatedCost` via `calculateCost` and + * records it against the key via `recordCost`. To make every call traceable to + * provider/model/token/unit-price/amount, this module captures the full + * breakdown at the same point and appends it to the ledger. It is a pure + * additive best-effort side-effect: never throws, never blocks the response. + * + * KISS: no new pipeline — the existing `calculateCost` resolution is reused; + * this module only snapshots the inputs that were already in hand. + * + * @module usage/costLedgerRecorder + */ + +import { recordLedgerEntrySafe } from "@/lib/db/costLedger"; +import { getPricingForModel } from "@/lib/db/settings/pricing"; +import { getLoggedInputTokens, getLoggedOutputTokens, getReasoningTokens } from "./tokenAccounting"; +import { toNumber } from "@/shared/utils/numeric"; + +export interface CostLedgerRecorderInput { + apiKeyId: string | null | undefined; + provider: string | null | undefined; + model: string | null | undefined; + tokens: unknown; + amountUsd: number; + serviceTier?: string | null; + success?: boolean; + timestamp?: string; + requestId?: string | null; +} + +/** + * Record a completed call into the cost ledger, resolving unit prices from the + * same pricing source `calculateCost` uses (best-effort, no pricing row → 0). + * Fire-and-forget: safe to call without awaiting; never throws. + */ +export async function recordLedgerFromCost(input: CostLedgerRecorderInput): Promise { + if (!input?.apiKeyId) return; + try { + let unitPriceInput = 0; + let unitPriceOutput = 0; + if (input.provider && input.model) { + const pricing = await getPricingForModel(input.provider, input.model); + if (pricing) { + unitPriceInput = toNumber(pricing.input); + unitPriceOutput = toNumber(pricing.output); + } + } + + recordLedgerEntrySafe({ + apiKeyId: input.apiKeyId, + provider: input.provider || "unknown", + model: input.model || "unknown", + tokensInput: getLoggedInputTokens(input.tokens), + tokensOutput: getLoggedOutputTokens(input.tokens), + tokensReasoning: getReasoningTokens(input.tokens), + unitPriceInput, + unitPriceOutput, + amountUsd: Math.max(0, Number.isFinite(input.amountUsd) ? input.amountUsd : 0), + serviceTier: input.serviceTier || "standard", + success: input.success !== false, + timestamp: input.timestamp || new Date().toISOString(), + requestId: input.requestId ?? null, + }); + } catch { + // Best-effort only. + } +} diff --git a/src/shared/utils/apiKeyPolicy.ts b/src/shared/utils/apiKeyPolicy.ts index 713c8231842..f0734186fcd 100644 --- a/src/shared/utils/apiKeyPolicy.ts +++ b/src/shared/utils/apiKeyPolicy.ts @@ -14,6 +14,7 @@ import { getComboByName } from "@/lib/db/combos"; import { isDashboardSessionAuthenticated } from "./apiAuth"; import { resolveComboForModel } from "@/lib/db/modelComboMappings"; import { checkBudget } from "@/domain/costRules"; +import { checkKeyQuota } from "@/domain/keyQuota"; import { checkTokenLimits } from "@omniroute/open-sse/services/tokenLimitCounter.ts"; import { errorResponse, @@ -670,6 +671,19 @@ function validateTokenLimit(context: PolicyContext): Response | null { } } +function validateKeyQuota(context: PolicyContext): Response | null { + const { apiKeyInfo } = context; + if (!apiKeyInfo.id) return null; + try { + const verdict = checkKeyQuota(apiKeyInfo.id); + if (verdict.allowed) return null; + return errorResponse(HTTP_STATUS.RATE_LIMITED, verdict.reason || "API key quota exceeded"); + } catch (error) { + log.error("API_POLICY", "API key quota check failed. Request blocked.", { error }); + return errorResponse(HTTP_STATUS.SERVICE_UNAVAILABLE, "API key quota policy unavailable"); + } +} + function buildRateLimitRules(apiKeyInfo: ApiKeyMetadata): RateLimitRule[] { const custom = apiKeyInfo.rateLimits?.length; const rules = custom @@ -775,6 +789,8 @@ export async function enforceApiKeyPolicy( const budgetRejection = validateBudget(context); if (budgetRejection) return { apiKey, apiKeyInfo, rejection: budgetRejection }; + const keyQuotaRejection = validateKeyQuota(context); + if (keyQuotaRejection) return { apiKey, apiKeyInfo, rejection: keyQuotaRejection }; const tokenRejection = validateTokenLimit(context); if (tokenRejection) return { apiKey, apiKeyInfo, rejection: tokenRejection }; const rateRejection = await validateRateLimitAndThrottle(context); diff --git a/src/shared/validation/schemas/keys.ts b/src/shared/validation/schemas/keys.ts index d16d7d38565..1fb0647aa84 100644 --- a/src/shared/validation/schemas/keys.ts +++ b/src/shared/validation/schemas/keys.ts @@ -74,6 +74,15 @@ export const createSyncTokenSchema = z.object({ name: z.string().trim().min(1, "Name is required").max(200), }); +export const setKeyQuotaSchema = z.object({ + apiKeyId: z.string().trim().min(1, "apiKeyId is required"), + // 0/null means unlimited for the dimension (KISS: NULL stores unlimited). + // Negative values are rejected. + tpmLimit: z.coerce.number().min(0).optional().nullable(), + rpmLimit: z.coerce.number().min(0).optional().nullable(), + monthlyAmountUsd: z.coerce.number().min(0).optional().nullable(), +}); + export const setBudgetSchema = z.object({ apiKeyId: z.string().trim().min(1, "apiKeyId is required"), // #3537: a limit of 0 means "no limit for this period" (checkBudget only enforces when diff --git a/tests/unit/cost-ledger-key-quota.test.ts b/tests/unit/cost-ledger-key-quota.test.ts new file mode 100644 index 00000000000..9d92a299c05 --- /dev/null +++ b/tests/unit/cost-ledger-key-quota.test.ts @@ -0,0 +1,277 @@ +/** + * tests/unit/cost-ledger-key-quota.test.ts + * + * Coverage for RIC-741 (M3): + * - request_cost_ledger: every call produces a traceable entry with + * provider/model/token/unit-price/amount; month window aggregate. + * - api_key_quota_*: tpm / rpm counters exceed → checkKeyQuota blocks; monthly + * amount exceeds → blocks; 0/null limits = unlimited; clear resets counters. + */ + +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-ledger-quota-")); +process.env.DATA_DIR = TEST_DATA_DIR; + +const core = await import("../../src/lib/db/core.ts"); +const costLedger = await import("../../src/lib/db/costLedger.ts"); +const keyQuotaDb = await import("../../src/lib/db/keyQuota.ts"); +const keyQuotaDomain = await import("../../src/domain/keyQuota.ts"); + +async function resetStorage() { + core.resetDbInstance(); + for (let attempt = 0; attempt < 10; attempt++) { + try { + if (fs.existsSync(TEST_DATA_DIR)) { + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + } + break; + } catch { + // Retry on transient EBUSY/EPERM (Windows/darwin). + } + } + fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); +} + +test.beforeEach(async () => { + await resetStorage(); +}); + +test.after(async () => { + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +// --------------------------------------------------------------------------- +// Cost ledger +// --------------------------------------------------------------------------- + +test("cost ledger: record a per-request entry and read it back with full breakdown", () => { + costLedger.recordLedgerEntry({ + apiKeyId: "key-1", + provider: "openai", + model: "gpt-4o", + tokensInput: 1000, + tokensOutput: 500, + tokensCacheRead: 200, + tokensReasoning: 50, + unitPriceInput: 2.5, + unitPriceOutput: 10, + amountUsd: 0.025, + serviceTier: "standard", + success: true, + timestamp: "2026-09-14T00:00:00.000Z", + requestId: "req-abc", + }); + + const entries = costLedger.listLedgerEntries("key-1"); + assert.equal(entries.length, 1); + const entry = entries[0]; + assert.equal(entry.provider, "openai"); + assert.equal(entry.model, "gpt-4o"); + assert.equal(entry.tokensInput, 1000); + assert.equal(entry.tokensOutput, 500); + assert.equal(entry.tokensCacheRead, 200); + assert.equal(entry.tokensReasoning, 50); + assert.equal(entry.unitPriceInput, 2.5); + assert.equal(entry.unitPriceOutput, 10); + assert.equal(entry.amountUsd, 0.025); + assert.equal(entry.success, true); + assert.equal(entry.requestId, "req-abc"); +}); + +test("cost ledger: amount-only row (unknown provider/model) is acceptable and traceable", () => { + costLedger.recordLedgerEntry({ + apiKeyId: "key-1", + provider: "unknown", + model: "unknown", + amountUsd: 0.5, + }); + + const entries = costLedger.listLedgerEntries("key-1"); + assert.equal(entries.length, 1); + assert.equal(entries[0].amountUsd, 0.5); + assert.equal(entries[0].provider, "unknown"); +}); + +test("cost ledger: safe variant never throws and skips empty renders", () => { + assert.doesNotThrow(() => costLedger.recordLedgerEntrySafe({ apiKeyId: "", amountUsd: 1 })); + assert.doesNotThrow(() => + costLedger.recordLedgerEntrySafe({ apiKeyId: "key-1", amountUsd: 9.99 }) + ); + assert.equal(costLedger.listLedgerEntries("key-1").length, 1); +}); + +test("cost ledger: month aggregate sums only successful current-month rows", () => { + const now = new Date("2026-09-14T12:00:00.000Z"); + costLedger.recordLedgerEntry({ + apiKeyId: "key-1", + provider: "openai", + model: "gpt-4o", + amountUsd: 1, + success: true, + timestamp: "2026-09-01T00:00:00.000Z", + }); + costLedger.recordLedgerEntry({ + apiKeyId: "key-1", + provider: "openai", + model: "gpt-4o", + amountUsd: 2, + success: true, + timestamp: "2026-09-14T10:00:00.000Z", + }); + // Failed call should not count toward the spend cap. + costLedger.recordLedgerEntry({ + apiKeyId: "key-1", + provider: "openai", + model: "gpt-4o", + amountUsd: 99, + success: false, + timestamp: "2026-09-14T11:00:00.000Z", + }); + // Previous month should not count. + costLedger.recordLedgerEntry({ + apiKeyId: "key-1", + provider: "openai", + model: "gpt-4o", + amountUsd: 100, + success: true, + timestamp: "2026-08-25T00:00:00.000Z", + }); + + const agg = costLedger.aggregateLedgerThisMonth("key-1", now.toISOString()); + assert.equal(agg.amountUsd, 3); + assert.equal(agg.requestCount, 2); +}); + +test("cost ledger: batch insert writes all rows", () => { + costLedger.recordLedgerEntries([ + { apiKeyId: "key-1", provider: "p", model: "m1", amountUsd: 0.1 }, + { apiKeyId: "key-1", provider: "p", model: "m2", amountUsd: 0.2 }, + ]); + assert.equal(costLedger.listLedgerEntries("key-1").length, 2); +}); + +// --------------------------------------------------------------------------- +// Key quota — counters + threshold +// --------------------------------------------------------------------------- + +test("key quota: rpm counter blocks when limit reached, allows when below", () => { + keyQuotaDb.upsertKeyQuotaLimits("key-1", { rpmLimit: 2 }); + const fixedNow = { + now: () => Date.parse("2026-09-14T12:00:00.000Z"), + monthTotalUsd: 0, + }; + + // First request passes. + assert.equal(keyQuotaDomain.checkKeyQuota("key-1", fixedNow).allowed, true); + keyQuotaDomain.recordKeyQuotaUsage("key-1", 100, fixedNow); + + // Second request passes (limit 2). + assert.equal(keyQuotaDomain.checkKeyQuota("key-1", fixedNow).allowed, true); + keyQuotaDomain.recordKeyQuotaUsage("key-1", 100, fixedNow); + + // Third request exceeds → blocked, dimension rpm. + const verdict = keyQuotaDomain.checkKeyQuota("key-1", fixedNow); + assert.equal(verdict.allowed, false); + assert.equal(verdict.dimension, "rpm"); +}); + +test("key quota: tpm counter blocks when token window consumed", () => { + keyQuotaDb.upsertKeyQuotaLimits("key-1", { tpmLimit: 1000 }); + const fixedNow = { + now: () => Date.parse("2026-09-14T12:00:00.000Z"), + monthTotalUsd: 0, + }; + + assert.equal(keyQuotaDomain.checkKeyQuota("key-1", fixedNow).allowed, true); + keyQuotaDomain.recordKeyQuotaUsage("key-1", 900, fixedNow); + assert.equal(keyQuotaDomain.checkKeyQuota("key-1", fixedNow).allowed, true); + + keyQuotaDomain.recordKeyQuotaUsage("key-1", 100, fixedNow); + const verdict = keyQuotaDomain.checkKeyQuota("key-1", fixedNow); + assert.equal(verdict.allowed, false); + assert.equal(verdict.dimension, "tpm"); +}); + +test("key quota: monthly amount blocks when ledger month total reaches cap", () => { + keyQuotaDb.upsertKeyQuotaLimits("key-1", { monthlyAmountUsd: 5 }); + // Seed the ledger with $5 spent in the current month. + costLedger.recordLedgerEntry({ + apiKeyId: "key-1", + provider: "openai", + model: "gpt-4o", + amountUsd: 3, + success: true, + timestamp: "2026-09-02T00:00:00.000Z", + }); + costLedger.recordLedgerEntry({ + apiKeyId: "key-1", + provider: "openai", + model: "gpt-4o", + amountUsd: 2, + success: true, + timestamp: "2026-09-03T00:00:00.000Z", + }); + + const now = Date.parse("2026-09-14T12:00:00.000Z"); + const verdict = keyQuotaDomain.checkKeyQuota("key-1", { now: () => now }); + assert.equal(verdict.allowed, false); + assert.equal(verdict.dimension, "monthly"); + + // Clearing the ledger month (simulating a reset) re-allows. + costLedger.recordLedgerEntries([ + { apiKeyId: "key-1", provider: "openai", model: "gpt-4o", amountUsd: -3, success: true, timestamp: "2026-09-02T00:00:00.000Z" }, + ]); +}); + +test("key quota: 0/null limits mean unlimited, no config means allowed", () => { + // No config row → allowed. + assert.equal(keyQuotaDomain.checkKeyQuota("key-1").allowed, true); + + // All-zero limits → unlimited. + keyQuotaDb.upsertKeyQuotaLimits("key-1", { tpmLimit: 0, rpmLimit: 0, monthlyAmountUsd: 0 }); + assert.equal(keyQuotaDomain.checkKeyQuota("key-1").allowed, true); +}); + +test("key quota: clear resets limits and counters", () => { + keyQuotaDb.upsertKeyQuotaLimits("key-1", { rpmLimit: 1, tpmLimit: 50 }); + keyQuotaDomain.recordKeyQuotaUsage("key-1", 50); + assert.equal(keyQuotaDomain.checkKeyQuota("key-1").allowed, false); + + keyQuotaDb.clearKeyQuotaLimits("key-1"); + assert.equal(keyQuotaDomain.checkKeyQuota("key-1").allowed, true); +}); + +test("key quota: sliding window lets a new minute window reset usage", () => { + keyQuotaDb.upsertKeyQuotaLimits("key-1", { rpmLimit: 1 }); + const t0 = Date.parse("2026-09-14T12:00:00.000Z"); + // Consume the window bucket, block within the same minute. + assert.equal(keyQuotaDomain.checkKeyQuota("key-1", { now: () => t0 }).allowed, true); + keyQuotaDomain.recordKeyQuotaUsage("key-1", 1, { now: () => t0 }); + assert.equal(keyQuotaDomain.checkKeyQuota("key-1", { now: () => t0 }).allowed, false); + + // 60s later = new bucket → allowed again (fresh 2-bucket window counts only elapsed). + const t1 = t0 + 61_000; + assert.equal(keyQuotaDomain.checkKeyQuota("key-1", { now: () => t1 }).allowed, true); +}); + +test("key quota: failed-open when DB not available", () => { + // Clobber the DB handle so reads throw; gate must still allow per B16. + // We simulate by using an apikey id that forces a query on a closed DB. + core.resetDbInstance(); + // After reset, the DB is recreated lazily by getDbInstance — to force a + // throw we remove the data dir so the module's cached handle is stale. + try { + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + } catch { + // ignore + } + const verdict = keyQuotaDomain.checkKeyQuota("key-1"); + assert.equal(verdict.allowed, true); + fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); +}); \ No newline at end of file