diff --git a/config/quality/file-size-baseline.json b/config/quality/file-size-baseline.json index 85d8712212..735f121e37 100644 --- a/config/quality/file-size-baseline.json +++ b/config/quality/file-size-baseline.json @@ -149,7 +149,7 @@ "src/app/api/usage/analytics/route.ts": 941, "src/app/api/v1/models/catalog.ts": 1478, "src/lib/cloudflaredTunnel.ts": 934, - "src/lib/db/apiKeys.ts": 1661, + "src/lib/db/apiKeys.ts": 1662, "src/lib/db/core.ts": 1820, "src/lib/db/migrationRunner.ts": 1125, "src/lib/db/models.ts": 1184, diff --git a/src/app/(dashboard)/dashboard/api-manager/ApiManagerPageClient.tsx b/src/app/(dashboard)/dashboard/api-manager/ApiManagerPageClient.tsx index c182ff9c84..b31e683d2c 100644 --- a/src/app/(dashboard)/dashboard/api-manager/ApiManagerPageClient.tsx +++ b/src/app/(dashboard)/dashboard/api-manager/ApiManagerPageClient.tsx @@ -14,11 +14,14 @@ import { isRestricted as isKeyRestricted, classifyKeyStatus, computeApiKeyCounts, + formatUsdCost, + toLocalDateTimeInputValue, } from "./apiManagerPageUtils"; import type { KeyStatus, KeyType } from "./apiManagerPageUtils"; import { readActiveOnlyPreference, writeActiveOnlyPreference } from "./apiManagerPageStorage"; import { buildApiKeyCreateScopes, mergeApiKeyPermissionScopes } from "./apiManagerScopes"; import { SELF_ACCOUNT_QUOTA_SCOPE, SELF_USAGE_SCOPE } from "@/shared/constants/selfServiceScopes"; +import { UsageLimitSettings } from "./components/UsageLimitSettings"; // Constants for validation const MAX_KEY_NAME_LENGTH = 200; @@ -44,16 +47,6 @@ const CLAUDE_CODE_BLOCK_PATTERN_SET = new Set( Object.values(CLAUDE_CODE_FAMILY_BLOCK_PATTERNS).flat() ); -function toLocalDateTimeInputValue(value: string | null | undefined): string { - if (!value) return ""; - const date = new Date(value); - if (Number.isNaN(date.getTime())) return ""; - const pad = (n: number) => String(n).padStart(2, "0"); - return `${date.getFullYear()}-${pad(date.getMonth() + 1)}-${pad(date.getDate())}T${pad( - date.getHours() - )}:${pad(date.getMinutes())}`; -} - // Debounce hook for search optimization function useDebouncedValue(value: T, delay: number): T { const [debouncedValue, setDebouncedValue] = useState(value); @@ -129,6 +122,9 @@ interface ApiKey { streamDefaultMode?: StreamDefaultMode; disableNonPublicModels?: boolean; allowUsageCommand?: boolean; + usageLimitEnabled?: boolean; + dailyUsageLimitUsd?: number | null; + weeklyUsageLimitUsd?: number | null; allowedQuotas?: string[] | null; createdAt: string; } @@ -146,16 +142,6 @@ interface KeyUsageStats { lastUsed: string | null; } -function formatUsdCost(value: number, locale: string): string { - const amount = Number.isFinite(value) ? value : 0; - return new Intl.NumberFormat(locale, { - style: "currency", - currency: "USD", - minimumFractionDigits: amount > 0 && amount < 1 ? 4 : 2, - maximumFractionDigits: amount > 0 && amount < 1 ? 4 : 2, - }).format(amount); -} - interface Model { id: string; owned_by: string; @@ -664,6 +650,9 @@ export default function ApiManagerPageClient() { streamDefaultMode: StreamDefaultMode, disableNonPublicModels: boolean, allowUsageCommand: boolean, + usageLimitEnabled: boolean, + dailyUsageLimitUsd: number | null, + weeklyUsageLimitUsd: number | null, blockedModels: string[] ) => { if (!editingKey || !editingKey.id) return; @@ -732,6 +721,9 @@ export default function ApiManagerPageClient() { streamDefaultMode, disableNonPublicModels, allowUsageCommand, + usageLimitEnabled, + dailyUsageLimitUsd, + weeklyUsageLimitUsd, }), }); @@ -1046,6 +1038,12 @@ export default function ApiManagerPageClient() { {t("localUsageCommandBadge")} )} + {key.usageLimitEnabled === true && ( + + paid + USD quota + + )} {hasSessionLimit && ( group @@ -1469,6 +1467,9 @@ const PermissionsModal = memo(function PermissionsModal({ streamDefaultMode: StreamDefaultMode, disableNonPublicModels: boolean, allowUsageCommand: boolean, + usageLimitEnabled: boolean, + dailyUsageLimitUsd: number | null, + weeklyUsageLimitUsd: number | null, blockedModels: string[] ) => void; }) { @@ -1552,6 +1553,17 @@ const PermissionsModal = memo(function PermissionsModal({ const [usageCommandEnabled, setUsageCommandEnabled] = useState( apiKey?.allowUsageCommand === true ); + const [usageLimitEnabled, setUsageLimitEnabled] = useState(apiKey?.usageLimitEnabled === true); + const [dailyUsageLimitUsd, setDailyUsageLimitUsd] = useState( + typeof apiKey?.dailyUsageLimitUsd === "number" && apiKey.dailyUsageLimitUsd > 0 + ? String(apiKey.dailyUsageLimitUsd) + : "" + ); + const [weeklyUsageLimitUsd, setWeeklyUsageLimitUsd] = useState( + typeof apiKey?.weeklyUsageLimitUsd === "number" && apiKey.weeklyUsageLimitUsd > 0 + ? String(apiKey.weeklyUsageLimitUsd) + : "" + ); const getModelDisplayName = useCallback( (modelId: string) => modelId === CLAUDE_CODE_DEFAULT_MODEL_ID ? CLAUDE_CODE_DEFAULT_MODEL_NAME : modelId, @@ -1670,6 +1682,11 @@ const PermissionsModal = memo(function PermissionsModal({ [allowAllEndpoints] ); + const parseUsdLimitInput = useCallback((value: string): number | null => { + const parsed = Number(value); + return Number.isFinite(parsed) && parsed > 0 ? parsed : null; + }, []); + const handleSave = useCallback(() => { // Clear previous inline errors setNameError(null); @@ -1736,6 +1753,9 @@ const PermissionsModal = memo(function PermissionsModal({ streamDefaultMode, disableNonPublicModels, usageCommandEnabled, + usageLimitEnabled, + parseUsdLimitInput(dailyUsageLimitUsd), + parseUsdLimitInput(weeklyUsageLimitUsd), blockedModels ); }, [ @@ -1768,6 +1788,10 @@ const PermissionsModal = memo(function PermissionsModal({ streamDefaultMode, disableNonPublicModels, usageCommandEnabled, + usageLimitEnabled, + dailyUsageLimitUsd, + weeklyUsageLimitUsd, + parseUsdLimitInput, blockedClaudeCodeFamilies, initialBlockedModels, apiKey?.scopes, @@ -2341,6 +2365,16 @@ const PermissionsModal = memo(function PermissionsModal({ {t("localUsageCommand")} - {usageCommandEnabled ? tc("enabled") : tc("disabled")}

{t("localUsageCommandDesc")}

+ {/* Disable Non-Public Models Toggle */} diff --git a/src/app/(dashboard)/dashboard/api-manager/apiManagerPageUtils.ts b/src/app/(dashboard)/dashboard/api-manager/apiManagerPageUtils.ts index d6dd756900..844a7ab784 100644 --- a/src/app/(dashboard)/dashboard/api-manager/apiManagerPageUtils.ts +++ b/src/app/(dashboard)/dashboard/api-manager/apiManagerPageUtils.ts @@ -83,3 +83,23 @@ export function computeApiKeyCounts(keys: ApiKeyShape[]): ApiKeyCounts { return counts; } + +export function toLocalDateTimeInputValue(value: string | null | undefined): string { + if (!value) return ""; + const date = new Date(value); + if (Number.isNaN(date.getTime())) return ""; + const pad = (n: number) => String(n).padStart(2, "0"); + return `${date.getFullYear()}-${pad(date.getMonth() + 1)}-${pad(date.getDate())}T${pad( + date.getHours() + )}:${pad(date.getMinutes())}`; +} + +export function formatUsdCost(value: number, locale: string): string { + const amount = Number.isFinite(value) ? value : 0; + return new Intl.NumberFormat(locale, { + style: "currency", + currency: "USD", + minimumFractionDigits: amount > 0 && amount < 1 ? 4 : 2, + maximumFractionDigits: amount > 0 && amount < 1 ? 4 : 2, + }).format(amount); +} diff --git a/src/app/(dashboard)/dashboard/api-manager/components/UsageLimitSettings.tsx b/src/app/(dashboard)/dashboard/api-manager/components/UsageLimitSettings.tsx new file mode 100644 index 0000000000..1449d7192c --- /dev/null +++ b/src/app/(dashboard)/dashboard/api-manager/components/UsageLimitSettings.tsx @@ -0,0 +1,79 @@ +"use client"; + +import { Input } from "@/shared/components"; + +export function UsageLimitSettings({ + enabled, + dailyLimitUsd, + weeklyLimitUsd, + enabledLabel, + disabledLabel, + onEnabledChange, + onDailyLimitUsdChange, + onWeeklyLimitUsdChange, +}: { + enabled: boolean; + dailyLimitUsd: string; + weeklyLimitUsd: string; + enabledLabel: string; + disabledLabel: string; + onEnabledChange: (enabled: boolean) => void; + onDailyLimitUsdChange: (value: string) => void; + onWeeklyLimitUsdChange: (value: string) => void; +}) { + return ( +
+
+
+

USD usage quota

+

+ Blocks this key with a 400 API error after its local USD spend reaches the configured + daily or weekly quota. +

+
+ +
+
+
+ + onDailyLimitUsdChange(event.target.value)} + placeholder="0.00" + /> +
+
+ + onWeeklyLimitUsdChange(event.target.value)} + placeholder="0.00" + /> +
+
+

