From ec42cde3fb6f9f6fb3fb52cb8254121898a5d9bc Mon Sep 17 00:00:00 2001 From: mahdiwafy Date: Wed, 29 Jul 2026 22:24:18 +0700 Subject: [PATCH] feat(keys): add granular API key limits --- open-sse/handlers/chatCore.js | 4 +- .../handlers/chatCore/nonStreamingHandler.js | 4 +- open-sse/handlers/chatCore/requestDetail.js | 9 +- .../handlers/chatCore/sseToJsonHandler.js | 6 +- .../handlers/chatCore/streamingHandler.js | 11 +- .../dashboard/endpoint/EndpointPageClient.js | 96 +++++- src/app/api/keys/[id]/route.js | 20 +- src/app/api/keys/route.js | 23 +- src/app/api/v1/models/route.js | 10 + .../db/migrations/005-add-api-key-limits.js | 24 ++ src/lib/db/migrations/index.js | 5 +- src/lib/db/repos/apiKeyUsageRepo.js | 217 +++++++++++++ src/lib/db/repos/apiKeysRepo.js | 286 +++++++++++------- src/lib/db/schema.js | 11 +- src/lib/localDb.js | 5 + src/sse/handlers/chat.js | 25 +- src/sse/services/auth.js | 28 +- 17 files changed, 643 insertions(+), 141 deletions(-) create mode 100644 src/lib/db/migrations/005-add-api-key-limits.js create mode 100644 src/lib/db/repos/apiKeyUsageRepo.js diff --git a/open-sse/handlers/chatCore.js b/open-sse/handlers/chatCore.js index 681a0b0590..2992a5632f 100644 --- a/open-sse/handlers/chatCore.js +++ b/open-sse/handlers/chatCore.js @@ -101,7 +101,7 @@ export function applyLoopGuard(translatedBody, finalFormat, provider, model, log * @param {object} options.credentials - Provider credentials * @param {string} options.sourceFormatOverride - Override detected source format (e.g. "openai-responses") */ -export async function handleChatCore({ body, modelInfo, credentials, log, onCredentialsRefreshed, onRequestSuccess, onDisconnect, clientRawRequest, connectionId, userAgent, apiKey, apiKeyName = null, ccFilterNaming, rtkEnabled, headroomEnabled, headroomUrl, headroomCompressUserMessages, cavemanEnabled, cavemanLevel, ponytailEnabled, ponytailLevel, pxpipeEnabled = false, pxpipeMinChars = 1000, pxpipeTimeoutMs = 10000, pxpipeTransform = "png", onPxpipeEvent = null, sourceFormatOverride, providerThinking, clientSignal, loopGuardEnabled = true }) { +export async function handleChatCore({ body, modelInfo, credentials, log, onCredentialsRefreshed, onRequestSuccess, onDisconnect, clientRawRequest, connectionId, userAgent, apiKey, apiKeyInfo = null, apiKeyName = null, ccFilterNaming, rtkEnabled, headroomEnabled, headroomUrl, headroomCompressUserMessages, cavemanEnabled, cavemanLevel, ponytailEnabled, ponytailLevel, pxpipeEnabled = false, pxpipeMinChars = 1000, pxpipeTimeoutMs = 10000, pxpipeTransform = "png", onPxpipeEvent = null, sourceFormatOverride, providerThinking, clientSignal, loopGuardEnabled = true }) { const { provider, model, accountCount = 0 } = modelInfo; const requestStartTime = Date.now(); @@ -490,7 +490,7 @@ export async function handleChatCore({ body, modelInfo, credentials, log, onCred return createErrorResult(statusCode, errMsg, resetsAtMs); } - const sharedCtx = { provider, model, body, stream, translatedBody, finalBody, requestStartTime, connectionId, apiKey, apiKeyName, clientRawRequest, onRequestSuccess, pxpipe: pxpipeSummary }; + const sharedCtx = { provider, model, body, stream, translatedBody, finalBody, requestStartTime, connectionId, apiKey, apiKeyInfo, apiKeyName, clientRawRequest, onRequestSuccess, pxpipe: pxpipeSummary }; const appendLog = (extra) => appendRequestLog({ model, provider, connectionId, ...extra }).catch(() => { }); const trackDone = () => trackPendingRequest(model, provider, connectionId, false); diff --git a/open-sse/handlers/chatCore/nonStreamingHandler.js b/open-sse/handlers/chatCore/nonStreamingHandler.js index 484b5951f6..c54ed34ca6 100644 --- a/open-sse/handlers/chatCore/nonStreamingHandler.js +++ b/open-sse/handlers/chatCore/nonStreamingHandler.js @@ -217,7 +217,7 @@ export function translateNonStreamingResponse(responseBody, targetFormat, source /** * Handle non-streaming response from provider. */ -export async function handleNonStreamingResponse({ providerResponse, provider, model, sourceFormat, targetFormat, body, stream, translatedBody, finalBody, requestStartTime, connectionId, apiKey, apiKeyName, clientRawRequest, onRequestSuccess, reqLogger, toolNameMap, trackDone, appendLog, pxpipe, comboName }) { +export async function handleNonStreamingResponse({ providerResponse, provider, model, sourceFormat, targetFormat, body, stream, translatedBody, finalBody, requestStartTime, connectionId, apiKey, apiKeyInfo, apiKeyName, clientRawRequest, onRequestSuccess, reqLogger, toolNameMap, trackDone, appendLog, pxpipe, comboName }) { trackDone(); const contentType = providerResponse.headers.get("content-type") || ""; let responseBody; @@ -264,7 +264,7 @@ export async function handleNonStreamingResponse({ providerResponse, provider, m const usage = extractUsageFromResponse(responseBody); appendLog({ tokens: usage, status: "200 OK" }); - saveUsageStats({ provider, model, tokens: usage, connectionId, apiKey, endpoint: clientRawRequest?.endpoint, comboName }); + saveUsageStats({ provider, model, tokens: usage, connectionId, apiKey, apiKeyInfo, endpoint: clientRawRequest?.endpoint, comboName }); const translatedResponse = needsTranslation(targetFormat, sourceFormat) ? translateNonStreamingResponse(responseBody, targetFormat, sourceFormat) diff --git a/open-sse/handlers/chatCore/requestDetail.js b/open-sse/handlers/chatCore/requestDetail.js index 6e7dd74cb5..98c7624715 100644 --- a/open-sse/handlers/chatCore/requestDetail.js +++ b/open-sse/handlers/chatCore/requestDetail.js @@ -1,4 +1,5 @@ import { saveRequestUsage, appendRequestLog, saveRequestDetail } from "@/lib/usageDb.js"; +import { recordApiKeyUsage } from "@/lib/db/repos/apiKeyUsageRepo.js"; import { COLORS } from "../../utils/stream.js"; import { canonicalizeUsage } from "../../utils/usageTracking.js"; @@ -95,7 +96,7 @@ export function buildRequestDetail(base, overrides = {}) { }; } -export function saveUsageStats({ provider, model, tokens, connectionId, apiKey, endpoint, label = "USAGE" }) { +export function saveUsageStats({ provider, model, tokens, connectionId, apiKey, apiKeyInfo, endpoint, label = "USAGE" }) { if (!tokens || typeof tokens !== "object") return; const inTokens = tokens.input_tokens ?? tokens.prompt_tokens ?? 0; @@ -123,4 +124,10 @@ export function saveUsageStats({ provider, model, tokens, connectionId, apiKey, apiKey: apiKey || undefined, endpoint: endpoint || null }).catch(() => {}); + + // Record per-key usage limits against the API key info (if provided). + if (apiKeyInfo) { + const totalTokens = (normalized.prompt_tokens || 0) + (normalized.completion_tokens || 0); + recordApiKeyUsage(apiKeyInfo, totalTokens); + } } diff --git a/open-sse/handlers/chatCore/sseToJsonHandler.js b/open-sse/handlers/chatCore/sseToJsonHandler.js index 09383971be..463c774e3d 100644 --- a/open-sse/handlers/chatCore/sseToJsonHandler.js +++ b/open-sse/handlers/chatCore/sseToJsonHandler.js @@ -125,7 +125,7 @@ export function parseSSEToOpenAIResponse(rawSSE, fallbackModel, validToolNames = * Handle case: provider forced streaming but client wants JSON. * Supports both Codex/Responses API SSE and standard Chat Completions SSE. */ -export async function handleForcedSSEToJson({ providerResponse, sourceFormat, provider, model, body, stream, translatedBody, finalBody, requestStartTime, connectionId, apiKey, apiKeyName, clientRawRequest, onRequestSuccess, trackDone, appendLog, comboName, toolNameMap }) { +export async function handleForcedSSEToJson({ providerResponse, sourceFormat, provider, model, body, stream, translatedBody, finalBody, requestStartTime, connectionId, apiKey, apiKeyInfo, apiKeyName, clientRawRequest, onRequestSuccess, trackDone, appendLog, comboName, toolNameMap }) { const contentType = providerResponse.headers.get("content-type") || ""; const isSSE = contentType.includes("text/event-stream") || (contentType === "" && isResponsesProvider(provider)); if (!isSSE) return null; // not handled here @@ -147,7 +147,7 @@ export async function handleForcedSSEToJson({ providerResponse, sourceFormat, pr const usage = jsonResponse.usage || {}; appendLog({ tokens: usage, status: "200 OK" }); - saveUsageStats({ provider, model, tokens: usage, connectionId, apiKey, endpoint: clientRawRequest?.endpoint, comboName }); + saveUsageStats({ provider, model, tokens: usage, connectionId, apiKey, apiKeyInfo, endpoint: clientRawRequest?.endpoint, comboName }); const { msgItem, textContent } = pickAssistantMessageForChatCompletion(jsonResponse.output); const totalLatency = Date.now() - requestStartTime; @@ -244,7 +244,7 @@ export async function handleForcedSSEToJson({ providerResponse, sourceFormat, pr const usage = parsed.usage || {}; appendLog({ tokens: usage, status: "200 OK" }); - saveUsageStats({ provider, model, tokens: usage, connectionId, apiKey, endpoint: clientRawRequest?.endpoint }); + saveUsageStats({ provider, model, tokens: usage, connectionId, apiKey, apiKeyInfo, endpoint: clientRawRequest?.endpoint }); const totalLatency = Date.now() - requestStartTime; saveRequestDetail(buildRequestDetail({ diff --git a/open-sse/handlers/chatCore/streamingHandler.js b/open-sse/handlers/chatCore/streamingHandler.js index 2d540778da..505ce2d495 100644 --- a/open-sse/handlers/chatCore/streamingHandler.js +++ b/open-sse/handlers/chatCore/streamingHandler.js @@ -103,7 +103,12 @@ function buildTransformStream({ provider, sourceFormat, targetFormat, userAgent, * Includes a readiness gate: if upstream closes before any byte arrives, * return STREAM_EARLY_EOF so the caller can retry once on the same connection. */ -export async function handleStreamingResponse({ providerResponse, provider, model, sourceFormat, targetFormat, userAgent, body, stream, translatedBody, finalBody, requestStartTime, connectionId, apiKey, apiKeyName, clientRawRequest, onRequestSuccess, reqLogger, toolNameMap, streamController, onStreamComplete, streamDetailId, pxpipe }) { +export async function handleStreamingResponse({ + providerResponse, provider, model, sourceFormat, targetFormat, userAgent, + body, stream, translatedBody, finalBody, requestStartTime, connectionId, + apiKey, apiKeyInfo, apiKeyName, clientRawRequest, onRequestSuccess, + reqLogger, toolNameMap, streamController, onStreamComplete, streamDetailId, pxpipe, +}) { if (onRequestSuccess) { Promise.resolve() .then(onRequestSuccess) @@ -180,7 +185,7 @@ export async function handleStreamingResponse({ providerResponse, provider, mode /** * Build onStreamComplete callback for streaming usage tracking. */ -export function buildOnStreamComplete({ provider, model, connectionId, apiKey, apiKeyName, requestStartTime, body, stream, finalBody, translatedBody, clientRawRequest, pxpipe }) { +export function buildOnStreamComplete({ provider, model, connectionId, apiKey, apiKeyInfo, apiKeyName, requestStartTime, body, stream, finalBody, translatedBody, clientRawRequest, pxpipe }) { const streamDetailId = `${Date.now()}-${Math.random().toString(36).slice(2, 11)}`; const onStreamComplete = (contentObj, usage, ttftAt) => { @@ -205,7 +210,7 @@ export function buildOnStreamComplete({ provider, model, connectionId, apiKey, a console.error("[RequestDetail] Failed to update streaming content:", err.message); }); - saveUsageStats({ provider, model, tokens: usage, connectionId, apiKey, endpoint: clientRawRequest?.endpoint, label: "STREAM USAGE" }); + saveUsageStats({ provider, model, tokens: usage, connectionId, apiKey, apiKeyInfo, endpoint: clientRawRequest?.endpoint, label: "STREAM USAGE" }); }; return { onStreamComplete, streamDetailId }; diff --git a/src/app/(dashboard)/dashboard/endpoint/EndpointPageClient.js b/src/app/(dashboard)/dashboard/endpoint/EndpointPageClient.js index 76125dcca6..da89236f56 100644 --- a/src/app/(dashboard)/dashboard/endpoint/EndpointPageClient.js +++ b/src/app/(dashboard)/dashboard/endpoint/EndpointPageClient.js @@ -39,6 +39,16 @@ export default function APIPageClient({ machineId }) { const [editProvidersAll, setEditProvidersAll] = useState(true); const [editCombosAll, setEditCombosAll] = useState(true); const [editSaving, setEditSaving] = useState(false); + // Limit edit state (empty string = unlimited) + const [editExpiresAt, setEditExpiresAt] = useState(""); + const [editMaxTokens, setEditMaxTokens] = useState(""); + const [editMaxTokensDaily, setEditMaxTokensDaily] = useState(""); + const [editRpm, setEditRpm] = useState(""); + const [editRph, setEditRph] = useState(""); + const [editRpd, setEditRpd] = useState(""); + const [editTokens5h, setEditTokens5h] = useState(""); + const [editTokensWeekly, setEditTokensWeekly] = useState(""); + const [editTokensMonthly, setEditTokensMonthly] = useState(""); const [providerList, setProviderList] = useState([]); const [aliasMap, setAliasMap] = useState({}); // alias → provider ID const [comboList, setComboList] = useState([]); @@ -478,6 +488,15 @@ export default function APIPageClient({ machineId }) { const handleOpenEditKey = (key) => { setEditingKey(key); setEditName(key.name || ""); + setEditExpiresAt(key.expiresAt ? key.expiresAt.slice(0, 16) : ""); + setEditMaxTokens(key.maxTokens ?? ""); + setEditMaxTokensDaily(key.maxTokensDaily ?? ""); + setEditRpm(key.rpm ?? ""); + setEditRph(key.rph ?? ""); + setEditRpd(key.rpd ?? ""); + setEditTokens5h(key.tokens5h ?? ""); + setEditTokensWeekly(key.tokensWeekly ?? ""); + setEditTokensMonthly(key.tokensMonthly ?? ""); const ap = key.allowedProviders; const ac = key.allowedCombos; const ak = key.allowedKinds; @@ -514,10 +533,28 @@ export default function APIPageClient({ machineId }) { if (!editingKey) return; setEditSaving(true); try { + const parseNum = (v) => { + const s = String(v).trim(); + if (s === "" || s === "null" || s === "undefined") return null; + const n = Number(s); + return Number.isFinite(n) && n >= 0 ? n : null; + }; const body = { + name: editName.trim() || editingKey.name, allowedProviders: editProvidersAll ? null : editProviders, allowedCombos: editCombosAll ? null : editCombos, allowedKinds: editKindsAll ? null : editKinds, + limits: { + expiresAt: editExpiresAt ? new Date(editExpiresAt).toISOString() : null, + maxTokens: parseNum(editMaxTokens), + maxTokensDaily: parseNum(editMaxTokensDaily), + rpm: parseNum(editRpm), + rph: parseNum(editRph), + rpd: parseNum(editRpd), + tokens5h: parseNum(editTokens5h), + tokensWeekly: parseNum(editTokensWeekly), + tokensMonthly: parseNum(editTokensMonthly), + }, }; const res = await fetch(`/api/keys/${editingKey.id}`, { method: "PUT", @@ -1542,6 +1579,11 @@ export default function APIPageClient({ machineId }) {

{key.isActive === false && (

Paused

+ )} + {key.expiresAt && ( +

+ Expires {new Date(key.expiresAt).toLocaleString()} +

)} {/* ACL badges */} {(key.allowedProviders || key.allowedCombos || key.allowedKinds) && ( @@ -1571,6 +1613,14 @@ export default function APIPageClient({ machineId }) { {key.allowedKinds.length === 0 ? "No kinds" : key.allowedKinds.join(", ")} )} + {key.rpm != null && {key.rpm} RPM} + {key.rph != null && {key.rph} RPH} + {key.rpd != null && {key.rpd} RPD} + {key.maxTokens != null && max {key.maxTokens.toLocaleString()} tokens} + {key.maxTokensDaily != null && {key.maxTokensDaily.toLocaleString()} tokens/day} + {key.tokens5h != null && {key.tokens5h.toLocaleString()} tokens/5h} + {key.tokensWeekly != null && {key.tokensWeekly.toLocaleString()} tokens/week} + {key.tokensMonthly != null && {key.tokensMonthly.toLocaleString()} tokens/month} )} @@ -1785,6 +1835,49 @@ export default function APIPageClient({ machineId }) { {editCombosAll &&

This key can access all combos.

} + {/* Limits */} +
+ +
+
+ + setEditExpiresAt(e.target.value)} /> +
+
+ + setEditMaxTokens(e.target.value)} placeholder="∞" /> +
+
+ + setEditMaxTokensDaily(e.target.value)} placeholder="∞" /> +
+
+ + setEditRpm(e.target.value)} placeholder="∞" /> +
+
+ + setEditRph(e.target.value)} placeholder="∞" /> +
+
+ + setEditRpd(e.target.value)} placeholder="∞" /> +
+
+ + setEditTokens5h(e.target.value)} placeholder="∞" /> +
+
+ + setEditTokensWeekly(e.target.value)} placeholder="∞" /> +
+
+ + setEditTokensMonthly(e.target.value)} placeholder="∞" /> +
+
+
+ {/* ACL info */}

How it works:

@@ -1792,6 +1885,7 @@ export default function APIPageClient({ machineId }) {
  • All allowed = unrestricted (default). No ACL filtering.
  • Unchecked + empty = deny everything of that type.
  • Checked items = only those items are accessible.
    • +
  • Usage limits = leave blank to disable. Applies to requests made with this key.
    • @@ -1801,7 +1895,7 @@ export default function APIPageClient({ machineId }) { Cancel diff --git a/src/app/api/keys/[id]/route.js b/src/app/api/keys/[id]/route.js index fdaa54b419..86713a2251 100644 --- a/src/app/api/keys/[id]/route.js +++ b/src/app/api/keys/[id]/route.js @@ -1,5 +1,5 @@ import { NextResponse } from "next/server"; -import { deleteApiKey, getApiKeyById, updateApiKey } from "@/lib/localDb"; +import { deleteApiKey, getApiKeyById, updateApiKey, getApiKeyUsageSnapshot } from "@/lib/localDb"; // GET /api/keys/[id] - Get single key export async function GET(request, { params }) { @@ -9,7 +9,7 @@ export async function GET(request, { params }) { if (!key) { return NextResponse.json({ error: "Key not found" }, { status: 404 }); } - return NextResponse.json({ key }); + return NextResponse.json({ key: { ...key, usage: getApiKeyUsageSnapshot(key) } }); } catch (error) { console.log("Error fetching key:", error); return NextResponse.json({ error: "Failed to fetch key" }, { status: 500 }); @@ -21,7 +21,7 @@ export async function PUT(request, { params }) { try { const { id } = await params; const body = await request.json(); - const { isActive, allowedProviders, allowedCombos, allowedKinds } = body; + const { isActive, name, allowedProviders, allowedCombos, allowedKinds, limits = {} } = body; const existing = await getApiKeyById(id); if (!existing) { @@ -30,6 +30,7 @@ export async function PUT(request, { params }) { const updateData = {}; if (isActive !== undefined) updateData.isActive = isActive; + if (name !== undefined) updateData.name = name; // null = all allowed, [] = none, [x] = specific. Only update if key present in body. if ("allowedProviders" in body) { updateData.allowedProviders = allowedProviders === null ? null : (Array.isArray(allowedProviders) ? allowedProviders : null); @@ -40,9 +41,20 @@ export async function PUT(request, { params }) { if ("allowedKinds" in body) { updateData.allowedKinds = allowedKinds === null ? null : (Array.isArray(allowedKinds) ? allowedKinds : null); } + if ("expiresAt" in limits || "expiresAt" in body) { + updateData.expiresAt = limits.expiresAt ?? body.expiresAt ?? null; + } + if ("maxTokens" in limits) updateData.maxTokens = limits.maxTokens; + if ("maxTokensDaily" in limits) updateData.maxTokensDaily = limits.maxTokensDaily; + if ("rpm" in limits) updateData.rpm = limits.rpm; + if ("rph" in limits) updateData.rph = limits.rph; + if ("rpd" in limits) updateData.rpd = limits.rpd; + if ("tokens5h" in limits) updateData.tokens5h = limits.tokens5h; + if ("tokensWeekly" in limits) updateData.tokensWeekly = limits.tokensWeekly; + if ("tokensMonthly" in limits) updateData.tokensMonthly = limits.tokensMonthly; const updated = await updateApiKey(id, updateData); - return NextResponse.json({ key: updated }); + return NextResponse.json({ key: { ...updated, usage: getApiKeyUsageSnapshot(updated) } }); } catch (error) { console.log("Error updating key:", error); return NextResponse.json({ error: "Failed to update key" }, { status: 500 }); diff --git a/src/app/api/keys/route.js b/src/app/api/keys/route.js index a389504c50..a32907c8e7 100644 --- a/src/app/api/keys/route.js +++ b/src/app/api/keys/route.js @@ -1,5 +1,5 @@ import { NextResponse } from "next/server"; -import { getApiKeys, createApiKey } from "@/lib/localDb"; +import { getApiKeys, createApiKey, getApiKeyUsageSnapshot } from "@/lib/localDb"; import { getConsistentMachineId } from "@/shared/utils/machineId"; export const dynamic = "force-dynamic"; @@ -8,7 +8,11 @@ export const dynamic = "force-dynamic"; export async function GET() { try { const keys = await getApiKeys(); - return NextResponse.json({ keys }); + const keysWithUsage = keys.map((k) => ({ + ...k, + usage: getApiKeyUsageSnapshot(k), + })); + return NextResponse.json({ keys: keysWithUsage }); } catch (error) { console.log("Error fetching keys:", error); return NextResponse.json({ error: "Failed to fetch keys" }, { status: 500 }); @@ -19,7 +23,7 @@ export async function GET() { export async function POST(request) { try { const body = await request.json(); - const { name } = body; + const { name, limits = {} } = body; if (!name) { return NextResponse.json({ error: "Name is required" }, { status: 400 }); @@ -27,13 +31,24 @@ export async function POST(request) { // Always get machineId from server const machineId = await getConsistentMachineId(); - const apiKey = await createApiKey(name, machineId); + const apiKey = await createApiKey(name, machineId, limits); return NextResponse.json({ key: apiKey.key, name: apiKey.name, id: apiKey.id, machineId: apiKey.machineId, + limits: { + expiresAt: apiKey.expiresAt, + maxTokens: apiKey.maxTokens, + maxTokensDaily: apiKey.maxTokensDaily, + rpm: apiKey.rpm, + rph: apiKey.rph, + rpd: apiKey.rpd, + tokens5h: apiKey.tokens5h, + tokensWeekly: apiKey.tokensWeekly, + tokensMonthly: apiKey.tokensMonthly, + }, }, { status: 201 }); } catch (error) { console.log("Error creating key:", error); diff --git a/src/app/api/v1/models/route.js b/src/app/api/v1/models/route.js index 0f8de25b42..730700575c 100644 --- a/src/app/api/v1/models/route.js +++ b/src/app/api/v1/models/route.js @@ -3,6 +3,7 @@ import { getSettings } from "@/lib/localDb"; import { stripComboPrefix } from "open-sse/services/combo.js"; import { buildModelsList } from "@/sse/services/allowedModels.js"; import { capabilitiesFromServiceKind } from "open-sse/providers/capabilities.js"; +import { checkApiKeyLimits, recordApiKeyUsage } from "@/lib/db/repos/apiKeyUsageRepo.js"; const parseOpenAIStyleModels = (data) => { if (Array.isArray(data)) return data; @@ -54,6 +55,15 @@ export async function GET(request) { { status: 401, headers: { "Access-Control-Allow-Origin": "*" } } ); } + const limitCheck = checkApiKeyLimits(apiKeyInfo, 0); + if (!limitCheck.allowed) { + const retryAfter = limitCheck.retryAfterMs ? Math.ceil(limitCheck.retryAfterMs / 1000).toString() : undefined; + return Response.json( + { error: { message: limitCheck.reason, type: "rate_limit_error", code: "rate_limit_exceeded" } }, + { status: 429, headers: { "Access-Control-Allow-Origin": "*", ...(retryAfter ? { "Retry-After": retryAfter } : {}) } } + ); + } + recordApiKeyUsage(apiKeyInfo, 0); } let data = await buildModelsList([LLM_KIND], { skipDynamicFetch }); diff --git a/src/lib/db/migrations/005-add-api-key-limits.js b/src/lib/db/migrations/005-add-api-key-limits.js new file mode 100644 index 0000000000..0236cf0cf3 --- /dev/null +++ b/src/lib/db/migrations/005-add-api-key-limits.js @@ -0,0 +1,24 @@ +// Migration 005: add optional per-key usage/expiration limits. +// All columns are nullable; null means unlimited / no restriction. +export default { + version: 5, + name: "add-api-key-limits", + up(db) { + const rows = db.all("PRAGMA table_info(apiKeys)"); + const columns = Array.isArray(rows) ? rows.map((row) => row.name) : []; + const add = (col, type) => { + if (!columns.includes(col)) { + db.exec(`ALTER TABLE apiKeys ADD COLUMN ${col} ${type}`); + } + }; + add("expiresAt", "TEXT"); + add("maxTokens", "INTEGER"); + add("maxTokensDaily", "INTEGER"); + add("rpm", "INTEGER"); + add("rph", "INTEGER"); // requests per hour + add("rpd", "INTEGER"); // requests per day + add("tokens5h", "INTEGER"); // max tokens in rolling 5-hour window + add("tokensWeekly", "INTEGER"); + add("tokensMonthly", "INTEGER"); + }, +}; diff --git a/src/lib/db/migrations/index.js b/src/lib/db/migrations/index.js index c041147b37..4e4521e0b0 100644 --- a/src/lib/db/migrations/index.js +++ b/src/lib/db/migrations/index.js @@ -3,10 +3,11 @@ // Versions MUST be unique and monotonically increasing. import m001 from "./001-initial.js"; import m002 from "./002-fix-empty-allowed-lists.js"; -import m003 from "./003-add-allowed-lists-columns.js"; +import m003 from "./003-add-allowed-lists-columns.js"; import m004 from "./004-add-request-details-apikey.js"; +import m005 from "./005-add-api-key-limits.js"; -export const MIGRATIONS = [m001, m002, m003, m004].sort((a, b) => a.version - b.version); +export const MIGRATIONS = [m001, m002, m003, m004, m005].sort((a, b) => a.version - b.version); export function latestVersion() { return MIGRATIONS.length ? MIGRATIONS[MIGRATIONS.length - 1].version : 0; diff --git a/src/lib/db/repos/apiKeyUsageRepo.js b/src/lib/db/repos/apiKeyUsageRepo.js new file mode 100644 index 0000000000..e0919d1455 --- /dev/null +++ b/src/lib/db/repos/apiKeyUsageRepo.js @@ -0,0 +1,217 @@ +import { getAdapter } from "../driver.js"; + +// In-memory rolling counters. We intentionally keep these in process memory +// (not the DB) to avoid write amplification on every request and because +// limit enforcement is best-effort across restarts (a restart resets counters). +// For stricter accounting, callers can persist usage via usageRepo after the fact. + +if (!global._apiKeyCounters) { + global._apiKeyCounters = { + rpm: new Map(), // keyId -> { ts: minuteTimestamp, count } + rph: new Map(), // keyId -> { ts: hourTimestamp, count } + rpd: new Map(), // keyId -> { dateKey, count } + tokens5h: new Map(), // keyId -> [{ ts, tokens }] + tokensDaily: new Map(),// keyId -> { dateKey, tokens } + tokensWeekly: new Map(), // keyId -> { weekKey, tokens } + tokensMonthly: new Map(),// keyId -> { monthKey, tokens } + }; +} +const counters = global._apiKeyCounters; + +function getMinuteTs() { + return Math.floor(Date.now() / 60000); +} +function getHourTs() { + return Math.floor(Date.now() / 3600000); +} +function getDateKey() { + const d = new Date(); + return `${d.getFullYear()}-${String(d.getMonth() + 1).padStart(2, "0")}-${String(d.getDate()).padStart(2, "0")}`; +} +function getWeekKey() { + const d = new Date(); + const start = new Date(d); + start.setHours(0, 0, 0, 0); + start.setDate(d.getDate() - d.getDay()); // Sunday-based week + return `${start.getFullYear()}-${String(start.getMonth() + 1).padStart(2, "0")}-${String(start.getDate()).padStart(2, "0")}`; +} +function getMonthKey() { + const d = new Date(); + return `${d.getFullYear()}-${String(d.getMonth() + 1).padStart(2, "0")}`; +} + +function getRollingTokenCount(map, keyId, windowMs) { + const now = Date.now(); + const entries = map.get(keyId) || []; + const fresh = entries.filter((e) => now - e.ts <= windowMs); + if (fresh.length !== entries.length) map.set(keyId, fresh); + return fresh.reduce((sum, e) => sum + e.tokens, 0); +} + +function bumpCounter(map, keyId, bucketKey, amount = 1) { + let entry = map.get(keyId); + if (!entry || entry.key !== bucketKey) { + entry = { key: bucketKey, count: 0 }; + map.set(keyId, entry); + } + entry.count += amount; + return entry.count; +} + +function bumpTokens(map, keyId, amount) { + let entries = map.get(keyId); + if (!entries) { + entries = []; + map.set(keyId, entries); + } + entries.push({ ts: Date.now(), tokens: amount }); +} + +/** + * Check per-key limits BEFORE processing a request. + * @param {object} apiKeyInfo - key row from apiKeysRepo + * @param {number} requestedTokens - estimated tokens for this request (0 if unknown) + * @returns {{ allowed: boolean, reason?: string, retryAfterMs?: number }} + */ +export function checkApiKeyLimits(apiKeyInfo, requestedTokens = 0) { + if (!apiKeyInfo) return { allowed: true }; + const keyId = apiKeyInfo.id; + const tokens = Math.max(0, Number(requestedTokens) || 0); + + // RPM + const rpmLimit = apiKeyInfo.rpm; + if (rpmLimit != null) { + const minute = getMinuteTs(); + const current = bumpCounter(counters.rpm, keyId, minute, 0); + if (current + 1 > rpmLimit) { + return { allowed: false, reason: `Rate limit exceeded: ${rpmLimit} requests per minute`, retryAfterMs: (minute + 1) * 60000 - Date.now() }; + } + } + + // RPH + const rphLimit = apiKeyInfo.rph; + if (rphLimit != null) { + const hour = getHourTs(); + const current = bumpCounter(counters.rph, keyId, hour, 0); + if (current + 1 > rphLimit) { + return { allowed: false, reason: `Rate limit exceeded: ${rphLimit} requests per hour`, retryAfterMs: (hour + 1) * 3600000 - Date.now() }; + } + } + + // RPD + const rpdLimit = apiKeyInfo.rpd; + if (rpdLimit != null) { + const date = getDateKey(); + const current = bumpCounter(counters.rpd, keyId, date, 0); + if (current + 1 > rpdLimit) { + return { allowed: false, reason: `Rate limit exceeded: ${rpdLimit} requests per day`, retryAfterMs: 86400000 - (Date.now() % 86400000) }; + } + } + + // Max tokens per request + if (apiKeyInfo.maxTokens != null && tokens > apiKeyInfo.maxTokens) { + return { allowed: false, reason: `Token limit exceeded: max ${apiKeyInfo.maxTokens} tokens per request` }; + } + + // Daily tokens + if (apiKeyInfo.maxTokensDaily != null) { + const date = getDateKey(); + const current = counters.tokensDaily.get(keyId); + const used = current?.key === date ? current.tokens : 0; + if (used + tokens > apiKeyInfo.maxTokensDaily) { + return { allowed: false, reason: `Daily token limit exceeded: ${apiKeyInfo.maxTokensDaily}` }; + } + } + + // 5-hour rolling tokens + if (apiKeyInfo.tokens5h != null) { + const used = getRollingTokenCount(counters.tokens5h, keyId, 5 * 3600000); + if (used + tokens > apiKeyInfo.tokens5h) { + return { allowed: false, reason: `5-hour token window exceeded: ${apiKeyInfo.tokens5h}` }; + } + } + + // Weekly tokens + if (apiKeyInfo.tokensWeekly != null) { + const week = getWeekKey(); + const current = counters.tokensWeekly.get(keyId); + const used = current?.key === week ? current.tokens : 0; + if (used + tokens > apiKeyInfo.tokensWeekly) { + return { allowed: false, reason: `Weekly token limit exceeded: ${apiKeyInfo.tokensWeekly}` }; + } + } + + // Monthly tokens + if (apiKeyInfo.tokensMonthly != null) { + const month = getMonthKey(); + const current = counters.tokensMonthly.get(keyId); + const used = current?.key === month ? current.tokens : 0; + if (used + tokens > apiKeyInfo.tokensMonthly) { + return { allowed: false, reason: `Monthly token limit exceeded: ${apiKeyInfo.tokensMonthly}` }; + } + } + + return { allowed: true }; +} + +/** + * Record request usage against a key's counters. + * @param {object} apiKeyInfo + * @param {number} tokensUsed - total tokens consumed (prompt + completion) + */ +export function recordApiKeyUsage(apiKeyInfo, tokensUsed = 0) { + if (!apiKeyInfo) return; + const keyId = apiKeyInfo.id; + const tokens = Math.max(0, Number(tokensUsed) || 0); + + bumpCounter(counters.rpm, keyId, getMinuteTs(), 1); + bumpCounter(counters.rph, keyId, getHourTs(), 1); + bumpCounter(counters.rpd, keyId, getDateKey(), 1); + + if (tokens > 0) { + bumpTokens(counters.tokens5h, keyId, tokens); + + const date = getDateKey(); + const daily = counters.tokensDaily.get(keyId); + if (!daily || daily.key !== date) { + counters.tokensDaily.set(keyId, { key: date, tokens }); + } else { + daily.tokens += tokens; + } + + const week = getWeekKey(); + const weekly = counters.tokensWeekly.get(keyId); + if (!weekly || weekly.key !== week) { + counters.tokensWeekly.set(keyId, { key: week, tokens }); + } else { + weekly.tokens += tokens; + } + + const month = getMonthKey(); + const monthly = counters.tokensMonthly.get(keyId); + if (!monthly || monthly.key !== month) { + counters.tokensMonthly.set(keyId, { key: month, tokens }); + } else { + monthly.tokens += tokens; + } + } +} + +/** + * Get current usage snapshot for a key (for dashboard display). + * @param {object} apiKeyInfo + */ +export function getApiKeyUsageSnapshot(apiKeyInfo) { + if (!apiKeyInfo) return null; + const keyId = apiKeyInfo.id; + return { + rpm: { limit: apiKeyInfo.rpm, used: (counters.rpm.get(keyId)?.count || 0) }, + rph: { limit: apiKeyInfo.rph, used: (counters.rph.get(keyId)?.count || 0) }, + rpd: { limit: apiKeyInfo.rpd, used: (counters.rpd.get(keyId)?.count || 0) }, + tokens5h: { limit: apiKeyInfo.tokens5h, used: getRollingTokenCount(counters.tokens5h, keyId, 5 * 3600000) }, + maxTokens: { limit: apiKeyInfo.maxTokens, used: null }, + maxTokensDaily: { limit: apiKeyInfo.maxTokensDaily, used: (counters.tokensDaily.get(keyId)?.tokens || 0) }, + tokensWeekly: { limit: apiKeyInfo.tokensWeekly, used: (counters.tokensWeekly.get(keyId)?.tokens || 0) }, + tokensMonthly: { limit: apiKeyInfo.tokensMonthly, used: (counters.tokensMonthly.get(keyId)?.tokens || 0) }, + }; +} diff --git a/src/lib/db/repos/apiKeysRepo.js b/src/lib/db/repos/apiKeysRepo.js index 97541cf649..2177946f64 100644 --- a/src/lib/db/repos/apiKeysRepo.js +++ b/src/lib/db/repos/apiKeysRepo.js @@ -1,112 +1,174 @@ -import { randomUUID } from "node:crypto"; -import { getAdapter } from "../driver.js"; - -// Parse a JSON TEXT column with null=all / []=none semantics. -// DB NULL → null (all allowed). DB "[]" → [] (none). DB "[x]" → [x]. -function parsePermList(raw) { - if (raw === null || raw === undefined) return null; - try { return JSON.parse(raw); } catch { return null; } -} - -// Serialize back: null → null (DB NULL), [] → "[]", [x] → "[x]" -function serializePermList(val) { - if (val === null || val === undefined) return null; - return JSON.stringify(Array.isArray(val) ? val : []); -} - -function rowToKey(row) { - if (!row) return null; - return { - id: row.id, - key: row.key, - name: row.name, - machineId: row.machineId, - isActive: row.isActive === 1 || row.isActive === true, - createdAt: row.createdAt, - allowedProviders: parsePermList(row.allowedProviders), - allowedCombos: parsePermList(row.allowedCombos), - allowedKinds: parsePermList(row.allowedKinds), - }; -} - -export async function getApiKeys() { - const db = await getAdapter(); - const rows = db.all(`SELECT * FROM apiKeys ORDER BY createdAt ASC`); - return rows.map(rowToKey); -} - -export async function getApiKeyById(id) { - const db = await getAdapter(); - const row = db.get(`SELECT * FROM apiKeys WHERE id = ?`, [id]); - return rowToKey(row); -} - -export async function createApiKey(name, machineId) { - if (!machineId) throw new Error("machineId is required"); - const [db, { generateApiKeyWithMachine }] = await Promise.all([ - getAdapter(), - import("@/shared/utils/apiKey"), - ]); - const result = generateApiKeyWithMachine(machineId); - const apiKey = { - id: randomUUID(), - name, - key: result.key, - machineId, - isActive: true, - createdAt: new Date().toISOString(), - allowedProviders: null, - allowedCombos: null, - allowedKinds: null, - }; - db.run( - `INSERT INTO apiKeys(id, key, name, machineId, isActive, createdAt, allowedProviders, allowedCombos, allowedKinds) VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?)`, - [apiKey.id, apiKey.key, apiKey.name, apiKey.machineId, 1, apiKey.createdAt, null, null, null] - ); - return apiKey; -} - -export async function updateApiKey(id, data) { - const db = await getAdapter(); - let result = null; - db.transaction(() => { - const row = db.get(`SELECT * FROM apiKeys WHERE id = ?`, [id]); - if (!row) return; - const current = rowToKey(row); - // Merge: only override fields explicitly present in data - const merged = { ...current }; - if (data.isActive !== undefined) merged.isActive = data.isActive; - if (data.name !== undefined) merged.name = data.name; - if ("allowedProviders" in data) merged.allowedProviders = data.allowedProviders; - if ("allowedCombos" in data) merged.allowedCombos = data.allowedCombos; - if ("allowedKinds" in data) merged.allowedKinds = data.allowedKinds; - db.run( - `UPDATE apiKeys SET key = ?, name = ?, machineId = ?, isActive = ?, allowedProviders = ?, allowedCombos = ?, allowedKinds = ? WHERE id = ?`, - [ - merged.key, - merged.name, - merged.machineId, - merged.isActive ? 1 : 0, - serializePermList(merged.allowedProviders), - serializePermList(merged.allowedCombos), - serializePermList(merged.allowedKinds), - id, - ] - ); - result = merged; - }); - return result; -} - -export async function deleteApiKey(id) { - const db = await getAdapter(); - const res = db.run(`DELETE FROM apiKeys WHERE id = ?`, [id]); - return (res?.changes ?? 0) > 0; -} - -export async function validateApiKey(key) { - const db = await getAdapter(); - const row = db.get(`SELECT * FROM apiKeys WHERE key = ?`, [key]); - if (!row || row.isActive !== 1 && row.isActive !== true) return null; - return rowToKey(row); -} +import { randomUUID } from "node:crypto"; +import { getAdapter } from "../driver.js"; + +// Parse a JSON TEXT column with null=all / []=none semantics. +// DB NULL → null (all allowed). DB "[]" → [] (none). DB "[x]" → [x]. +function parsePermList(raw) { + if (raw === null || raw === undefined) return null; + try { return JSON.parse(raw); } catch { return null; } +} + +// Serialize back: null → null (DB NULL), [] → "[]", [x] → "[x]" +function serializePermList(val) { + if (val === null || val === undefined) return null; + return JSON.stringify(Array.isArray(val) ? val : []); +} + +// Parse optional integer limit: null/undefined stays null, otherwise integer or null on bad input. +function parseLimitInt(raw) { + if (raw === null || raw === undefined) return null; + const n = Number(raw); + return Number.isFinite(n) && n >= 0 ? Math.floor(n) : null; +} + +function rowToKey(row) { + if (!row) return null; + return { + id: row.id, + key: row.key, + name: row.name, + machineId: row.machineId, + isActive: row.isActive === 1 || row.isActive === true, + createdAt: row.createdAt, + allowedProviders: parsePermList(row.allowedProviders), + allowedCombos: parsePermList(row.allowedCombos), + allowedKinds: parsePermList(row.allowedKinds), + expiresAt: row.expiresAt || null, + maxTokens: parseLimitInt(row.maxTokens), + maxTokensDaily: parseLimitInt(row.maxTokensDaily), + rpm: parseLimitInt(row.rpm), + rph: parseLimitInt(row.rph), + rpd: parseLimitInt(row.rpd), + tokens5h: parseLimitInt(row.tokens5h), + tokensWeekly: parseLimitInt(row.tokensWeekly), + tokensMonthly: parseLimitInt(row.tokensMonthly), + }; +} + +export async function getApiKeys() { + const db = await getAdapter(); + const rows = db.all(`SELECT * FROM apiKeys ORDER BY createdAt ASC`); + return rows.map(rowToKey); +} + +export async function getApiKeyById(id) { + const db = await getAdapter(); + const row = db.get(`SELECT * FROM apiKeys WHERE id = ?`, [id]); + return rowToKey(row); +} + +export async function createApiKey(name, machineId, limits = {}) { + if (!machineId) throw new Error("machineId is required"); + const [db, { generateApiKeyWithMachine }] = await Promise.all([ + getAdapter(), + import("@/shared/utils/apiKey"), + ]); + const result = generateApiKeyWithMachine(machineId); + const apiKey = { + id: randomUUID(), + name, + key: result.key, + machineId, + isActive: true, + createdAt: new Date().toISOString(), + allowedProviders: null, + allowedCombos: null, + allowedKinds: null, + expiresAt: limits.expiresAt || null, + maxTokens: parseLimitInt(limits.maxTokens), + maxTokensDaily: parseLimitInt(limits.maxTokensDaily), + rpm: parseLimitInt(limits.rpm), + rph: parseLimitInt(limits.rph), + rpd: parseLimitInt(limits.rpd), + tokens5h: parseLimitInt(limits.tokens5h), + tokensWeekly: parseLimitInt(limits.tokensWeekly), + tokensMonthly: parseLimitInt(limits.tokensMonthly), + }; + db.run( + `INSERT INTO apiKeys( + id, key, name, machineId, isActive, createdAt, + allowedProviders, allowedCombos, allowedKinds, + expiresAt, maxTokens, maxTokensDaily, rpm, rph, rpd, tokens5h, tokensWeekly, tokensMonthly + ) VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + [ + apiKey.id, apiKey.key, apiKey.name, apiKey.machineId, 1, apiKey.createdAt, + null, null, null, + apiKey.expiresAt, apiKey.maxTokens, apiKey.maxTokensDaily, apiKey.rpm, apiKey.rph, apiKey.rpd, + apiKey.tokens5h, apiKey.tokensWeekly, apiKey.tokensMonthly, + ] + ); + return apiKey; +} + +export async function updateApiKey(id, data) { + const db = await getAdapter(); + let result = null; + db.transaction(() => { + const row = db.get(`SELECT * FROM apiKeys WHERE id = ?`, [id]); + if (!row) return; + const current = rowToKey(row); + // Merge: only override fields explicitly present in data + const merged = { ...current }; + if (data.isActive !== undefined) merged.isActive = data.isActive; + if (data.name !== undefined) merged.name = data.name; + if ("allowedProviders" in data) merged.allowedProviders = data.allowedProviders; + if ("allowedCombos" in data) merged.allowedCombos = data.allowedCombos; + if ("allowedKinds" in data) merged.allowedKinds = data.allowedKinds; + if ("expiresAt" in data) merged.expiresAt = data.expiresAt || null; + if ("maxTokens" in data) merged.maxTokens = parseLimitInt(data.maxTokens); + if ("maxTokensDaily" in data) merged.maxTokensDaily = parseLimitInt(data.maxTokensDaily); + if ("rpm" in data) merged.rpm = parseLimitInt(data.rpm); + if ("rph" in data) merged.rph = parseLimitInt(data.rph); + if ("rpd" in data) merged.rpd = parseLimitInt(data.rpd); + if ("tokens5h" in data) merged.tokens5h = parseLimitInt(data.tokens5h); + if ("tokensWeekly" in data) merged.tokensWeekly = parseLimitInt(data.tokensWeekly); + if ("tokensMonthly" in data) merged.tokensMonthly = parseLimitInt(data.tokensMonthly); + + db.run( + `UPDATE apiKeys SET key = ?, name = ?, machineId = ?, isActive = ?, + allowedProviders = ?, allowedCombos = ?, allowedKinds = ?, + expiresAt = ?, maxTokens = ?, maxTokensDaily = ?, rpm = ?, rph = ?, rpd = ?, + tokens5h = ?, tokensWeekly = ?, tokensMonthly = ? + WHERE id = ?`, + [ + merged.key, + merged.name, + merged.machineId, + merged.isActive ? 1 : 0, + serializePermList(merged.allowedProviders), + serializePermList(merged.allowedCombos), + serializePermList(merged.allowedKinds), + merged.expiresAt, + merged.maxTokens, + merged.maxTokensDaily, + merged.rpm, + merged.rph, + merged.rpd, + merged.tokens5h, + merged.tokensWeekly, + merged.tokensMonthly, + id, + ] + ); + result = merged; + }); + return result; +} + +export async function deleteApiKey(id) { + const db = await getAdapter(); + const res = db.run(`DELETE FROM apiKeys WHERE id = ?`, [id]); + return (res?.changes ?? 0) > 0; +} + +export async function validateApiKey(key) { + const db = await getAdapter(); + const row = db.get(`SELECT * FROM apiKeys WHERE key = ?`, [key]); + if (!row || (row.isActive !== 1 && row.isActive !== true)) return null; + const apiKey = rowToKey(row); + if (apiKey.expiresAt) { + const expiry = new Date(apiKey.expiresAt).getTime(); + if (expiry && expiry <= Date.now()) return null; + } + return apiKey; +} diff --git a/src/lib/db/schema.js b/src/lib/db/schema.js index 095187814b..3d53c0fd61 100644 --- a/src/lib/db/schema.js +++ b/src/lib/db/schema.js @@ -3,7 +3,7 @@ // pre-change safety backup in migrate.js: when the stored version is lower, // one lightweight DB backup is taken before applying schema changes. Forgetting // to bump only skips that backup — it does NOT break the additive auto-sync. -export const SCHEMA_VERSION = 2; +export const SCHEMA_VERSION = 5; export const PRAGMA_SQL = ` PRAGMA journal_mode = WAL; @@ -83,6 +83,15 @@ export const TABLES = { machineId: "TEXT", isActive: "INTEGER DEFAULT 1", createdAt: "TEXT NOT NULL", + expiresAt: "TEXT", + maxTokens: "INTEGER", + maxTokensDaily: "INTEGER", + rpm: "INTEGER", + rph: "INTEGER", + rpd: "INTEGER", + tokens5h: "INTEGER", + tokensWeekly: "INTEGER", + tokensMonthly: "INTEGER", }, indexes: ["CREATE INDEX IF NOT EXISTS idx_ak_key ON apiKeys(key)"], }, diff --git a/src/lib/localDb.js b/src/lib/localDb.js index 12427dfeb3..2480cf194c 100644 --- a/src/lib/localDb.js +++ b/src/lib/localDb.js @@ -11,6 +11,11 @@ export { getProxyPools, getProxyPoolById, createProxyPool, updateProxyPool, deleteProxyPool, getApiKeys, getApiKeyById, createApiKey, updateApiKey, deleteApiKey, validateApiKey, +} from "@/lib/db/index.js"; +export { + checkApiKeyLimits, recordApiKeyUsage, getApiKeyUsageSnapshot, +} from "@/lib/db/repos/apiKeyUsageRepo.js"; +export { getCombos, getComboById, getComboByName, createCombo, updateCombo, deleteCombo, getModelAliases, setModelAlias, deleteModelAlias, diff --git a/src/sse/handlers/chat.js b/src/sse/handlers/chat.js index e0f7cf9fa0..98679f32a2 100644 --- a/src/sse/handlers/chat.js +++ b/src/sse/handlers/chat.js @@ -11,6 +11,7 @@ import { isKindAllowed, isTrustedInternalRequest, } from "../services/auth.js"; +import { checkApiKeyLimits, recordApiKeyUsage } from "@/lib/db/repos/apiKeyUsageRepo.js"; import { isKimchiQuotaExhausted, buildKimchiQuotaExhaustedUpdate, @@ -123,6 +124,27 @@ export async function handleChat(request, clientRawRequest = null) { } } + // Enforce per-key usage limits (applies even when requireApiKey is false but a key was supplied) + if (apiKeyInfo) { + const estimatedTokens = body.max_tokens || 0; + const limitCheck = checkApiKeyLimits(apiKeyInfo, estimatedTokens); + if (!limitCheck.allowed) { + log.warn("AUTH", `API key limit exceeded: ${limitCheck.reason}`); + const retryAfter = limitCheck.retryAfterMs ? Math.ceil(limitCheck.retryAfterMs / 1000).toString() : undefined; + const body = JSON.stringify({ + error: { message: limitCheck.reason, type: "rate_limit_error", code: "rate_limit_exceeded" }, + }); + return new Response(body, { + status: HTTP_STATUS.RATE_LIMITED, + headers: { + "Content-Type": "application/json", + "Access-Control-Allow-Origin": "*", + ...(retryAfter ? { "Retry-After": retryAfter } : {}), + }, + }); + } + } + if (!modelStr) { log.warn("CHAT", "Missing model"); return errorResponse(HTTP_STATUS.BAD_REQUEST, "Missing model"); @@ -472,7 +494,8 @@ async function handleSingleModelChat(body, modelStr, clientRawRequest = null, re onRequestSuccess: async () => { await clearAccountError(credentials.connectionId, credentials, model); clearProviderFailure(provider, proxyHash); - } + }, + apiKeyInfo }); } finally { // Always release the semaphore slot, even if handleChatCore throws diff --git a/src/sse/services/auth.js b/src/sse/services/auth.js index 1c71f3b9c3..78e822b645 100644 --- a/src/sse/services/auth.js +++ b/src/sse/services/auth.js @@ -389,7 +389,10 @@ export async function isProviderAllowed(apiKeyInfo, providerIdOrAlias) { if (!apiKeyInfo) return true; const allowed = apiKeyInfo.allowedProviders; if (allowed === null || allowed === undefined) return true; // null = all - if (!Array.isArray(allowed) || allowed.length === 0) return false; // [] = none + if (!Array.isArray(allowed) || allowed.length === 0) { + log.warn("AUTH", `Provider "${providerIdOrAlias}" blocked for key ${apiKeyInfo.id?.slice(0, 8)}: allowedProviders is empty`); + return false; // [] = none + } if (allowed.includes(providerIdOrAlias)) return true; const alias = getProviderAlias(providerIdOrAlias); if (alias !== providerIdOrAlias && allowed.includes(alias)) return true; @@ -399,6 +402,7 @@ export async function isProviderAllowed(apiKeyInfo, providerIdOrAlias) { const prefix = await getNodePrefix(providerIdOrAlias); if (prefix && allowed.includes(prefix)) return true; } + log.warn("AUTH", `Provider "${providerIdOrAlias}" blocked for key ${apiKeyInfo.id?.slice(0, 8)}: allowed=[${allowed.join(", ")}]`); return false; } @@ -411,8 +415,15 @@ export function isComboAllowed(apiKeyInfo, comboName) { const name = comboName.startsWith("combo/") ? comboName.slice(6) : comboName; const allowed = apiKeyInfo.allowedCombos; if (allowed === null || allowed === undefined) return true; - if (!Array.isArray(allowed) || allowed.length === 0) return false; - return allowed.includes(name); + if (!Array.isArray(allowed) || allowed.length === 0) { + log.warn("AUTH", `Combo "${name}" blocked for key ${apiKeyInfo.id?.slice(0, 8)}: allowedCombos is empty`); + return false; + } + const ok = allowed.includes(name); + if (!ok) { + log.warn("AUTH", `Combo "${name}" blocked for key ${apiKeyInfo.id?.slice(0, 8)}: allowed=[${allowed.join(", ")}]`); + } + return ok; } /** @@ -424,7 +435,14 @@ export function isKindAllowed(apiKeyInfo, kind) { if (!apiKeyInfo) return true; const allowed = apiKeyInfo.allowedKinds; if (allowed === null || allowed === undefined) return true; // null = all - if (!Array.isArray(allowed) || allowed.length === 0) return false; // [] = none - return allowed.includes(kind); + if (!Array.isArray(allowed) || allowed.length === 0) { + log.warn("AUTH", `Kind "${kind}" blocked for key ${apiKeyInfo.id?.slice(0, 8)}: allowedKinds is empty`); + return false; // [] = none + } + const ok = allowed.includes(kind); + if (!ok) { + log.warn("AUTH", `Kind "${kind}" blocked for key ${apiKeyInfo.id?.slice(0, 8)}: allowed=[${allowed.join(", ")}]`); + } + return ok; }