+ Weekly quota follows the cached Claude weekly reset when available; otherwise it falls back + to a rolling 7 day window. Daily quota uses the Fortaleza calendar day. +

+
+ ); +} diff --git a/src/app/(dashboard)/dashboard/costs/CostOverviewTab.tsx b/src/app/(dashboard)/dashboard/costs/CostOverviewTab.tsx index b9d06c66f1..829fb5060f 100644 --- a/src/app/(dashboard)/dashboard/costs/CostOverviewTab.tsx +++ b/src/app/(dashboard)/dashboard/costs/CostOverviewTab.tsx @@ -37,6 +37,9 @@ import { parseExplorerGroupBy, type CostRange, } from "./costExplorerParams"; +import { ApiKeyUsageLimitCard } from "./components/ApiKeyUsageLimitCard"; +import { MetricCard } from "./components/MetricCard"; +import { useApiKeyUsageLimits } from "./useApiKeyUsageLimits"; interface UsageAnalyticsSummary { totalCost: number; @@ -281,7 +284,9 @@ export default function CostOverviewTab() { const locale = useLocale(); const searchParams = useSearchParams(); const apiKeyIdsParam = searchParams.get("apiKeyIds"); - const apiKeyFilter = useMemo(() => parseApiKeyIds(apiKeyIdsParam).join(","), [apiKeyIdsParam]); + const selectedApiKeyIds = useMemo(() => parseApiKeyIds(apiKeyIdsParam), [apiKeyIdsParam]); + const selectedApiKeyId = selectedApiKeyIds.length === 1 ? selectedApiKeyIds[0] : null; + const apiKeyFilter = useMemo(() => selectedApiKeyIds.join(","), [selectedApiKeyIds]); const currencyFormatter = useMemo(() => createCurrencyFormatter(locale), [locale]); const [range, setRange] = useState(() => parseCostRange(searchParams.get("range"))); const [analytics, setAnalytics] = useState(null); @@ -300,6 +305,11 @@ export default function CostOverviewTab() { const [explorerSortKey, setExplorerSortKey] = useState("cost"); const [explorerSortDirection, setExplorerSortDirection] = useState("desc"); + const { + payload: apiKeyUsageLimits, + loading: apiKeyUsageLimitsLoading, + save: saveApiKeyUsageLimits, + } = useApiKeyUsageLimits(selectedApiKeyId); useEffect(() => { let active = true; @@ -536,6 +546,15 @@ export default function CostOverviewTab() { /> + {selectedApiKeyId && ( + + )} +
-

{label}

-

{loading ? "…" : value}

- {subValue ?

{subValue}

: null} - - ); -} - function CostExplorerCard({ rows, totalRows, diff --git a/src/app/(dashboard)/dashboard/costs/components/ApiKeyUsageLimitCard.tsx b/src/app/(dashboard)/dashboard/costs/components/ApiKeyUsageLimitCard.tsx new file mode 100644 index 0000000000..7916ac1633 --- /dev/null +++ b/src/app/(dashboard)/dashboard/costs/components/ApiKeyUsageLimitCard.tsx @@ -0,0 +1,214 @@ +"use client"; + +import { useEffect, useMemo, useState } from "react"; +import { Card } from "@/shared/components"; + +export interface ApiKeyUsageLimitPayload { + key: { + id: string; + name: string; + usageLimitEnabled: boolean; + dailyUsageLimitUsd: number | null; + weeklyUsageLimitUsd: number | null; + }; + status: { + enabled: boolean; + dailyLimitUsd: number | null; + weeklyLimitUsd: number | null; + dailySpentUsd: number; + weeklySpentUsd: number; + dailyWindowStartIso: string; + weeklyWindowStartIso: string; + weeklyResetAtIso: string | null; + dailyExceeded: boolean; + weeklyExceeded: boolean; + }; +} + +export interface ApiKeyUsageLimitSavePayload { + usageLimitEnabled: boolean; + dailyUsageLimitUsd: number | null; + weeklyUsageLimitUsd: number | null; +} + +function createCurrencyFormatter(locale: string) { + return new Intl.NumberFormat(locale, { + style: "currency", + currency: "USD", + minimumFractionDigits: 2, + maximumFractionDigits: 2, + }); +} + +function parseUsdInput(value: string): number | null { + const parsed = Number(value); + return Number.isFinite(parsed) && parsed > 0 ? parsed : null; +} + +function formatResetHint(resetAtIso: string | null): string { + if (!resetAtIso) return "fallback: rolling 7 days"; + const resetMs = Date.parse(resetAtIso); + if (!Number.isFinite(resetMs)) return "fallback: rolling 7 days"; + const deltaMs = resetMs - Date.now(); + if (deltaMs <= 0) return "reset due now"; + const hourMs = 60 * 60 * 1000; + const dayMs = 24 * hourMs; + if (deltaMs < dayMs) return `resets in ${Math.max(1, Math.ceil(deltaMs / hourMs))}h`; + return `resets in ${Math.max(1, Math.ceil(deltaMs / dayMs))}d`; +} + +function UsageQuotaMetric({ label, value }: { label: string; value: string }) { + return ( +
+

{label}

+

{value}

+
+ ); +} + +export function ApiKeyUsageLimitCard({ + payload, + loading, + locale, + onSave, +}: { + payload: ApiKeyUsageLimitPayload | null; + loading: boolean; + locale: string; + onSave: (next: ApiKeyUsageLimitSavePayload) => Promise; +}) { + const [enabled, setEnabled] = useState(false); + const [dailyLimit, setDailyLimit] = useState(""); + const [weeklyLimit, setWeeklyLimit] = useState(""); + const [saving, setSaving] = useState(false); + const [error, setError] = useState(null); + + useEffect(() => { + if (!payload) return; + setEnabled(payload.key.usageLimitEnabled); + setDailyLimit( + typeof payload.key.dailyUsageLimitUsd === "number" + ? String(payload.key.dailyUsageLimitUsd) + : "" + ); + setWeeklyLimit( + typeof payload.key.weeklyUsageLimitUsd === "number" + ? String(payload.key.weeklyUsageLimitUsd) + : "" + ); + setError(null); + }, [payload]); + + const formatter = useMemo(() => createCurrencyFormatter(locale), [locale]); + const status = payload?.status; + + const handleSave = async () => { + setSaving(true); + setError(null); + try { + await onSave({ + usageLimitEnabled: enabled, + dailyUsageLimitUsd: parseUsdInput(dailyLimit), + weeklyUsageLimitUsd: parseUsdInput(weeklyLimit), + }); + } catch (saveError) { + setError(saveError instanceof Error ? saveError.message : "Failed to save usage limits"); + } finally { + setSaving(false); + } + }; + + return ( + +
+
+
+ paid +

API key USD quota

+ {payload?.key.name && ( + + {payload.key.name} + + )} +
+

+ When enabled, @@om-usage returns daily quota, weekly quota, daily spend, and weekly + spend in USD. Weekly follows the cached Claude reset when available. +

+
+ +
+ +
+ + +
+ + setDailyLimit(event.target.value)} + className="mt-2 w-full rounded-md border border-border bg-background px-2 py-1.5 text-sm text-text-main" + placeholder="0.00" + /> +
+
+ + setWeeklyLimit(event.target.value)} + className="mt-2 w-full rounded-md border border-border bg-background px-2 py-1.5 text-sm text-text-main" + placeholder="0.00" + /> +

+ {formatResetHint(status?.weeklyResetAtIso ?? null)} +

+
+
+ + {error &&

{error}

} +
+ +
+ + ); +} diff --git a/src/app/(dashboard)/dashboard/costs/components/MetricCard.tsx b/src/app/(dashboard)/dashboard/costs/components/MetricCard.tsx new file mode 100644 index 0000000000..361bb25054 --- /dev/null +++ b/src/app/(dashboard)/dashboard/costs/components/MetricCard.tsx @@ -0,0 +1,23 @@ +import { Card } from "@/shared/components"; + +export function MetricCard({ + label, + value, + subValue, + color = "text-text-main", + loading = false, +}: { + label: string; + value: string; + subValue?: string; + color?: string; + loading?: boolean; +}) { + return ( + +

{label}

+

{loading ? "…" : value}

+ {subValue ?

{subValue}

: null} + + ); +} diff --git a/src/app/(dashboard)/dashboard/costs/useApiKeyUsageLimits.ts b/src/app/(dashboard)/dashboard/costs/useApiKeyUsageLimits.ts new file mode 100644 index 0000000000..82bf8d3b21 --- /dev/null +++ b/src/app/(dashboard)/dashboard/costs/useApiKeyUsageLimits.ts @@ -0,0 +1,51 @@ +"use client"; + +import { useCallback, useEffect, useState } from "react"; +import type { + ApiKeyUsageLimitPayload, + ApiKeyUsageLimitSavePayload, +} from "./components/ApiKeyUsageLimitCard"; + +export function useApiKeyUsageLimits(selectedApiKeyId: string | null) { + const [payload, setPayload] = useState(null); + const [loading, setLoading] = useState(false); + + const load = useCallback(async () => { + if (!selectedApiKeyId) { + setPayload(null); + return; + } + setLoading(true); + try { + const response = await fetch( + `/api/keys/${encodeURIComponent(selectedApiKeyId)}/usage-limits` + ); + if (!response.ok) throw new Error("Failed to load API key usage limits"); + setPayload((await response.json()) as ApiKeyUsageLimitPayload); + } catch { + setPayload(null); + } finally { + setLoading(false); + } + }, [selectedApiKeyId]); + + const save = useCallback( + async (next: ApiKeyUsageLimitSavePayload) => { + if (!selectedApiKeyId) return; + const response = await fetch(`/api/keys/${encodeURIComponent(selectedApiKeyId)}`, { + method: "PATCH", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(next), + }); + if (!response.ok) throw new Error("Failed to save API key usage limits"); + await load(); + }, + [load, selectedApiKeyId] + ); + + useEffect(() => { + void load(); + }, [load]); + + return { payload, loading, save }; +} diff --git a/src/app/api/keys/[id]/route.ts b/src/app/api/keys/[id]/route.ts index d4161f59b6..8fd10cc338 100644 --- a/src/app/api/keys/[id]/route.ts +++ b/src/app/api/keys/[id]/route.ts @@ -83,6 +83,9 @@ export async function PATCH(request, { params }) { streamDefaultMode, disableNonPublicModels, allowUsageCommand, + usageLimitEnabled, + dailyUsageLimitUsd, + weeklyUsageLimitUsd, } = validation.data; const payload: Parameters[1] = {}; @@ -106,6 +109,9 @@ export async function PATCH(request, { params }) { if (disableNonPublicModels !== undefined) payload.disableNonPublicModels = disableNonPublicModels; if (allowUsageCommand !== undefined) payload.allowUsageCommand = allowUsageCommand; + if (usageLimitEnabled !== undefined) payload.usageLimitEnabled = usageLimitEnabled; + if (dailyUsageLimitUsd !== undefined) payload.dailyUsageLimitUsd = dailyUsageLimitUsd; + if (weeklyUsageLimitUsd !== undefined) payload.weeklyUsageLimitUsd = weeklyUsageLimitUsd; const updated = await updateApiKeyPermissions(id, payload); if (!updated) { @@ -136,6 +142,9 @@ export async function PATCH(request, { params }) { ...(streamDefaultMode !== undefined && { streamDefaultMode }), ...(disableNonPublicModels !== undefined && { disableNonPublicModels }), ...(allowUsageCommand !== undefined && { allowUsageCommand }), + ...(usageLimitEnabled !== undefined && { usageLimitEnabled }), + ...(dailyUsageLimitUsd !== undefined && { dailyUsageLimitUsd }), + ...(weeklyUsageLimitUsd !== undefined && { weeklyUsageLimitUsd }), }); } catch (error) { log.error("keys", "Error updating key permissions", error); diff --git a/src/app/api/keys/[id]/usage-limits/route.ts b/src/app/api/keys/[id]/usage-limits/route.ts new file mode 100644 index 0000000000..a35dc89e1b --- /dev/null +++ b/src/app/api/keys/[id]/usage-limits/route.ts @@ -0,0 +1,44 @@ +import { NextResponse } from "next/server"; +import { requireManagementAuth } from "@/lib/api/requireManagementAuth"; +import { getApiKeyById } from "@/lib/db/apiKeys"; +import { getApiKeyUsageLimitStatus } from "@/lib/usage/apiKeyUsageLimits"; +import * as log from "@/sse/utils/logger"; + +export async function GET(request: Request, { params }: { params: Promise<{ id: string }> }) { + const authError = await requireManagementAuth(request); + if (authError) return authError; + + try { + const { id } = await params; + const key = await getApiKeyById(id); + if (!key || typeof key.id !== "string") { + return NextResponse.json({ error: "Key not found" }, { status: 404 }); + } + + const status = await getApiKeyUsageLimitStatus({ + id: key.id, + allowedConnections: Array.isArray(key.allowedConnections) ? key.allowedConnections : [], + usageLimitEnabled: key.usageLimitEnabled === true, + dailyUsageLimitUsd: + typeof key.dailyUsageLimitUsd === "number" ? key.dailyUsageLimitUsd : null, + weeklyUsageLimitUsd: + typeof key.weeklyUsageLimitUsd === "number" ? key.weeklyUsageLimitUsd : null, + }); + + return NextResponse.json({ + key: { + id: key.id, + name: typeof key.name === "string" ? key.name : "", + usageLimitEnabled: key.usageLimitEnabled === true, + dailyUsageLimitUsd: + typeof key.dailyUsageLimitUsd === "number" ? key.dailyUsageLimitUsd : null, + weeklyUsageLimitUsd: + typeof key.weeklyUsageLimitUsd === "number" ? key.weeklyUsageLimitUsd : null, + }, + status, + }); + } catch (error) { + log.error("keys", "Error fetching API key usage limits", error); + return NextResponse.json({ error: "Failed to fetch usage limits" }, { status: 500 }); + } +} diff --git a/src/app/api/keys/route.ts b/src/app/api/keys/route.ts index 32e2954f16..b4611943f6 100644 --- a/src/app/api/keys/route.ts +++ b/src/app/api/keys/route.ts @@ -63,16 +63,33 @@ export async function POST(request) { if (isValidationFailure(validation)) { return NextResponse.json({ error: validation.error }, { status: 400 }); } - const { name, noLog, scopes, allowUsageCommand } = validation.data; + const { + name, + noLog, + scopes, + allowUsageCommand, + usageLimitEnabled, + dailyUsageLimitUsd, + weeklyUsageLimitUsd, + } = validation.data; // Always get machineId from server const machineId = await getConsistentMachineId(); const normalizedScopes = normalizeSelfServiceScopesForCreate(scopes); const apiKey = await createApiKey(name, machineId, normalizedScopes); - if (noLog === true || allowUsageCommand === true) { + if ( + noLog === true || + allowUsageCommand === true || + usageLimitEnabled === true || + dailyUsageLimitUsd !== undefined || + weeklyUsageLimitUsd !== undefined + ) { await updateApiKeyPermissions(apiKey.id, { ...(noLog === true && { noLog: true }), ...(allowUsageCommand === true && { allowUsageCommand: true }), + ...(usageLimitEnabled === true && { usageLimitEnabled: true }), + ...(dailyUsageLimitUsd !== undefined && { dailyUsageLimitUsd }), + ...(weeklyUsageLimitUsd !== undefined && { weeklyUsageLimitUsd }), }); } @@ -87,6 +104,9 @@ export async function POST(request) { machineId: apiKey.machineId, noLog: noLog === true, allowUsageCommand: allowUsageCommand === true, + usageLimitEnabled: usageLimitEnabled === true, + dailyUsageLimitUsd: dailyUsageLimitUsd ?? null, + weeklyUsageLimitUsd: weeklyUsageLimitUsd ?? null, streamDefaultMode: "legacy", }, { status: 201 } diff --git a/src/lib/db/apiKeyColumnFallbacks.ts b/src/lib/db/apiKeyColumnFallbacks.ts new file mode 100644 index 0000000000..398672c5fc --- /dev/null +++ b/src/lib/db/apiKeyColumnFallbacks.ts @@ -0,0 +1,47 @@ +export const API_KEY_COLUMN_FALLBACKS = [ + { name: "allowed_models", definition: "allowed_models TEXT" }, + { name: "blocked_models", definition: "blocked_models TEXT" }, + { name: "allowed_combos", definition: "allowed_combos TEXT" }, + { name: "no_log", definition: "no_log INTEGER NOT NULL DEFAULT 0" }, + { name: "allowed_connections", definition: "allowed_connections TEXT" }, + { name: "auto_resolve", definition: "auto_resolve INTEGER NOT NULL DEFAULT 0" }, + { name: "is_active", definition: "is_active INTEGER NOT NULL DEFAULT 1" }, + { name: "access_schedule", definition: "access_schedule TEXT" }, + { name: "max_requests_per_day", definition: "max_requests_per_day INTEGER" }, + { name: "max_requests_per_minute", definition: "max_requests_per_minute INTEGER" }, + { name: "throttle_delay_ms", definition: "throttle_delay_ms INTEGER" }, + { name: "max_sessions", definition: "max_sessions INTEGER NOT NULL DEFAULT 0" }, + { name: "revoked_at", definition: "revoked_at TEXT" }, + { name: "expires_at", definition: "expires_at TEXT" }, + { name: "last_used_at", definition: "last_used_at TEXT" }, + { name: "key_prefix", definition: "key_prefix TEXT" }, + { name: "ip_allowlist", definition: "ip_allowlist TEXT" }, + { name: "scopes", definition: "scopes TEXT" }, + { name: "rate_limits", definition: "rate_limits TEXT" }, + { name: "is_banned", definition: "is_banned INTEGER NOT NULL DEFAULT 0" }, + { name: "key_hash", definition: "key_hash TEXT" }, + { name: "proxy_id", definition: "proxy_id TEXT" }, + { name: "allowed_endpoints", definition: "allowed_endpoints TEXT" }, + { name: "allowed_quotas", definition: "allowed_quotas TEXT NOT NULL DEFAULT '[]'" }, + { name: "stream_default_mode", definition: "stream_default_mode TEXT NOT NULL DEFAULT 'legacy'" }, + { + name: "disable_non_public_models", + definition: "disable_non_public_models INTEGER NOT NULL DEFAULT 0", + }, + { + name: "allow_usage_command", + definition: "allow_usage_command INTEGER NOT NULL DEFAULT 0", + }, + { + name: "usage_limit_enabled", + definition: "usage_limit_enabled INTEGER NOT NULL DEFAULT 0", + }, + { + name: "daily_usage_limit_usd", + definition: "daily_usage_limit_usd REAL", + }, + { + name: "weekly_usage_limit_usd", + definition: "weekly_usage_limit_usd REAL", + }, +] as const; diff --git a/src/lib/db/apiKeyUsageLimitFields.ts b/src/lib/db/apiKeyUsageLimitFields.ts new file mode 100644 index 0000000000..57f0c60b37 --- /dev/null +++ b/src/lib/db/apiKeyUsageLimitFields.ts @@ -0,0 +1,65 @@ +type UsageLimitRecord = Record; + +export interface ApiKeyUsageLimitFields { + usageLimitEnabled: boolean; + dailyUsageLimitUsd: number | null; + weeklyUsageLimitUsd: number | null; +} + +export function parseUsageLimitEnabled(value: unknown): boolean { + return value === true || value === 1 || value === "1"; +} + +export function parseNullablePositiveNumber(value: unknown): number | null { + if (typeof value === "number" && Number.isFinite(value) && value > 0) return value; + if (typeof value === "string" && value.trim()) { + const parsed = Number(value); + return Number.isFinite(parsed) && parsed > 0 ? parsed : null; + } + return null; +} + +export function parseApiKeyUsageLimitFields(record: UsageLimitRecord): ApiKeyUsageLimitFields { + return { + usageLimitEnabled: parseUsageLimitEnabled( + record.usage_limit_enabled ?? record.usageLimitEnabled + ), + dailyUsageLimitUsd: parseNullablePositiveNumber( + record.daily_usage_limit_usd ?? record.dailyUsageLimitUsd + ), + weeklyUsageLimitUsd: parseNullablePositiveNumber( + record.weekly_usage_limit_usd ?? record.weeklyUsageLimitUsd + ), + }; +} + +export function hasUsageLimitUpdate(update: UsageLimitRecord): boolean { + return ( + update.usageLimitEnabled !== undefined || + update.dailyUsageLimitUsd !== undefined || + update.weeklyUsageLimitUsd !== undefined + ); +} + +export function appendUsageLimitUpdates( + update: UsageLimitRecord, + updates: string[], + params: { + usageLimitEnabled?: number; + dailyUsageLimitUsd?: number | null; + weeklyUsageLimitUsd?: number | null; + } +) { + if (update.usageLimitEnabled !== undefined) { + updates.push("usage_limit_enabled = @usageLimitEnabled"); + params.usageLimitEnabled = update.usageLimitEnabled === true ? 1 : 0; + } + if (update.dailyUsageLimitUsd !== undefined) { + updates.push("daily_usage_limit_usd = @dailyUsageLimitUsd"); + params.dailyUsageLimitUsd = parseNullablePositiveNumber(update.dailyUsageLimitUsd); + } + if (update.weeklyUsageLimitUsd !== undefined) { + updates.push("weekly_usage_limit_usd = @weeklyUsageLimitUsd"); + params.weeklyUsageLimitUsd = parseNullablePositiveNumber(update.weeklyUsageLimitUsd); + } +} diff --git a/src/lib/db/apiKeys.ts b/src/lib/db/apiKeys.ts index 47b82d7784..810ca53300 100644 --- a/src/lib/db/apiKeys.ts +++ b/src/lib/db/apiKeys.ts @@ -8,6 +8,12 @@ import { getDbInstance, rowToCamel } from "./core"; import { backupDbFile } from "./backup"; import { registerDbStateResetter } from "./stateReset"; import { getKeyGroupsForApiKey, checkKeyModelAccess } from "./apiKeyGroups"; +import { API_KEY_COLUMN_FALLBACKS } from "./apiKeyColumnFallbacks"; +import { + appendUsageLimitUpdates, + hasUsageLimitUpdate, + parseApiKeyUsageLimitFields, +} from "./apiKeyUsageLimitFields"; import { setNoLog } from "../compliance/noLog"; import { resolveModelAlias } from "@omniroute/open-sse/services/modelDeprecation.ts"; import { getSyncedAvailableModelsByConnection, getCustomModels, getModelIsHidden } from "./models"; @@ -68,6 +74,9 @@ interface ApiKeyMetadata { streamDefaultMode: "legacy" | "json"; disableNonPublicModels: boolean; allowUsageCommand: boolean; + usageLimitEnabled: boolean; + dailyUsageLimitUsd: number | null; + weeklyUsageLimitUsd: number | null; } interface ApiKeyRow extends JsonRecord { @@ -101,6 +110,12 @@ interface ApiKeyRow extends JsonRecord { streamDefaultMode?: unknown; allow_usage_command?: unknown; allowUsageCommand?: unknown; + usage_limit_enabled?: unknown; + usageLimitEnabled?: unknown; + daily_usage_limit_usd?: unknown; + dailyUsageLimitUsd?: unknown; + weekly_usage_limit_usd?: unknown; + weeklyUsageLimitUsd?: unknown; } interface StatementLike { @@ -144,6 +159,9 @@ interface ApiKeyView extends JsonRecord { streamDefaultMode: "legacy" | "json"; disableNonPublicModels?: boolean; allowUsageCommand?: boolean; + usageLimitEnabled?: boolean; + dailyUsageLimitUsd?: number | null; + weeklyUsageLimitUsd?: number | null; } // LRU cache for API key validation (valid keys only) @@ -159,42 +177,6 @@ const CLAUDE_CODE_SHORT_ALIASES = new Set(["sonnet", "opus", "haiku", "fable"]); // Wildcard scope matching is now handled by `matchesWildcardPattern` // (deterministic, no RegExp from dynamic strings). -const API_KEY_COLUMN_FALLBACKS = [ - { name: "allowed_models", definition: "allowed_models TEXT" }, - { name: "blocked_models", definition: "blocked_models TEXT" }, - { name: "allowed_combos", definition: "allowed_combos TEXT" }, - { name: "no_log", definition: "no_log INTEGER NOT NULL DEFAULT 0" }, - { name: "allowed_connections", definition: "allowed_connections TEXT" }, - { name: "auto_resolve", definition: "auto_resolve INTEGER NOT NULL DEFAULT 0" }, - { name: "is_active", definition: "is_active INTEGER NOT NULL DEFAULT 1" }, - { name: "access_schedule", definition: "access_schedule TEXT" }, - { name: "max_requests_per_day", definition: "max_requests_per_day INTEGER" }, - { name: "max_requests_per_minute", definition: "max_requests_per_minute INTEGER" }, - { name: "throttle_delay_ms", definition: "throttle_delay_ms INTEGER" }, - { name: "max_sessions", definition: "max_sessions INTEGER NOT NULL DEFAULT 0" }, - { name: "revoked_at", definition: "revoked_at TEXT" }, - { name: "expires_at", definition: "expires_at TEXT" }, - { name: "last_used_at", definition: "last_used_at TEXT" }, - { name: "key_prefix", definition: "key_prefix TEXT" }, - { name: "ip_allowlist", definition: "ip_allowlist TEXT" }, - { name: "scopes", definition: "scopes TEXT" }, - { name: "rate_limits", definition: "rate_limits TEXT" }, - { name: "is_banned", definition: "is_banned INTEGER NOT NULL DEFAULT 0" }, - { name: "key_hash", definition: "key_hash TEXT" }, - { name: "proxy_id", definition: "proxy_id TEXT" }, - { name: "allowed_endpoints", definition: "allowed_endpoints TEXT" }, - { name: "allowed_quotas", definition: "allowed_quotas TEXT NOT NULL DEFAULT '[]'" }, - { name: "stream_default_mode", definition: "stream_default_mode TEXT NOT NULL DEFAULT 'legacy'" }, - { - name: "disable_non_public_models", - definition: "disable_non_public_models INTEGER NOT NULL DEFAULT 0", - }, - { - name: "allow_usage_command", - definition: "allow_usage_command INTEGER NOT NULL DEFAULT 0", - }, -] as const; - // Cache for model permission checks const _modelPermissionCache = new Map(); @@ -510,7 +492,7 @@ function getPreparedStatements(db: ApiKeysDbLike): ApiKeysStatements { "SELECT id, expires_at, revoked_at, is_active, is_banned FROM api_keys WHERE key = ? OR key_hash = ?" ); _stmtGetKeyMetadata = db.prepare( - "SELECT id, name, machine_id, allowed_models, blocked_models, allowed_combos, allowed_connections, allowed_quotas, no_log, auto_resolve, is_active, access_schedule, max_requests_per_day, max_requests_per_minute, throttle_delay_ms, max_sessions, revoked_at, expires_at, ip_allowlist, scopes, rate_limits, is_banned, key_hash, allowed_endpoints, stream_default_mode, disable_non_public_models, allow_usage_command, proxy_id FROM api_keys WHERE key = ? OR key_hash = ?" + "SELECT id, name, machine_id, allowed_models, blocked_models, allowed_combos, allowed_connections, allowed_quotas, no_log, auto_resolve, is_active, access_schedule, max_requests_per_day, max_requests_per_minute, throttle_delay_ms, max_sessions, revoked_at, expires_at, ip_allowlist, scopes, rate_limits, is_banned, key_hash, allowed_endpoints, stream_default_mode, disable_non_public_models, allow_usage_command, usage_limit_enabled, daily_usage_limit_usd, weekly_usage_limit_usd, proxy_id FROM api_keys WHERE key = ? OR key_hash = ?" ); _stmtInsertKey = db.prepare( "INSERT INTO api_keys (id, name, key, machine_id, allowed_models, no_log, created_at, key_prefix, key_hash, scopes) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)" @@ -563,6 +545,7 @@ export async function getApiKeys() { (camelRow as JsonRecord).disableNonPublicModels ); camelRow.allowUsageCommand = parseAllowUsageCommand((camelRow as JsonRecord).allowUsageCommand); + Object.assign(camelRow, parseApiKeyUsageLimitFields(camelRow)); if (typeof camelRow.id === "string" && camelRow.id.length > 0) { setNoLog(camelRow.id, camelRow.noLog === true); } @@ -594,6 +577,7 @@ export async function getApiKeyById(id: string) { (camelRow as JsonRecord).disableNonPublicModels ); camelRow.allowUsageCommand = parseAllowUsageCommand((camelRow as JsonRecord).allowUsageCommand); + Object.assign(camelRow, parseApiKeyUsageLimitFields(camelRow)); if (typeof camelRow.id === "string" && camelRow.id.length > 0) { setNoLog(camelRow.id, camelRow.noLog === true); } @@ -866,6 +850,9 @@ export async function updateApiKeyPermissions( streamDefaultMode?: "legacy" | "json" | null; disableNonPublicModels?: boolean; allowUsageCommand?: boolean; + usageLimitEnabled?: boolean; + dailyUsageLimitUsd?: number | null; + weeklyUsageLimitUsd?: number | null; } ) { const db = getDbInstance() as ApiKeysDbLike; @@ -900,6 +887,10 @@ export async function updateApiKeyPermissions( disableNonPublicModels: (update as { disableNonPublicModels?: boolean }) .disableNonPublicModels, allowUsageCommand: (update as { allowUsageCommand?: boolean }).allowUsageCommand, + usageLimitEnabled: (update as { usageLimitEnabled?: boolean }).usageLimitEnabled, + dailyUsageLimitUsd: (update as { dailyUsageLimitUsd?: number | null }).dailyUsageLimitUsd, + weeklyUsageLimitUsd: (update as { weeklyUsageLimitUsd?: number | null }) + .weeklyUsageLimitUsd, }; if ( @@ -925,7 +916,8 @@ export async function updateApiKeyPermissions( (normalized as Record).allowedEndpoints === undefined && (normalized as Record).streamDefaultMode === undefined && normalized.disableNonPublicModels === undefined && - normalized.allowUsageCommand === undefined + normalized.allowUsageCommand === undefined && + !hasUsageLimitUpdate(normalized as Record) ) { return false; } @@ -955,6 +947,9 @@ export async function updateApiKeyPermissions( streamDefaultMode?: "legacy" | "json"; disableNonPublicModels?: number; allowUsageCommand?: number; + usageLimitEnabled?: number; + dailyUsageLimitUsd?: number | null; + weeklyUsageLimitUsd?: number | null; } = { id }; if (normalized.name !== undefined) { @@ -1058,6 +1053,8 @@ export async function updateApiKeyPermissions( params.allowUsageCommand = normalized.allowUsageCommand ? 1 : 0; } + appendUsageLimitUpdates(normalized as Record, updates, params); + const maxSessionsUpdate = (normalized as Record).maxSessions; if (maxSessionsUpdate !== undefined) { updates.push("max_sessions = @maxSessions"); @@ -1439,6 +1436,9 @@ export async function getApiKeyMetadata( streamDefaultMode: "legacy", disableNonPublicModels: false, allowUsageCommand: false, + usageLimitEnabled: false, + dailyUsageLimitUsd: null, + weeklyUsageLimitUsd: null, }; } @@ -1512,6 +1512,7 @@ export async function getApiKeyMetadata( allowUsageCommand: parseAllowUsageCommand( (record as JsonRecord).allow_usage_command ?? (record as JsonRecord).allowUsageCommand ), + ...parseApiKeyUsageLimitFields(record as JsonRecord), }; if (!metadata.id) { diff --git a/src/lib/db/migrations/101_api_key_usage_limits.sql b/src/lib/db/migrations/101_api_key_usage_limits.sql new file mode 100644 index 0000000000..2d2d271f3e --- /dev/null +++ b/src/lib/db/migrations/101_api_key_usage_limits.sql @@ -0,0 +1,4 @@ +-- Migration: per-API-key USD usage limits +ALTER TABLE api_keys ADD COLUMN usage_limit_enabled INTEGER NOT NULL DEFAULT 0; +ALTER TABLE api_keys ADD COLUMN daily_usage_limit_usd REAL; +ALTER TABLE api_keys ADD COLUMN weekly_usage_limit_usd REAL; diff --git a/src/lib/usage/apiKeyUsageLimits.ts b/src/lib/usage/apiKeyUsageLimits.ts new file mode 100644 index 0000000000..166c717bf4 --- /dev/null +++ b/src/lib/usage/apiKeyUsageLimits.ts @@ -0,0 +1,354 @@ +import { getDbInstance } from "@/lib/db/core"; +import type { ProviderLimitsCacheEntry } from "@/lib/db/providerLimits"; +import { calculateCost } from "./costCalculator"; +import { buildErrorBody, sanitizeErrorMessage } from "@omniroute/open-sse/utils/error.ts"; + +const FORTALEZA_UTC_OFFSET_MS = 3 * 60 * 60 * 1000; +const WEEK_MS = 7 * 24 * 60 * 60 * 1000; + +export interface ApiKeyUsageLimitMetadata { + id: string; + allowedConnections?: string[] | null; + usageLimitEnabled?: boolean; + dailyUsageLimitUsd?: number | null; + weeklyUsageLimitUsd?: number | null; +} + +export interface ApiKeyUsageLimitStatus { + enabled: boolean; + dailyLimitUsd: number | null; + weeklyLimitUsd: number | null; + dailySpentUsd: number; + weeklySpentUsd: number; + dailyWindowStartIso: string; + weeklyWindowStartIso: string; + weeklyResetAtIso: string | null; + dailyExceeded: boolean; + weeklyExceeded: boolean; +} + +export interface ApiKeyUsageLimitDeps { + now?: () => number; + getProviderConnectionById?: (connectionId: string) => Promise; + getProviderConnections?: (filter?: Record) => Promise; + getProviderLimitsCache?: (connectionId: string) => ProviderLimitsCacheEntry | null; + getAllProviderLimitsCache?: () => Record; +} + +interface UsageCostRow { + provider: string | null; + model: string | null; + serviceTier: string | null; + promptTokens: number | null; + completionTokens: number | null; + cacheReadTokens: number | null; + cacheCreationTokens: number | null; + reasoningTokens: number | null; +} + +function toNumber(value: unknown): number { + if (typeof value === "number" && Number.isFinite(value)) return value; + if (typeof value === "string" && value.trim()) { + const parsed = Number(value); + return Number.isFinite(parsed) ? parsed : 0; + } + return 0; +} + +function asRecord(value: unknown): Record | null { + return value && typeof value === "object" && !Array.isArray(value) + ? (value as Record) + : null; +} + +function normalizeLimitUsd(value: unknown): number | null { + const numeric = toNumber(value); + return numeric > 0 ? numeric : null; +} + +function roundUsd(value: number): number { + return Math.round(value * 1_000_000) / 1_000_000; +} + +function formatUsd(value: number | null): string { + if (value === null || !Number.isFinite(value)) return "Not configured"; + return `$${value.toFixed(2)}`; +} + +export function getFortalezaDayStartIso(nowMs = Date.now()): string { + const fortalezaLocal = new Date(nowMs - FORTALEZA_UTC_OFFSET_MS); + return new Date( + Date.UTC( + fortalezaLocal.getUTCFullYear(), + fortalezaLocal.getUTCMonth(), + fortalezaLocal.getUTCDate(), + 3, + 0, + 0, + 0 + ) + ).toISOString(); +} + +export function getRollingWeekStartIso(nowMs = Date.now()): string { + return new Date(nowMs - WEEK_MS).toISOString(); +} + +function normalizeQuotaName(value: string): string { + return value + .toLowerCase() + .replace(/[^a-z0-9]+/g, " ") + .trim(); +} + +function findWeeklyQuotaResetAt(quotas: unknown, nowMs: number): string | null { + const quotaEntries: Array<[string, Record]> = []; + if (Array.isArray(quotas)) { + for (const item of quotas) { + const quota = asRecord(item); + if (!quota) continue; + const name = typeof quota.name === "string" ? quota.name : ""; + quotaEntries.push([name, quota]); + } + } else { + const quotaMap = asRecord(quotas); + if (quotaMap) { + for (const [name, value] of Object.entries(quotaMap)) { + const quota = asRecord(value); + if (quota) quotaEntries.push([name, quota]); + } + } + } + + for (const [name, quota] of quotaEntries) { + const label = normalizeQuotaName(`${name} ${typeof quota.name === "string" ? quota.name : ""}`); + if (!label) continue; + const isWeekly = label.includes("weekly") || label.includes("7d"); + if (!isWeekly || label.includes("sonnet")) continue; + const resetAt = typeof quota.resetAt === "string" && quota.resetAt.trim() ? quota.resetAt : ""; + const resetMs = Date.parse(resetAt); + if (Number.isFinite(resetMs) && resetMs > nowMs) { + return new Date(resetMs).toISOString(); + } + } + + return null; +} + +function connectionFromValue(value: unknown): { id: string; provider: string } | null { + const record = asRecord(value); + if (!record) return null; + const id = typeof record.id === "string" ? record.id : ""; + const provider = typeof record.provider === "string" ? record.provider : ""; + if (!id || !provider || record.isActive === false) return null; + return { id, provider }; +} + +async function resolveDeps(deps: ApiKeyUsageLimitDeps): Promise> { + const providers = + deps.getProviderConnectionById && deps.getProviderConnections + ? null + : await import("@/lib/db/providers"); + const providerLimits = + deps.getProviderLimitsCache && deps.getAllProviderLimitsCache + ? null + : await import("@/lib/db/providerLimits"); + + return { + now: deps.now ?? Date.now, + getProviderConnectionById: + deps.getProviderConnectionById ?? providers!.getProviderConnectionById, + getProviderConnections: deps.getProviderConnections ?? providers!.getProviderConnections, + getProviderLimitsCache: deps.getProviderLimitsCache ?? providerLimits!.getProviderLimitsCache, + getAllProviderLimitsCache: + deps.getAllProviderLimitsCache ?? providerLimits!.getAllProviderLimitsCache, + }; +} + +async function getProviderWeeklyResetAt( + metadata: ApiKeyUsageLimitMetadata, + deps: Required, + nowMs: number +): Promise { + const allowedConnections = Array.isArray(metadata.allowedConnections) + ? metadata.allowedConnections.filter((id) => typeof id === "string" && id.trim()) + : []; + + const resetCandidates: string[] = []; + if (allowedConnections.length > 0) { + for (const connectionId of allowedConnections) { + const connection = connectionFromValue(await deps.getProviderConnectionById(connectionId)); + if (!connection || connection.provider.toLowerCase() !== "claude") continue; + const resetAt = findWeeklyQuotaResetAt( + deps.getProviderLimitsCache(connection.id)?.quotas, + nowMs + ); + if (resetAt) resetCandidates.push(resetAt); + } + } else { + const caches = deps.getAllProviderLimitsCache(); + const connections = await deps.getProviderConnections({ isActive: true }); + for (const rawConnection of connections) { + const connection = connectionFromValue(rawConnection); + if (!connection || connection.provider.toLowerCase() !== "claude") continue; + const resetAt = findWeeklyQuotaResetAt(caches[connection.id]?.quotas, nowMs); + if (resetAt) resetCandidates.push(resetAt); + } + } + + return resetCandidates.sort((left, right) => Date.parse(left) - Date.parse(right)).at(0) ?? null; +} + +async function getApiKeyUsdSpendSince(apiKeyId: string, sinceIso: string): Promise { + if (!apiKeyId) return 0; + const db = getDbInstance(); + const rows = db + .prepare( + ` + SELECT + LOWER(provider) as provider, + LOWER(model) as model, + COALESCE(NULLIF(service_tier, ''), 'standard') as serviceTier, + COALESCE(SUM(tokens_input), 0) as promptTokens, + COALESCE(SUM(tokens_output), 0) as completionTokens, + COALESCE(SUM(tokens_cache_read), 0) as cacheReadTokens, + COALESCE(SUM(tokens_cache_creation), 0) as cacheCreationTokens, + COALESCE(SUM(tokens_reasoning), 0) as reasoningTokens + FROM usage_history + WHERE api_key_id = @apiKeyId + AND timestamp >= @sinceIso + AND success = 1 + GROUP BY LOWER(provider), LOWER(model), serviceTier + ` + ) + .all({ apiKeyId, sinceIso }) as UsageCostRow[]; + + let total = 0; + for (const row of rows) { + const provider = typeof row.provider === "string" ? row.provider : ""; + const model = typeof row.model === "string" ? row.model : ""; + if (!provider || !model) continue; + + total += await calculateCost( + provider, + model, + { + input: toNumber(row.promptTokens), + output: toNumber(row.completionTokens), + cacheRead: toNumber(row.cacheReadTokens), + cacheCreation: toNumber(row.cacheCreationTokens), + reasoning: toNumber(row.reasoningTokens), + }, + { + provider, + model, + serviceTier: row.serviceTier || "standard", + } + ); + } + + return roundUsd(total); +} + +export async function getApiKeyUsageLimitStatus( + metadata: ApiKeyUsageLimitMetadata, + deps: ApiKeyUsageLimitDeps = {} +): Promise { + const resolvedDeps = await resolveDeps(deps); + const now = resolvedDeps.now(); + const dailyWindowStartIso = getFortalezaDayStartIso(now); + const weeklyResetAtIso = await getProviderWeeklyResetAt(metadata, resolvedDeps, now); + const weeklyWindowStartIso = weeklyResetAtIso + ? new Date(Date.parse(weeklyResetAtIso) - WEEK_MS).toISOString() + : getRollingWeekStartIso(now); + const dailyLimitUsd = normalizeLimitUsd(metadata.dailyUsageLimitUsd); + const weeklyLimitUsd = normalizeLimitUsd(metadata.weeklyUsageLimitUsd); + const enabled = metadata.usageLimitEnabled === true; + + const [dailySpentUsd, weeklySpentUsd] = await Promise.all([ + getApiKeyUsdSpendSince(metadata.id, dailyWindowStartIso), + getApiKeyUsdSpendSince(metadata.id, weeklyWindowStartIso), + ]); + + return { + enabled, + dailyLimitUsd, + weeklyLimitUsd, + dailySpentUsd, + weeklySpentUsd, + dailyWindowStartIso, + weeklyWindowStartIso, + weeklyResetAtIso, + dailyExceeded: enabled && dailyLimitUsd !== null && dailySpentUsd >= dailyLimitUsd, + weeklyExceeded: enabled && weeklyLimitUsd !== null && weeklySpentUsd >= weeklyLimitUsd, + }; +} + +export function buildApiKeyUsageLimitText(status: ApiKeyUsageLimitStatus): string { + return [ + "Cota diaria", + formatUsd(status.dailyLimitUsd), + "Gasto diario", + formatUsd(status.dailySpentUsd), + "", + "Cota semanal", + formatUsd(status.weeklyLimitUsd), + "Gasto semanal", + formatUsd(status.weeklySpentUsd), + ].join("\n"); +} + +function buildUsageLimitExceededMessage(status: ApiKeyUsageLimitStatus): string { + if (status.dailyExceeded && status.dailyLimitUsd !== null) { + return `This API key reached its daily USD usage quota (${formatUsd(status.dailySpentUsd)} of ${formatUsd(status.dailyLimitUsd)}). Choose another allowed model or wait for quota reset.`; + } + if (status.weeklyExceeded && status.weeklyLimitUsd !== null) { + return `This API key reached its weekly USD usage quota (${formatUsd(status.weeklySpentUsd)} of ${formatUsd(status.weeklyLimitUsd)}). Choose another allowed model or wait for quota reset.`; + } + return "This API key reached its USD usage quota. Choose another allowed model or wait for quota reset."; +} + +function isAnthropicMessagesRequest(request: Request): boolean { + if (request.headers.has("anthropic-version")) return true; + try { + return new URL(request.url).pathname.endsWith("/v1/messages"); + } catch { + return false; + } +} + +export function buildApiKeyUsageLimitRejection( + request: Request, + status: ApiKeyUsageLimitStatus +): Response { + const message = sanitizeErrorMessage(buildUsageLimitExceededMessage(status)); + if (isAnthropicMessagesRequest(request)) { + return new Response( + JSON.stringify({ + type: "error", + error: { + type: "invalid_request_error", + message, + }, + }), + { + status: 400, + headers: { "Content-Type": "application/json" }, + } + ); + } + + return new Response(JSON.stringify(buildErrorBody(400, message)), { + status: 400, + headers: { "Content-Type": "application/json" }, + }); +} + +export async function buildApiKeyUsageLimitPolicyRejection( + request: Request, + metadata: ApiKeyUsageLimitMetadata +): Promise { + const status = await getApiKeyUsageLimitStatus(metadata); + if (!status.enabled || (!status.dailyExceeded && !status.weeklyExceeded)) return null; + return buildApiKeyUsageLimitRejection(request, status); +} diff --git a/src/lib/usage/internalUsageCommand.ts b/src/lib/usage/internalUsageCommand.ts index 03eae21532..a76b307b49 100644 --- a/src/lib/usage/internalUsageCommand.ts +++ b/src/lib/usage/internalUsageCommand.ts @@ -1,4 +1,8 @@ import type { ProviderLimitsCacheEntry } from "@/lib/db/providerLimits"; +import { + buildApiKeyUsageLimitText, + type ApiKeyUsageLimitStatus, +} from "@/lib/usage/apiKeyUsageLimits"; export const INTERNAL_USAGE_COMMAND = "@@om-usage"; export const USAGE_COMMAND_DISABLED_MESSAGE = "Usage command is disabled for this API key."; @@ -12,6 +16,9 @@ interface UsageCommandApiKeyMetadata { name?: string; allowedConnections?: string[] | null; allowUsageCommand?: boolean; + usageLimitEnabled?: boolean; + dailyUsageLimitUsd?: number | null; + weeklyUsageLimitUsd?: number | null; } interface ProviderConnectionLike { @@ -35,6 +42,10 @@ export interface InternalUsageCommandDeps { getProviderConnections?: (filter?: JsonRecord) => Promise; getProviderLimitsCache?: (connectionId: string) => ProviderLimitsCacheEntry | null; getAllProviderLimitsCache?: () => Record; + getApiKeyUsageLimitStatus?: ( + metadata: UsageCommandApiKeyMetadata, + deps?: { now?: () => number } + ) => Promise; } type RequiredDeps = Required; @@ -50,6 +61,9 @@ async function normalizeDeps(deps: InternalUsageCommandDeps = {}): Promise { const resolvedDeps = await normalizeDeps(deps); + if (metadata.usageLimitEnabled === true) { + return buildApiKeyUsageLimitText( + await resolvedDeps.getApiKeyUsageLimitStatus(metadata, { now: resolvedDeps.now }) + ); + } + const snapshot = selectUsageSnapshot(await collectUsageSnapshots(metadata, resolvedDeps)); if (!snapshot) { diff --git a/src/shared/utils/apiKeyPolicy.ts b/src/shared/utils/apiKeyPolicy.ts index 77b207a8dd..7e93fb027d 100644 --- a/src/shared/utils/apiKeyPolicy.ts +++ b/src/shared/utils/apiKeyPolicy.ts @@ -30,6 +30,7 @@ import { checkRateLimit, RateLimitRule } from "./rateLimiter"; import { resolveEndpointCategory } from "@/shared/constants/endpointCategories"; import { resolveQuotaKeyScope } from "@/lib/quota/quotaKey"; import { isQuotaModelName, parseQuotaModelName } from "@/lib/quota/quotaModelNaming"; +import { buildApiKeyUsageLimitPolicyRejection } from "@/lib/usage/apiKeyUsageLimits"; // Default to no per-key request cap. API keys can still opt into explicit // limits via Settings/API Keys, while provider/account quota controls remain @@ -92,6 +93,9 @@ export interface ApiKeyMetadata { allowedEndpoints?: string[]; disableNonPublicModels?: boolean; allowUsageCommand?: boolean; + usageLimitEnabled?: boolean; + dailyUsageLimitUsd?: number | null; + weeklyUsageLimitUsd?: number | null; } /** @@ -379,6 +383,31 @@ export async function enforceApiKeyPolicy( } } + // ── Check 2.1: per-key USD fair usage cap ── + if (apiKeyInfo.usageLimitEnabled === true) { + try { + const usageLimitRejection = await buildApiKeyUsageLimitPolicyRejection(request, { + id: apiKeyInfo.id, + usageLimitEnabled: apiKeyInfo.usageLimitEnabled, + dailyUsageLimitUsd: apiKeyInfo.dailyUsageLimitUsd, + weeklyUsageLimitUsd: apiKeyInfo.weeklyUsageLimitUsd, + }); + if (usageLimitRejection) { + return { apiKey, apiKeyInfo, rejection: usageLimitRejection }; + } + } catch (error) { + log.error("API_POLICY", "API key USD usage limit check failed. Request blocked.", { error }); + return { + apiKey, + apiKeyInfo, + rejection: errorResponse( + HTTP_STATUS.SERVICE_UNAVAILABLE, + "API key usage limit unavailable" + ), + }; + } + } + // ── Check 2.5: Endpoint restriction ── if (apiKeyInfo.allowedEndpoints && apiKeyInfo.allowedEndpoints.length > 0) { try { diff --git a/src/shared/validation/schemas/keys.ts b/src/shared/validation/schemas/keys.ts index 119287297b..e3db9b7928 100644 --- a/src/shared/validation/schemas/keys.ts +++ b/src/shared/validation/schemas/keys.ts @@ -22,6 +22,9 @@ export const createKeySchema = z.object({ name: z.string().min(1, "Name is required").max(200), noLog: z.boolean().optional(), allowUsageCommand: z.boolean().optional(), + usageLimitEnabled: z.boolean().optional(), + dailyUsageLimitUsd: z.coerce.number().min(0).optional().nullable(), + weeklyUsageLimitUsd: z.coerce.number().min(0).optional().nullable(), scopes: z.array(z.string().trim().min(1).max(64)).max(32).optional(), }); @@ -104,6 +107,9 @@ export const updateKeyPermissionsSchema = z streamDefaultMode: z.enum(["legacy", "json"]).optional(), disableNonPublicModels: z.boolean().optional(), allowUsageCommand: z.boolean().optional(), + usageLimitEnabled: z.boolean().optional(), + dailyUsageLimitUsd: z.coerce.number().min(0).optional().nullable(), + weeklyUsageLimitUsd: z.coerce.number().min(0).optional().nullable(), }) .superRefine((value, ctx) => { if ( @@ -124,7 +130,10 @@ export const updateKeyPermissionsSchema = z value.allowedEndpoints === undefined && value.streamDefaultMode === undefined && value.disableNonPublicModels === undefined && - value.allowUsageCommand === undefined + value.allowUsageCommand === undefined && + value.usageLimitEnabled === undefined && + value.dailyUsageLimitUsd === undefined && + value.weeklyUsageLimitUsd === undefined ) { ctx.addIssue({ code: z.ZodIssueCode.custom, @@ -132,4 +141,4 @@ export const updateKeyPermissionsSchema = z path: [], }); } - }); \ No newline at end of file + }); diff --git a/tests/unit/api-key-usage-limits.test.ts b/tests/unit/api-key-usage-limits.test.ts new file mode 100644 index 0000000000..a928179c0a --- /dev/null +++ b/tests/unit/api-key-usage-limits.test.ts @@ -0,0 +1,197 @@ +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-api-key-usage-limits-")); +process.env.DATA_DIR = TEST_DATA_DIR; +process.env.API_KEY_SECRET = process.env.API_KEY_SECRET || "usage-limit-test-secret"; + +const core = await import("../../src/lib/db/core.ts"); +const localDb = await import("../../src/lib/localDb.ts"); +const apiKeysDb = await import("../../src/lib/db/apiKeys.ts"); +const usageHistory = await import("../../src/lib/usage/usageHistory.ts"); +const usageLimits = await import("../../src/lib/usage/apiKeyUsageLimits.ts"); + +const NOW = Date.parse("2026-06-19T20:00:00.000Z"); + +async function resetStorage() { + core.resetDbInstance(); + apiKeysDb.resetApiKeyState(); + usageHistory.clearPendingRequests(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); +} + +test.beforeEach(async () => { + await resetStorage(); +}); + +test.after(() => { + core.resetDbInstance(); + apiKeysDb.resetApiKeyState(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +test("API key USD usage limits persist and default off", async () => { + const created = await apiKeysDb.createApiKey("Usage Limit Key", "machine-limit-01"); + + let metadata = await apiKeysDb.getApiKeyMetadata(created.key); + assert.equal(metadata?.usageLimitEnabled, false); + assert.equal(metadata?.dailyUsageLimitUsd, null); + assert.equal(metadata?.weeklyUsageLimitUsd, null); + + await apiKeysDb.updateApiKeyPermissions(created.id, { + usageLimitEnabled: true, + dailyUsageLimitUsd: 10.5, + weeklyUsageLimitUsd: 50, + }); + apiKeysDb.clearApiKeyCaches(); + + metadata = await apiKeysDb.getApiKeyMetadata(created.key); + assert.equal(metadata?.usageLimitEnabled, true); + assert.equal(metadata?.dailyUsageLimitUsd, 10.5); + assert.equal(metadata?.weeklyUsageLimitUsd, 50); +}); + +test("getApiKeyUsageLimitStatus aligns weekly USD spend with provider resetAt when available", async () => { + await localDb.updatePricing({ + claude: { + "claude-opus-4-8": { + input: 1, + cached: 1, + output: 1, + reasoning: 1, + cache_creation: 1, + }, + }, + }); + + const created = await apiKeysDb.createApiKey("Metered Key", "machine-limit-02"); + await apiKeysDb.updateApiKeyPermissions(created.id, { + usageLimitEnabled: true, + dailyUsageLimitUsd: 10, + weeklyUsageLimitUsd: 20, + }); + + await usageHistory.saveRequestUsage({ + provider: "claude", + model: "claude-opus-4-8", + apiKeyId: created.id, + apiKeyName: "Metered Key", + tokens: { input: 2_000_000, output: 0 }, + success: true, + timestamp: "2026-06-19T12:00:00.000Z", + }); + await usageHistory.saveRequestUsage({ + provider: "claude", + model: "claude-opus-4-8", + apiKeyId: created.id, + apiKeyName: "Metered Key", + tokens: { input: 3_000_000, output: 0 }, + success: true, + timestamp: "2026-06-18T21:00:00.000Z", + }); + await usageHistory.saveRequestUsage({ + provider: "claude", + model: "claude-opus-4-8", + apiKeyId: created.id, + apiKeyName: "Metered Key", + tokens: { input: 7_000_000, output: 0 }, + success: true, + timestamp: "2026-06-18T12:00:00.000Z", + }); + + const metadata = await apiKeysDb.getApiKeyMetadata(created.key); + assert.ok(metadata); + + const weeklyResetAt = "2026-06-25T20:00:00.000Z"; + const status = await usageLimits.getApiKeyUsageLimitStatus( + { ...metadata, allowedConnections: ["conn-claude"] }, + { + now: () => NOW, + getProviderConnectionById: async () => ({ + id: "conn-claude", + provider: "claude", + isActive: true, + }), + getProviderConnections: async () => [], + getProviderLimitsCache: () => ({ + plan: "Claude Max", + quotas: { + "weekly (7d)": { + used: 27, + total: 100, + resetAt: weeklyResetAt, + }, + }, + message: null, + fetchedAt: new Date(NOW).toISOString(), + }), + getAllProviderLimitsCache: () => ({}), + } + ); + + assert.equal(status.enabled, true); + assert.equal(status.dailySpentUsd, 2); + assert.equal(status.weeklySpentUsd, 5); + assert.equal(status.dailyLimitUsd, 10); + assert.equal(status.weeklyLimitUsd, 20); + assert.equal(status.weeklyWindowStartIso, "2026-06-18T20:00:00.000Z"); + assert.equal(status.weeklyResetAtIso, weeklyResetAt); + assert.equal(status.dailyExceeded, false); + assert.equal(status.weeklyExceeded, false); +}); + +test("buildApiKeyUsageLimitText returns only the quota and spent USD lines", async () => { + const text = usageLimits.buildApiKeyUsageLimitText({ + enabled: true, + dailyLimitUsd: 10, + weeklyLimitUsd: 50, + dailySpentUsd: 2, + weeklySpentUsd: 5.25, + dailyWindowStartIso: "2026-06-19T03:00:00.000Z", + weeklyWindowStartIso: "2026-06-12T20:00:00.000Z", + weeklyResetAtIso: "2026-06-19T20:00:00.000Z", + dailyExceeded: false, + weeklyExceeded: false, + }); + + assert.equal( + text, + [ + "Cota diaria", + "$10.00", + "Gasto diario", + "$2.00", + "", + "Cota semanal", + "$50.00", + "Gasto semanal", + "$5.25", + ].join("\n") + ); +}); + +test("buildApiKeyUsageLimitRejection uses 400 so Claude Code does not trigger login", () => { + const response = usageLimits.buildApiKeyUsageLimitRejection( + new Request("http://localhost/v1/messages", { + headers: { "anthropic-version": "2023-06-01" }, + }), + { + enabled: true, + dailyLimitUsd: 10, + weeklyLimitUsd: 50, + dailySpentUsd: 12, + weeklySpentUsd: 20, + dailyWindowStartIso: "2026-06-19T03:00:00.000Z", + weeklyWindowStartIso: "2026-06-12T20:00:00.000Z", + weeklyResetAtIso: "2026-06-19T20:00:00.000Z", + dailyExceeded: true, + weeklyExceeded: false, + } + ); + + assert.equal(response.status, 400); +}); diff --git a/tests/unit/internal-usage-command.test.ts b/tests/unit/internal-usage-command.test.ts index 35f24ed3de..0188f638b8 100644 --- a/tests/unit/internal-usage-command.test.ts +++ b/tests/unit/internal-usage-command.test.ts @@ -113,6 +113,61 @@ test("buildUsageCommandText formats cached Claude usage windows exactly", async ); }); +test("buildUsageCommandText formats API key USD limits when fair usage is enabled", async () => { + const text = await buildUsageCommandText( + { + id: "key-limited", + name: "limited", + allowedConnections: ["conn-claude"], + usageLimitEnabled: true, + dailyUsageLimitUsd: 10, + weeklyUsageLimitUsd: 50, + }, + { + now: () => NOW, + getApiKeyUsageLimitStatus: async () => ({ + enabled: true, + dailyLimitUsd: 10, + weeklyLimitUsd: 50, + dailySpentUsd: 2, + weeklySpentUsd: 5.25, + dailyWindowStartIso: "2026-06-16T03:00:00.000Z", + weeklyWindowStartIso: "2026-06-09T12:00:00.000Z", + weeklyResetAtIso: "2026-06-16T12:00:00.000Z", + dailyExceeded: false, + weeklyExceeded: false, + }), + getProviderConnectionById: async () => { + throw new Error("provider connection lookup must not run for fair usage output"); + }, + getProviderConnections: async () => { + throw new Error("provider connection lookup must not run for fair usage output"); + }, + getProviderLimitsCache: () => null, + getAllProviderLimitsCache: () => { + throw new Error("provider cache lookup must not run for fair usage output"); + }, + isValidApiKey: async () => true, + getApiKeyMetadata: async () => null, + } + ); + + assert.equal( + text, + [ + "Cota diaria", + "$10.00", + "Gasto diario", + "$2.00", + "", + "Cota semanal", + "$50.00", + "Gasto semanal", + "$5.25", + ].join("\n") + ); +}); + test("handleInternalUsageCommand returns disabled response locally without provider routing", async () => { const response = await handleInternalUsageCommand( new Request("http://localhost/v1/chat/completions", {