@@ -205,13 +333,18 @@ export function KeysTable({ keys, onRevoke, onEdit }: KeysTableProps) {
return (
-
setSearchQuery(e.target.value)}
- className="max-w-md"
- />
+
+ setSearchQuery(e.target.value)}
+ className="max-w-sm"
+ />
+
+ {filteredKeys.length} / {keys.length}
+
+
{isMobileLayout ? (
{filteredKeys.map((key) => (
@@ -224,9 +357,16 @@ export function KeysTable({ keys, onRevoke, onEdit }: KeysTableProps) {
{key.name}
{key.description || "-"}
-
- {key.is_active ? t("enabled") : t("disabled")}
-
+
+
+ {key.is_active ? t("enabled") : t("disabled")}
+
+ {key.is_quota_exceeded ? (
+
+ {t("quotaExceeded")}
+
+ ) : null}
+
@@ -296,6 +436,13 @@ export function KeysTable({ keys, onRevoke, onEdit }: KeysTableProps) {
+ {key.spending_rules && key.spending_rules.length > 0 ? (
+
- {filteredKeys.map((key) => (
-
-
-
-
{key.name}
-
{
+ const hasQuota = !!(key.spending_rules && key.spending_rules.length > 0);
+ const isExpanded = expandedKeys.has(key.id);
+ return (
+
+ td]:border-b-0")}>
+ hasQuota && toggleExpand(key.id)}
>
- {key.is_active ? t("enabled") : t("disabled")}
-
-
-
- {visibleKeyIds.has(key.id)
- ? revealedKeys.get(key.id) || key.key_prefix
- : maskKey(key.key_prefix)}
-
-
-
-
-
- {key.description || -}
-
-
-
- {formatAccessModeLabel(key, t)}
-
-
-
- {formatExpiry(key.expires_at)}
-
-
- {formatDistanceToNow(new Date(key.created_at), {
- addSuffix: true,
- locale: dateLocale,
- })}
-
-
-
-
- {
- await handleToggleKeyActive(key, nextActive);
- }}
- disabled={
- toggleActiveMutation.isPending &&
- toggleActiveMutation.variables?.id === key.id
- }
- className="h-5 w-10"
- aria-label={`${key.is_active ? t("quickDisable") : t("quickEnable")}: ${key.name}`}
- />
+ {key.name}
+
+ {key.is_quota_exceeded ? (
+
+ {t("quotaExceeded")}
+
+ ) : null}
+
+
+
+
+
+ {visibleKeyIds.has(key.id)
+ ? revealedKeys.get(key.id) || key.key_prefix
+ : maskKey(key.key_prefix)}
+
+
+
+
+
+
+
+ {key.description || -}
+
+
- {key.is_active ? t("enabled") : t("disabled")}
+ {formatAccessModeLabel(key, t)}
-
+
+ {
+ await handleToggleKeyActive(key, nextActive);
+ }}
+ disabled={
+ toggleActiveMutation.isPending &&
+ toggleActiveMutation.variables?.id === key.id
+ }
+ className="h-5 w-10"
+ aria-label={`${key.is_active ? t("quickDisable") : t("quickEnable")}: ${key.name}`}
+ />
+
+ {key.is_active ? t("enabled") : t("disabled")}
+
+
+
+
+
diff --git a/src/lib/db/schema-pg.ts b/src/lib/db/schema-pg.ts
index 4788eca7..ac064d0a 100644
--- a/src/lib/db/schema-pg.ts
+++ b/src/lib/db/schema-pg.ts
@@ -27,6 +27,10 @@ export const apiKeys = pgTable(
description: text("description"),
userId: uuid("user_id"), // Reserved for future user system
accessMode: varchar("access_mode", { length: 16 }).notNull().default("unrestricted"),
+ spendingRules: json("spending_rules").$type<
+ | { period_type: "daily" | "monthly" | "rolling"; limit: number; period_hours?: number }[]
+ | null
+ >(),
isActive: boolean("is_active").notNull().default(true),
expiresAt: timestamp("expires_at", { withTimezone: true }),
createdAt: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(),
diff --git a/src/lib/db/schema-sqlite.ts b/src/lib/db/schema-sqlite.ts
index 5792de06..8f47ccce 100644
--- a/src/lib/db/schema-sqlite.ts
+++ b/src/lib/db/schema-sqlite.ts
@@ -18,6 +18,10 @@ export const apiKeys = sqliteTable(
description: text("description"),
userId: text("user_id"), // Reserved for future user system
accessMode: text("access_mode").notNull().default("unrestricted"),
+ spendingRules: text("spending_rules", { mode: "json" }).$type<
+ | { period_type: "daily" | "monthly" | "rolling"; limit: number; period_hours?: number }[]
+ | null
+ >(),
isActive: integer("is_active", { mode: "boolean" }).notNull().default(true),
expiresAt: integer("expires_at", { mode: "timestamp_ms" }),
createdAt: integer("created_at", { mode: "timestamp_ms" }).notNull().defaultNow(),
diff --git a/src/lib/services/api-key-quota-tracker.ts b/src/lib/services/api-key-quota-tracker.ts
new file mode 100644
index 00000000..aa18e8fa
--- /dev/null
+++ b/src/lib/services/api-key-quota-tracker.ts
@@ -0,0 +1,402 @@
+import { and, eq, gte, isNotNull, lt, sum } from "drizzle-orm";
+import { apiKeys, db, requestBillingSnapshots, type ApiKey } from "@/lib/db";
+import {
+ type RuleStatus,
+ type SpendingRule,
+ getPeriodStartForRule,
+ getResetsAtForRule,
+ getRollingWindowStart,
+} from "@/lib/services/upstream-quota-tracker";
+import { normalizeSpendingRules } from "@/lib/services/spending-rules";
+
+interface ApiKeyQuotaStatus {
+ apiKeyId: string;
+ apiKeyName: string;
+ rules: RuleStatus[];
+ isExceeded: boolean;
+}
+
+interface RuleCacheEntry {
+ periodType: SpendingRule["period_type"];
+ periodHours: number | null;
+ limit: number;
+ currentSpending: number;
+ lastSyncedAt: Date;
+}
+
+const NORMAL_SYNC_INTERVAL_MS = 5 * 60 * 1000;
+const URGENT_SYNC_INTERVAL_MS = 1 * 60 * 1000;
+const URGENT_THRESHOLD_PERCENT = 80;
+const HOUR_MS = 60 * 60 * 1000;
+
+function extractSpendingRulesFromApiKey(apiKey: Pick
): SpendingRule[] {
+ return normalizeSpendingRules(apiKey.spendingRules) ?? [];
+}
+
+class ApiKeyQuotaTracker {
+ private cache = new Map();
+ private rulesCache = new Map();
+ private nameCache = new Map();
+ private syncTimer: ReturnType | null = null;
+ private initialized = false;
+ private initPromise: Promise | null = null;
+
+ async initialize(): Promise {
+ if (this.initialized) {
+ return;
+ }
+ if (!this.initPromise) {
+ this.initPromise = (async () => {
+ await this.syncFromDb();
+ this.startSyncTimer();
+ this.initialized = true;
+ })().catch((error) => {
+ this.initPromise = null;
+ throw error;
+ });
+ }
+
+ await this.initPromise;
+ }
+
+ isWithinQuota(apiKeyId: string): boolean {
+ const rules = this.rulesCache.get(apiKeyId);
+ if (!rules || rules.length === 0) {
+ return true;
+ }
+
+ const entries = this.cache.get(apiKeyId);
+ if (!entries || entries.length === 0) {
+ return true;
+ }
+
+ return entries.every((entry) => entry.currentSpending < entry.limit);
+ }
+
+ adjustSpending(apiKeyId: string, delta: number): void {
+ if (!Number.isFinite(delta) || delta === 0) {
+ return;
+ }
+
+ const entries = this.cache.get(apiKeyId);
+ if (!entries) {
+ return;
+ }
+
+ for (const entry of entries) {
+ const nextSpending = entry.currentSpending + delta;
+ entry.currentSpending = nextSpending > 0 ? Number(nextSpending.toFixed(10)) : 0;
+ }
+ }
+
+ getQuotaStatus(apiKeyId: string): ApiKeyQuotaStatus | null {
+ const rules = this.rulesCache.get(apiKeyId);
+ if (!rules || rules.length === 0) {
+ return null;
+ }
+
+ const entries = this.cache.get(apiKeyId);
+ const now = new Date();
+ let isExceeded = false;
+
+ const statuses: RuleStatus[] = rules.map((rule) => {
+ const entry = entries?.find(
+ (candidate) =>
+ candidate.periodType === rule.period_type &&
+ candidate.periodHours === (rule.period_hours ?? null)
+ );
+ const currentSpending = entry?.currentSpending ?? 0;
+ const percentUsed = rule.limit > 0 ? Math.min(999, (currentSpending / rule.limit) * 100) : 0;
+ const ruleExceeded = currentSpending >= rule.limit;
+ if (ruleExceeded) {
+ isExceeded = true;
+ }
+
+ return {
+ periodType: rule.period_type,
+ periodHours: rule.period_hours ?? null,
+ currentSpending,
+ spendingLimit: rule.limit,
+ percentUsed: Number(percentUsed.toFixed(1)),
+ isExceeded: ruleExceeded,
+ resetsAt: getResetsAtForRule(rule, now),
+ estimatedRecoveryAt: null,
+ };
+ });
+
+ return {
+ apiKeyId,
+ apiKeyName: this.nameCache.get(apiKeyId) ?? apiKeyId,
+ rules: statuses,
+ isExceeded,
+ };
+ }
+
+ async syncFromDb(): Promise {
+ const quotaKeys = await db.query.apiKeys.findMany({
+ where: isNotNull(apiKeys.spendingRules),
+ columns: {
+ id: true,
+ name: true,
+ spendingRules: true,
+ },
+ });
+
+ const activeIds = new Set();
+ for (const key of quotaKeys) {
+ const rules = extractSpendingRulesFromApiKey(key);
+ if (rules.length === 0) {
+ continue;
+ }
+
+ this.rulesCache.set(key.id, rules);
+ this.nameCache.set(key.id, key.name);
+ activeIds.add(key.id);
+ }
+
+ for (const id of this.rulesCache.keys()) {
+ if (!activeIds.has(id)) {
+ this.rulesCache.delete(id);
+ this.cache.delete(id);
+ this.nameCache.delete(id);
+ }
+ }
+
+ const now = new Date();
+ for (const key of quotaKeys) {
+ const rules = this.rulesCache.get(key.id);
+ if (!rules) {
+ continue;
+ }
+
+ const entries: RuleCacheEntry[] = [];
+ for (const rule of rules) {
+ const periodStart = getPeriodStartForRule(rule, now);
+ const [row] = await db
+ .select({ totalCost: sum(requestBillingSnapshots.finalCost) })
+ .from(requestBillingSnapshots)
+ .where(
+ and(
+ eq(requestBillingSnapshots.apiKeyId, key.id),
+ eq(requestBillingSnapshots.billingStatus, "billed"),
+ gte(requestBillingSnapshots.billedAt, periodStart)
+ )
+ );
+
+ const totalCost = row?.totalCost ? Number(row.totalCost) : 0;
+ entries.push({
+ periodType: rule.period_type,
+ periodHours: rule.period_hours ?? null,
+ limit: rule.limit,
+ currentSpending: Number.isNaN(totalCost) ? 0 : totalCost,
+ lastSyncedAt: now,
+ });
+ }
+
+ this.cache.set(key.id, entries);
+ }
+ }
+
+ async syncApiKeyFromDb(
+ apiKeyId: string,
+ apiKeyName: string,
+ spendingRules: SpendingRule[] | null
+ ): Promise {
+ const rules = normalizeSpendingRules(spendingRules);
+ if (!rules || rules.length === 0) {
+ this.rulesCache.delete(apiKeyId);
+ this.cache.delete(apiKeyId);
+ this.nameCache.delete(apiKeyId);
+ return;
+ }
+
+ this.rulesCache.set(apiKeyId, rules);
+ this.nameCache.set(apiKeyId, apiKeyName);
+
+ const now = new Date();
+ const entries: RuleCacheEntry[] = [];
+ for (const rule of rules) {
+ const periodStart = getPeriodStartForRule(rule, now);
+ const [row] = await db
+ .select({ totalCost: sum(requestBillingSnapshots.finalCost) })
+ .from(requestBillingSnapshots)
+ .where(
+ and(
+ eq(requestBillingSnapshots.apiKeyId, apiKeyId),
+ eq(requestBillingSnapshots.billingStatus, "billed"),
+ gte(requestBillingSnapshots.billedAt, periodStart)
+ )
+ );
+
+ const totalCost = row?.totalCost ? Number(row.totalCost) : 0;
+ entries.push({
+ periodType: rule.period_type,
+ periodHours: rule.period_hours ?? null,
+ limit: rule.limit,
+ currentSpending: Number.isNaN(totalCost) ? 0 : totalCost,
+ lastSyncedAt: now,
+ });
+ }
+
+ this.cache.set(apiKeyId, entries);
+ }
+
+ async estimateRecoveryTime(apiKeyId: string, rule: SpendingRule): Promise {
+ if (rule.period_type !== "rolling") {
+ return null;
+ }
+
+ const entries = this.cache.get(apiKeyId);
+ const entry = entries?.find(
+ (candidate) =>
+ candidate.periodType === "rolling" && candidate.periodHours === (rule.period_hours ?? null)
+ );
+ if (!entry || entry.currentSpending < rule.limit) {
+ return null;
+ }
+
+ const configuredHours = rule.period_hours ?? 24;
+ const windowHours =
+ Number.isFinite(configuredHours) && configuredHours > 0 ? Math.floor(configuredHours) : 24;
+ const now = new Date();
+ const windowStart = getRollingWindowStart(windowHours, now);
+ const excessAmount = entry.currentSpending - rule.limit;
+ const scanHours = Math.min(windowHours, 8760);
+ if (scanHours <= 0) {
+ return null;
+ }
+
+ const rows = await db
+ .select({
+ billedAt: requestBillingSnapshots.billedAt,
+ finalCost: requestBillingSnapshots.finalCost,
+ })
+ .from(requestBillingSnapshots)
+ .where(
+ and(
+ eq(requestBillingSnapshots.apiKeyId, apiKeyId),
+ eq(requestBillingSnapshots.billingStatus, "billed"),
+ gte(requestBillingSnapshots.billedAt, windowStart),
+ lt(requestBillingSnapshots.billedAt, now)
+ )
+ );
+
+ const windowStartMs = windowStart.getTime();
+ const hourlySlideOutCosts = new Array(scanHours).fill(0);
+ for (const row of rows) {
+ const billedAtValue = row?.billedAt;
+ if (!billedAtValue) {
+ continue;
+ }
+
+ const billedAtMs =
+ billedAtValue instanceof Date ? billedAtValue.getTime() : new Date(billedAtValue).getTime();
+ if (!Number.isFinite(billedAtMs) || billedAtMs < windowStartMs) {
+ continue;
+ }
+
+ const cost = Number(row?.finalCost ?? 0);
+ if (!Number.isFinite(cost) || cost <= 0) {
+ continue;
+ }
+
+ const hourIndex = Math.floor((billedAtMs - windowStartMs) / HOUR_MS);
+ if (hourIndex < 0 || hourIndex >= scanHours) {
+ continue;
+ }
+
+ hourlySlideOutCosts[hourIndex] += cost;
+ }
+
+ let cumulativeSlideOut = 0;
+ for (let hourIndex = 0; hourIndex < scanHours; hourIndex += 1) {
+ cumulativeSlideOut += hourlySlideOutCosts[hourIndex] ?? 0;
+ if (cumulativeSlideOut > excessAmount) {
+ const sliceEnd = new Date(windowStartMs + (hourIndex + 1) * HOUR_MS);
+ return new Date(sliceEnd.getTime() + windowHours * HOUR_MS);
+ }
+ }
+
+ return null;
+ }
+
+ reset(): void {
+ this.stopSyncTimer();
+ this.cache.clear();
+ this.rulesCache.clear();
+ this.nameCache.clear();
+ this.initialized = false;
+ this.initPromise = null;
+ }
+
+ setRules(apiKeyId: string, rules: SpendingRule[], name?: string): void {
+ const normalizedRules = normalizeSpendingRules(rules) ?? [];
+ this.rulesCache.set(apiKeyId, normalizedRules);
+ if (name) {
+ this.nameCache.set(apiKeyId, name);
+ }
+ this.cache.set(
+ apiKeyId,
+ normalizedRules.map((rule) => ({
+ periodType: rule.period_type,
+ periodHours: rule.period_hours ?? null,
+ limit: rule.limit,
+ currentSpending: 0,
+ lastSyncedAt: new Date(),
+ }))
+ );
+ }
+
+ getCacheEntries(apiKeyId: string): RuleCacheEntry[] | undefined {
+ return this.cache.get(apiKeyId);
+ }
+
+ private startSyncTimer(): void {
+ if (this.syncTimer) {
+ return;
+ }
+ this.syncTimer = setInterval(() => {
+ this.tickSync().catch(() => undefined);
+ }, URGENT_SYNC_INTERVAL_MS);
+ }
+
+ private stopSyncTimer(): void {
+ if (!this.syncTimer) {
+ return;
+ }
+ clearInterval(this.syncTimer);
+ this.syncTimer = null;
+ }
+
+ private async tickSync(): Promise {
+ if (this.shouldSyncNow()) {
+ await this.syncFromDb();
+ }
+ }
+
+ private shouldSyncNow(): boolean {
+ const now = Date.now();
+ for (const [apiKeyId, entries] of this.cache) {
+ if (!this.rulesCache.has(apiKeyId)) {
+ continue;
+ }
+
+ for (const entry of entries) {
+ const elapsed = now - entry.lastSyncedAt.getTime();
+ const percentUsed = (entry.currentSpending / entry.limit) * 100;
+
+ if (percentUsed >= URGENT_THRESHOLD_PERCENT || entry.currentSpending >= entry.limit) {
+ if (elapsed >= URGENT_SYNC_INTERVAL_MS) {
+ return true;
+ }
+ } else if (elapsed >= NORMAL_SYNC_INTERVAL_MS) {
+ return true;
+ }
+ }
+ }
+
+ return false;
+ }
+}
+
+export const apiKeyQuotaTracker = new ApiKeyQuotaTracker();
diff --git a/src/lib/services/billing-cost-service.ts b/src/lib/services/billing-cost-service.ts
index 3f585b10..1b2c8430 100644
--- a/src/lib/services/billing-cost-service.ts
+++ b/src/lib/services/billing-cost-service.ts
@@ -7,6 +7,7 @@ import {
} from "@/lib/services/billing-price-service";
import { getProviderTypeForModel } from "@/lib/services/model-router";
import { quotaTracker } from "@/lib/services/upstream-quota-tracker";
+import { apiKeyQuotaTracker } from "@/lib/services/api-key-quota-tracker";
export type UnbillableReason =
| "model_missing"
@@ -47,6 +48,7 @@ interface NormalizedBillingUsage {
}
interface ExistingBillingSnapshot {
+ apiKeyId: string | null;
upstreamId: string | null;
billingStatus: string;
finalCost: number | null;
@@ -60,10 +62,12 @@ function parseSnapshotCost(value: number | null): number {
function applyQuotaDeltaAfterSnapshotUpsert(
previousSnapshot: ExistingBillingSnapshot | null,
nextBillingStatus: "billed" | "unbilled",
+ nextApiKeyId: string | null,
nextUpstreamId: string | null,
nextFinalCost: number | null
): void {
const deltaByUpstream = new Map();
+ const deltaByApiKey = new Map();
if (previousSnapshot?.billingStatus === "billed" && previousSnapshot.upstreamId) {
const previousCost = parseSnapshotCost(previousSnapshot.finalCost);
@@ -74,6 +78,15 @@ function applyQuotaDeltaAfterSnapshotUpsert(
);
}
}
+ if (previousSnapshot?.billingStatus === "billed" && previousSnapshot.apiKeyId) {
+ const previousCost = parseSnapshotCost(previousSnapshot.finalCost);
+ if (previousCost > 0) {
+ deltaByApiKey.set(
+ previousSnapshot.apiKeyId,
+ (deltaByApiKey.get(previousSnapshot.apiKeyId) ?? 0) - previousCost
+ );
+ }
+ }
const normalizedFinalCost = parseSnapshotCost(nextFinalCost);
if (nextBillingStatus === "billed" && nextUpstreamId && normalizedFinalCost > 0) {
@@ -82,6 +95,9 @@ function applyQuotaDeltaAfterSnapshotUpsert(
(deltaByUpstream.get(nextUpstreamId) ?? 0) + normalizedFinalCost
);
}
+ if (nextBillingStatus === "billed" && nextApiKeyId && normalizedFinalCost > 0) {
+ deltaByApiKey.set(nextApiKeyId, (deltaByApiKey.get(nextApiKeyId) ?? 0) + normalizedFinalCost);
+ }
for (const [upstreamId, delta] of deltaByUpstream) {
const normalizedDelta = Number(delta.toFixed(10));
@@ -89,6 +105,12 @@ function applyQuotaDeltaAfterSnapshotUpsert(
quotaTracker.adjustSpending(upstreamId, normalizedDelta);
}
}
+ for (const [apiKeyId, delta] of deltaByApiKey) {
+ const normalizedDelta = Number(delta.toFixed(10));
+ if (normalizedDelta !== 0) {
+ apiKeyQuotaTracker.adjustSpending(apiKeyId, normalizedDelta);
+ }
+ }
}
function normalizeUsage(usage: BillingUsageInput): NormalizedBillingUsage {
@@ -113,6 +135,7 @@ async function upsertUnbilledSnapshot(
const previousSnapshot = await db.query.requestBillingSnapshots.findFirst({
where: eq(requestBillingSnapshots.requestLogId, input.requestLogId),
columns: {
+ apiKeyId: true,
upstreamId: true,
billingStatus: true,
finalCost: true,
@@ -185,7 +208,13 @@ async function upsertUnbilledSnapshot(
},
});
- applyQuotaDeltaAfterSnapshotUpsert(previousSnapshot ?? null, "unbilled", input.upstreamId, null);
+ applyQuotaDeltaAfterSnapshotUpsert(
+ previousSnapshot ?? null,
+ "unbilled",
+ input.apiKeyId,
+ input.upstreamId,
+ null
+ );
return {
status: "unbilled",
@@ -293,6 +322,7 @@ export async function calculateAndPersistRequestBillingSnapshot(
const previousSnapshot = await db.query.requestBillingSnapshots.findFirst({
where: eq(requestBillingSnapshots.requestLogId, input.requestLogId),
columns: {
+ apiKeyId: true,
upstreamId: true,
billingStatus: true,
finalCost: true,
@@ -368,6 +398,7 @@ export async function calculateAndPersistRequestBillingSnapshot(
applyQuotaDeltaAfterSnapshotUpsert(
previousSnapshot ?? null,
"billed",
+ input.apiKeyId,
input.upstreamId,
cost.finalCost
);
diff --git a/src/lib/services/key-manager.ts b/src/lib/services/key-manager.ts
index b10b1873..891c0111 100644
--- a/src/lib/services/key-manager.ts
+++ b/src/lib/services/key-manager.ts
@@ -4,6 +4,9 @@ import { db, apiKeys, apiKeyUpstreams, upstreams, type ApiKey } from "../db";
import { hashApiKey, verifyApiKey } from "../utils/auth";
import { encrypt, decrypt, EncryptionError } from "../utils/encryption";
import { createLogger } from "../utils/logger";
+import { apiKeyQuotaTracker } from "@/lib/services/api-key-quota-tracker";
+import { parseSpendingRules } from "@/lib/services/spending-rules";
+import type { SpendingRule } from "@/lib/services/upstream-quota-tracker";
const log = createLogger("key-manager");
@@ -35,6 +38,18 @@ export interface ApiKeyCreateInput {
accessMode?: ApiKeyAccessMode;
description?: string | null;
expiresAt?: Date | null;
+ spendingRules?: SpendingRule[] | null;
+}
+
+export interface ApiKeySpendingRuleStatus {
+ periodType: SpendingRule["period_type"];
+ periodHours: number | null;
+ currentSpending: number;
+ spendingLimit: number;
+ percentUsed: number;
+ isExceeded: boolean;
+ resetsAt: Date | null;
+ estimatedRecoveryAt: Date | null;
}
export interface ApiKeyCreateResult {
@@ -45,6 +60,9 @@ export interface ApiKeyCreateResult {
description: string | null;
accessMode: ApiKeyAccessMode;
upstreamIds: string[];
+ spendingRules: SpendingRule[] | null;
+ spendingRuleStatuses: ApiKeySpendingRuleStatus[];
+ isQuotaExceeded: boolean;
isActive: boolean;
expiresAt: Date | null;
createdAt: Date;
@@ -58,6 +76,9 @@ export interface ApiKeyListItem {
description: string | null;
accessMode: ApiKeyAccessMode;
upstreamIds: string[];
+ spendingRules: SpendingRule[] | null;
+ spendingRuleStatuses: ApiKeySpendingRuleStatus[];
+ isQuotaExceeded: boolean;
isActive: boolean;
expiresAt: Date | null;
createdAt: Date;
@@ -86,6 +107,7 @@ export interface ApiKeyUpdateInput {
accessMode?: ApiKeyAccessMode;
expiresAt?: Date | null;
upstreamIds?: string[];
+ spendingRules?: SpendingRule[] | null;
}
function normalizeAccessMode(
@@ -99,6 +121,126 @@ function normalizeAccessMode(
return upstreamIds.length > 0 ? "restricted" : "unrestricted";
}
+function buildFallbackQuotaState(spendingRules: SpendingRule[] | null): {
+ spendingRuleStatuses: ApiKeySpendingRuleStatus[];
+ isQuotaExceeded: boolean;
+} {
+ return {
+ spendingRuleStatuses:
+ spendingRules?.map((rule) => ({
+ periodType: rule.period_type,
+ periodHours: rule.period_hours ?? null,
+ currentSpending: 0,
+ spendingLimit: rule.limit,
+ percentUsed: 0,
+ isExceeded: false,
+ resetsAt: null,
+ estimatedRecoveryAt: null,
+ })) ?? [],
+ isQuotaExceeded: false,
+ };
+}
+
+async function syncApiKeyQuotaStateBestEffort(
+ apiKeyId: string,
+ apiKeyName: string,
+ spendingRules: SpendingRule[] | null
+): Promise {
+ try {
+ await apiKeyQuotaTracker.syncApiKeyFromDb(apiKeyId, apiKeyName, spendingRules);
+ } catch (error) {
+ log.error(
+ { err: error, apiKeyId, apiKeyName },
+ "failed to sync API key quota tracker after persisted mutation"
+ );
+ }
+}
+
+async function resolveSpendingRuleStatuses(
+ apiKeyId: string,
+ spendingRules: SpendingRule[] | null
+): Promise<{
+ spendingRuleStatuses: ApiKeySpendingRuleStatus[];
+ isQuotaExceeded: boolean;
+}> {
+ if (!spendingRules || spendingRules.length === 0) {
+ return buildFallbackQuotaState(spendingRules);
+ }
+
+ try {
+ await apiKeyQuotaTracker.initialize();
+ const status = apiKeyQuotaTracker.getQuotaStatus(apiKeyId);
+ if (!status) {
+ return buildFallbackQuotaState(spendingRules);
+ }
+
+ const spendingRuleStatuses = await Promise.all(
+ status.rules.map(async (rule) => ({
+ periodType: rule.periodType,
+ periodHours: rule.periodHours,
+ currentSpending: rule.currentSpending,
+ spendingLimit: rule.spendingLimit,
+ percentUsed: rule.percentUsed,
+ isExceeded: rule.isExceeded,
+ resetsAt: rule.resetsAt,
+ estimatedRecoveryAt:
+ rule.periodType === "rolling" && rule.isExceeded
+ ? await apiKeyQuotaTracker.estimateRecoveryTime(apiKeyId, {
+ period_type: "rolling",
+ limit: rule.spendingLimit,
+ ...(rule.periodHours != null ? { period_hours: rule.periodHours } : {}),
+ })
+ : null,
+ }))
+ );
+
+ return {
+ spendingRuleStatuses,
+ isQuotaExceeded: status.isExceeded,
+ };
+ } catch (error) {
+ log.error({ err: error, apiKeyId }, "failed to resolve API key quota statuses");
+ return buildFallbackQuotaState(spendingRules);
+ }
+}
+
+async function buildApiKeyListItem(
+ key: Pick<
+ ApiKey,
+ | "id"
+ | "keyPrefix"
+ | "name"
+ | "description"
+ | "accessMode"
+ | "spendingRules"
+ | "isActive"
+ | "expiresAt"
+ | "createdAt"
+ | "updatedAt"
+ >,
+ upstreamIds: string[]
+): Promise {
+ const accessMode = normalizeAccessMode(key.accessMode, upstreamIds);
+ const spendingRules = parseSpendingRules(key.spendingRules);
+ const quotaState = await resolveSpendingRuleStatuses(key.id, spendingRules);
+
+ return {
+ id: key.id,
+ keyPrefix: key.keyPrefix,
+ name: key.name,
+ description: key.description,
+ accessMode,
+ upstreamIds: accessMode === "restricted" ? upstreamIds : [],
+ spendingRules,
+ spendingRuleStatuses: quotaState.spendingRuleStatuses,
+ isQuotaExceeded: quotaState.isQuotaExceeded,
+ isActive: key.isActive,
+ expiresAt: key.expiresAt,
+ createdAt: key.createdAt,
+ updatedAt: key.updatedAt,
+ };
+}
+
/**
* Generate a random API key using the `sk-auto-[base64-random]` format.
*/
@@ -114,6 +256,7 @@ export async function createApiKey(input: ApiKeyCreateInput): Promise {
}
await db.delete(apiKeys).where(eq(apiKeys.id, keyId));
+ await syncApiKeyQuotaStateBestEffort(keyId, existing.name, null);
log.info({ keyPrefix: existing.keyPrefix, name: existing.name }, "deleted API key");
}
@@ -229,6 +386,7 @@ export async function listApiKeys(
limit: pageSize,
offset,
});
+ await apiKeyQuotaTracker.initialize();
// For each API key, fetch authorized upstream IDs
const items: ApiKeyListItem[] = await Promise.all(
@@ -237,20 +395,7 @@ export async function listApiKeys(
where: eq(apiKeyUpstreams.apiKeyId, key.id),
});
const upstreamIds = upstreamLinks.map((link) => link.upstreamId);
- const accessMode = normalizeAccessMode(key.accessMode, upstreamIds);
-
- return {
- id: key.id,
- keyPrefix: key.keyPrefix,
- name: key.name,
- description: key.description,
- accessMode,
- upstreamIds: accessMode === "restricted" ? upstreamIds : [],
- isActive: key.isActive,
- expiresAt: key.expiresAt,
- createdAt: key.createdAt,
- updatedAt: key.updatedAt,
- };
+ return buildApiKeyListItem(key, upstreamIds);
})
);
@@ -318,20 +463,8 @@ export async function getApiKeyById(keyId: string): Promise link.upstreamId);
- const accessMode = normalizeAccessMode(apiKey.accessMode, upstreamIds);
-
- return {
- id: apiKey.id,
- keyPrefix: apiKey.keyPrefix,
- name: apiKey.name,
- description: apiKey.description,
- accessMode,
- upstreamIds: accessMode === "restricted" ? upstreamIds : [],
- isActive: apiKey.isActive,
- expiresAt: apiKey.expiresAt,
- createdAt: apiKey.createdAt,
- updatedAt: apiKey.updatedAt,
- };
+ await apiKeyQuotaTracker.initialize();
+ return buildApiKeyListItem(apiKey, upstreamIds);
}
/**
@@ -374,8 +507,10 @@ export async function updateApiKey(
): Promise {
const { name, description, isActive, accessMode, expiresAt, upstreamIds } = input;
const now = new Date();
+ const parsedSpendingRules =
+ input.spendingRules !== undefined ? parseSpendingRules(input.spendingRules) : undefined;
- return db.transaction(async (tx) => {
+ const updatedResult = await db.transaction(async (tx) => {
// Check if key exists
const existing = await tx.query.apiKeys.findFirst({
where: eq(apiKeys.id, keyId),
@@ -431,6 +566,9 @@ export async function updateApiKey(
isActive: boolean;
accessMode: ApiKeyAccessMode;
expiresAt: Date | null;
+ spendingRules:
+ | { period_type: "daily" | "monthly" | "rolling"; limit: number; period_hours?: number }[]
+ | null;
updatedAt: Date;
}> = { updatedAt: now };
@@ -449,6 +587,9 @@ export async function updateApiKey(
if (expiresAt !== undefined) {
updateData.expiresAt = expiresAt;
}
+ if (parsedSpendingRules !== undefined) {
+ updateData.spendingRules = parsedSpendingRules;
+ }
// Update the API key record
const [updatedKey] = await tx
@@ -492,7 +633,17 @@ export async function updateApiKey(
}
}
- log.info({ keyPrefix: updatedKey.keyPrefix, name: updatedKey.name }, "updated API key");
+ log.info(
+ {
+ keyPrefix: updatedKey.keyPrefix,
+ name: updatedKey.name,
+ spendingRules:
+ parsedSpendingRules !== undefined
+ ? (parsedSpendingRules?.length ?? 0)
+ : (existing.spendingRules?.length ?? 0),
+ },
+ "updated API key"
+ );
const resolvedAccessMode = normalizeAccessMode(updatedKey.accessMode, currentUpstreamIds);
@@ -503,10 +654,33 @@ export async function updateApiKey(
description: updatedKey.description,
accessMode: resolvedAccessMode,
upstreamIds: resolvedAccessMode === "restricted" ? currentUpstreamIds : [],
+ spendingRules: parseSpendingRules(updatedKey.spendingRules),
isActive: updatedKey.isActive,
expiresAt: updatedKey.expiresAt,
createdAt: updatedKey.createdAt,
updatedAt: updatedKey.updatedAt,
};
});
+
+ await syncApiKeyQuotaStateBestEffort(
+ updatedResult.id,
+ updatedResult.name,
+ updatedResult.spendingRules
+ );
+
+ return buildApiKeyListItem(
+ {
+ id: updatedResult.id,
+ keyPrefix: updatedResult.keyPrefix,
+ name: updatedResult.name,
+ description: updatedResult.description,
+ accessMode: updatedResult.accessMode,
+ spendingRules: updatedResult.spendingRules,
+ isActive: updatedResult.isActive,
+ expiresAt: updatedResult.expiresAt,
+ createdAt: updatedResult.createdAt,
+ updatedAt: updatedResult.updatedAt,
+ },
+ updatedResult.upstreamIds
+ );
}
diff --git a/src/lib/services/spending-rules.ts b/src/lib/services/spending-rules.ts
new file mode 100644
index 00000000..99ded995
--- /dev/null
+++ b/src/lib/services/spending-rules.ts
@@ -0,0 +1,64 @@
+import { z } from "zod";
+import type { SpendingRule } from "@/lib/services/upstream-quota-tracker";
+
+export const MAX_SPENDING_RULE_PERIOD_HOURS = 8760;
+
+export const spendingRuleSchema = z.object({
+ period_type: z.enum(["daily", "monthly", "rolling"]),
+ limit: z.number().positive(),
+ period_hours: z.number().int().min(1).max(MAX_SPENDING_RULE_PERIOD_HOURS).optional(),
+});
+
+export const nullableSpendingRulesSchema = z
+ .array(spendingRuleSchema)
+ .nullable()
+ .optional()
+ .superRefine((rules, ctx) => {
+ if (!rules) {
+ return;
+ }
+
+ for (const [index, rule] of rules.entries()) {
+ if (rule.period_type === "rolling" && rule.period_hours == null) {
+ ctx.addIssue({
+ code: z.ZodIssueCode.custom,
+ path: [index, "period_hours"],
+ message: "period_hours is required when period_type is 'rolling'",
+ });
+ }
+ }
+ });
+
+/**
+ * Normalize persisted or submitted spending rules into the canonical runtime shape.
+ */
+export function normalizeSpendingRules(
+ rules: SpendingRule[] | null | undefined
+): SpendingRule[] | null {
+ if (!rules || rules.length === 0) {
+ return null;
+ }
+
+ return rules.map((rule) => {
+ if (rule.period_type === "rolling") {
+ return {
+ period_type: "rolling",
+ limit: Number(rule.limit),
+ period_hours: rule.period_hours ?? 24,
+ };
+ }
+
+ return {
+ period_type: rule.period_type,
+ limit: Number(rule.limit),
+ };
+ });
+}
+
+/**
+ * Parse unknown input with shared validation and return normalized spending rules.
+ */
+export function parseSpendingRules(input: unknown): SpendingRule[] | null {
+ const parsed = nullableSpendingRulesSchema.parse(input);
+ return normalizeSpendingRules(parsed ?? null);
+}
diff --git a/src/lib/services/unified-error.ts b/src/lib/services/unified-error.ts
index f5b1f156..0d7beac3 100644
--- a/src/lib/services/unified-error.ts
+++ b/src/lib/services/unified-error.ts
@@ -14,6 +14,7 @@ export type UnifiedErrorCode =
| "ALL_UPSTREAMS_UNAVAILABLE"
| "NO_AUTHORIZED_UPSTREAMS"
| "NO_UPSTREAMS_CONFIGURED"
+ | "API_KEY_QUOTA_EXCEEDED"
| "SERVICE_UNAVAILABLE"
| "REQUEST_TIMEOUT"
| "CLIENT_DISCONNECTED"
@@ -28,6 +29,7 @@ export type UnifiedErrorType = "service_unavailable" | "timeout" | "client_error
* Optional error reason details for better diagnostics.
*/
export type UnifiedErrorReason =
+ | "API_KEY_QUOTA_EXCEEDED"
| "NO_AUTHORIZED_UPSTREAMS"
| "CLIENT_DISCONNECTED"
| "NO_HEALTHY_CANDIDATES"
@@ -65,6 +67,7 @@ const ERROR_MESSAGES: Record = {
ALL_UPSTREAMS_UNAVAILABLE: "服务暂时不可用，请稍后重试",
NO_AUTHORIZED_UPSTREAMS: "当前密钥未绑定可用上游，请先完成授权配置",
NO_UPSTREAMS_CONFIGURED: "服务暂时不可用，请稍后重试",
+ API_KEY_QUOTA_EXCEEDED: "当前密钥已达到消费限额，请稍后重试",
SERVICE_UNAVAILABLE: "服务暂时不可用，请稍后重试",
REQUEST_TIMEOUT: "请求超时，请稍后重试",
CLIENT_DISCONNECTED: "客户端已断开连接",
@@ -78,6 +81,7 @@ const ERROR_TYPES: Record = {
ALL_UPSTREAMS_UNAVAILABLE: "service_unavailable",
NO_AUTHORIZED_UPSTREAMS: "client_error",
NO_UPSTREAMS_CONFIGURED: "service_unavailable",
+ API_KEY_QUOTA_EXCEEDED: "client_error",
SERVICE_UNAVAILABLE: "service_unavailable",
REQUEST_TIMEOUT: "timeout",
CLIENT_DISCONNECTED: "client_error",
@@ -91,6 +95,7 @@ const HTTP_STATUS_CODES: Record = {
ALL_UPSTREAMS_UNAVAILABLE: 503,
NO_AUTHORIZED_UPSTREAMS: 403,
NO_UPSTREAMS_CONFIGURED: 503,
+ API_KEY_QUOTA_EXCEEDED: 429,
SERVICE_UNAVAILABLE: 503,
REQUEST_TIMEOUT: 504,
CLIENT_DISCONNECTED: 499, // Client Closed Request (nginx convention)
diff --git a/src/lib/utils/api-transformers.ts b/src/lib/utils/api-transformers.ts
index f83f5870..cacf6e0e 100644
--- a/src/lib/utils/api-transformers.ts
+++ b/src/lib/utils/api-transformers.ts
@@ -188,6 +188,20 @@ export interface ApiKeyApiResponse {
description: string | null;
access_mode: "unrestricted" | "restricted";
upstream_ids: string[];
+ spending_rules:
+ | { period_type: "daily" | "monthly" | "rolling"; limit: number; period_hours?: number }[]
+ | null;
+ spending_rule_statuses: {
+ period_type: "daily" | "monthly" | "rolling";
+ period_hours: number | null;
+ current_spending: number;
+ spending_limit: number;
+ percent_used: number;
+ is_exceeded: boolean;
+ resets_at: string | null;
+ estimated_recovery_at: string | null;
+ }[];
+ is_quota_exceeded: boolean;
is_active: boolean;
expires_at: string | null;
created_at: string;
@@ -227,6 +241,18 @@ export function transformApiKeyToApi(apiKey: ApiKeyListItem): ApiKeyApiResponse
description: apiKey.description,
access_mode: apiKey.accessMode,
upstream_ids: apiKey.upstreamIds,
+ spending_rules: apiKey.spendingRules,
+ spending_rule_statuses: apiKey.spendingRuleStatuses.map((rule) => ({
+ period_type: rule.periodType,
+ period_hours: rule.periodHours,
+ current_spending: rule.currentSpending,
+ spending_limit: rule.spendingLimit,
+ percent_used: rule.percentUsed,
+ is_exceeded: rule.isExceeded,
+ resets_at: toISOStringOrNull(rule.resetsAt),
+ estimated_recovery_at: toISOStringOrNull(rule.estimatedRecoveryAt),
+ })),
+ is_quota_exceeded: apiKey.isQuotaExceeded,
is_active: apiKey.isActive,
expires_at: toISOStringOrNull(apiKey.expiresAt),
created_at: apiKey.createdAt.toISOString(),
@@ -247,6 +273,18 @@ export function transformApiKeyCreateToApi(result: ApiKeyCreateResult): ApiKeyCr
description: result.description,
access_mode: result.accessMode,
upstream_ids: result.upstreamIds,
+ spending_rules: result.spendingRules,
+ spending_rule_statuses: result.spendingRuleStatuses.map((rule) => ({
+ period_type: rule.periodType,
+ period_hours: rule.periodHours,
+ current_spending: rule.currentSpending,
+ spending_limit: rule.spendingLimit,
+ percent_used: rule.percentUsed,
+ is_exceeded: rule.isExceeded,
+ resets_at: toISOStringOrNull(rule.resetsAt),
+ estimated_recovery_at: toISOStringOrNull(rule.estimatedRecoveryAt),
+ })),
+ is_quota_exceeded: result.isQuotaExceeded,
is_active: result.isActive,
expires_at: toISOStringOrNull(result.expiresAt),
created_at: result.createdAt.toISOString(),
diff --git a/src/messages/en.json b/src/messages/en.json
index 65b1e24e..4a302ec2 100644
--- a/src/messages/en.json
+++ b/src/messages/en.json
@@ -170,6 +170,29 @@
"deselectFilteredUpstreams": "Deselect visible results",
"filteredUpstreamsSelected": "{selected} of {total} visible upstreams selected",
"noMatchingUpstreams": "No upstreams match the current search",
+ "spendingRules": "Spending Rules",
+ "spendingRulesDesc": "Configure amount-based limits for this key. Any exceeded rule will block later requests.",
+ "spendingRulesEmpty": "No spending rules configured",
+ "addSpendingRule": "Add rule",
+ "spendingRuleLabel": "Rule {index}",
+ "quotaPeriodType": "Period Type",
+ "quotaPeriodType_daily": "Daily",
+ "quotaPeriodType_monthly": "Monthly",
+ "quotaPeriodType_rolling": "Rolling Window",
+ "quotaPeriodRollingWithHours": "Rolling {hours}h",
+ "quotaLimitUsd": "Limit Amount (USD)",
+ "quotaLimitPlaceholder": "e.g. 50",
+ "quotaLimitPositive": "Limit amount must be greater than 0",
+ "quotaPeriodHours": "Rolling Window Hours",
+ "quotaPeriodHoursPlaceholder": "e.g. 6",
+ "quotaPeriodHoursRequired": "Rolling rules require period hours",
+ "quotaStatusPending": "Quota status is initializing. Refresh in a moment.",
+ "quotaExceeded": "Exceeded",
+ "quotaWithinLimit": "Available",
+ "quotaPercentUsed": "{percent}% used",
+ "quotaResetTime": "Resets {time}",
+ "quotaRecoveryTime": "Recovers {time}",
+ "quotaRecoveryPending": "Recovery time is being calculated",
"expirationDate": "Expiration Date (Optional)",
"expirationDateDesc": "Leave empty for no expiration",
"selectDate": "Select date",
diff --git a/src/messages/zh-CN.json b/src/messages/zh-CN.json
index 21f2208f..5fc52710 100644
--- a/src/messages/zh-CN.json
+++ b/src/messages/zh-CN.json
@@ -170,6 +170,29 @@
"deselectFilteredUpstreams": "取消当前结果",
"filteredUpstreamsSelected": "当前结果已选 {selected} / {total}",
"noMatchingUpstreams": "没有匹配当前搜索的上游",
+ "spendingRules": "消费限额规则",
+ "spendingRulesDesc": "为此密钥配置金额限额。任一规则达到上限后，后续请求会被拒绝。",
+ "spendingRulesEmpty": "当前未配置消费限额规则",
+ "addSpendingRule": "添加规则",
+ "spendingRuleLabel": "规则 {index}",
+ "quotaPeriodType": "周期类型",
+ "quotaPeriodType_daily": "每日",
+ "quotaPeriodType_monthly": "每月",
+ "quotaPeriodType_rolling": "滚动窗口",
+ "quotaPeriodRollingWithHours": "滚动 {hours} 小时",
+ "quotaLimitUsd": "限额金额（USD）",
+ "quotaLimitPlaceholder": "例如 50",
+ "quotaLimitPositive": "限额金额必须大于 0",
+ "quotaPeriodHours": "滚动窗口时长（小时）",
+ "quotaPeriodHoursPlaceholder": "例如 6",
+ "quotaPeriodHoursRequired": "滚动窗口规则必须填写小时数",
+ "quotaStatusPending": "额度状态正在初始化，请稍后刷新",
+ "quotaExceeded": "已超额",
+ "quotaWithinLimit": "可用",
+ "quotaPercentUsed": "已用 {percent}%",
+ "quotaResetTime": "重置时间：{time}",
+ "quotaRecoveryTime": "预计恢复：{time}",
+ "quotaRecoveryPending": "预计恢复时间计算中",
"expirationDate": "过期时间（可选）",
"expirationDateDesc": "不设置则永不过期",
"selectDate": "选择日期",
diff --git a/src/types/api.ts b/src/types/api.ts
index 6a321c0b..742082a1 100644
--- a/src/types/api.ts
+++ b/src/types/api.ts
@@ -10,12 +10,30 @@ import type { RouteCapability, RouteMatchSource } from "@/lib/route-capabilities
export type APIKeyAccessMode = "unrestricted" | "restricted";
+export interface APIKeySpendingRule {
+ period_type: "daily" | "monthly" | "rolling";
+ limit: number;
+ period_hours?: number;
+}
+
+export interface APIKeySpendingRuleStatus {
+ period_type: "daily" | "monthly" | "rolling";
+ period_hours: number | null;
+ current_spending: number;
+ spending_limit: number;
+ percent_used: number;
+ is_exceeded: boolean;
+ resets_at: string | null;
+ estimated_recovery_at: string | null;
+}
+
export interface APIKeyCreate {
name: string;
description?: string | null;
access_mode?: APIKeyAccessMode;
upstream_ids: string[]; // UUID[]
expires_at?: string | null; // ISO 8601 date string
+ spending_rules?: APIKeySpendingRule[] | null;
}
export interface APIKeyUpdate {
@@ -25,6 +43,7 @@ export interface APIKeyUpdate {
access_mode?: APIKeyAccessMode;
expires_at?: string | null; // ISO 8601 date string
upstream_ids?: string[]; // UUID[]
+ spending_rules?: APIKeySpendingRule[] | null;
}
export interface APIKeyResponse {
@@ -34,6 +53,9 @@ export interface APIKeyResponse {
description: string | null;
access_mode: APIKeyAccessMode;
upstream_ids: string[]; // UUID[]
+ spending_rules: APIKeySpendingRule[] | null;
+ spending_rule_statuses: APIKeySpendingRuleStatus[];
+ is_quota_exceeded: boolean;
is_active: boolean;
expires_at: string | null; // ISO 8601 date string
created_at: string; // ISO 8601 date string
diff --git a/tests/components/create-key-dialog.test.tsx b/tests/components/create-key-dialog.test.tsx
index 0d44b012..05e7ce86 100644
--- a/tests/components/create-key-dialog.test.tsx
+++ b/tests/components/create-key-dialog.test.tsx
@@ -163,6 +163,9 @@ describe("CreateKeyDialog", () => {
description: null,
access_mode: "restricted",
upstream_ids: ["upstream-2"],
+ spending_rules: null,
+ spending_rule_statuses: [],
+ is_quota_exceeded: false,
is_active: true,
expires_at: null,
created_at: "2024-01-01T00:00:00Z",
@@ -196,6 +199,7 @@ describe("CreateKeyDialog", () => {
access_mode: "restricted",
upstream_ids: ["upstream-2"],
expires_at: null,
+ spending_rules: null,
});
});
});
@@ -210,6 +214,48 @@ describe("CreateKeyDialog", () => {
expect(screen.getByText("create")).toBeInTheDocument();
});
});
+
+ it("allows adding a spending rule and submits it", async () => {
+ mockCreateMutateAsync.mockResolvedValueOnce({
+ id: "key-2",
+ key_value: "sk-auto-quota",
+ key_prefix: "sk-auto-quota",
+ name: "Quota Key",
+ description: null,
+ access_mode: "unrestricted",
+ upstream_ids: [],
+ spending_rules: [{ period_type: "daily", limit: 12.5 }],
+ spending_rule_statuses: [],
+ is_quota_exceeded: false,
+ is_active: true,
+ expires_at: null,
+ created_at: "2024-01-01T00:00:00Z",
+ updated_at: "2024-01-01T00:00:00Z",
+ });
+
+ render(, { wrapper: Wrapper });
+
+ fireEvent.click(screen.getByText("createKey"));
+ fireEvent.click(screen.getByText("addSpendingRule"));
+ fireEvent.change(screen.getByPlaceholderText("keyNamePlaceholder"), {
+ target: { value: "Quota Key" },
+ });
+ fireEvent.change(screen.getByPlaceholderText("quotaLimitPlaceholder"), {
+ target: { value: "12.5" },
+ });
+ fireEvent.click(screen.getByText("create"));
+
+ await waitFor(() => {
+ expect(mockCreateMutateAsync).toHaveBeenCalledWith({
+ name: "Quota Key",
+ description: null,
+ access_mode: "unrestricted",
+ upstream_ids: [],
+ expires_at: null,
+ spending_rules: [{ period_type: "daily", limit: 12.5 }],
+ });
+ });
+ });
});
describe("Form Validation", () => {
@@ -241,6 +287,9 @@ describe("CreateKeyDialog", () => {
description: null,
access_mode: "unrestricted",
upstream_ids: [],
+ spending_rules: null,
+ spending_rule_statuses: [],
+ is_quota_exceeded: false,
is_active: true,
expires_at: null,
created_at: "2024-01-01T00:00:00Z",
@@ -267,6 +316,7 @@ describe("CreateKeyDialog", () => {
access_mode: "unrestricted",
upstream_ids: [],
expires_at: null,
+ spending_rules: null,
});
});
});
@@ -301,6 +351,9 @@ describe("CreateKeyDialog", () => {
description: null,
access_mode: "restricted",
upstream_ids: ["upstream-1"],
+ spending_rules: null,
+ spending_rule_statuses: [],
+ is_quota_exceeded: false,
is_active: true,
expires_at: null,
created_at: "2024-01-01T00:00:00Z",
@@ -332,6 +385,7 @@ describe("CreateKeyDialog", () => {
access_mode: "restricted",
upstream_ids: ["upstream-1"],
expires_at: null,
+ spending_rules: null,
});
});
});
diff --git a/tests/components/edit-key-dialog.test.tsx b/tests/components/edit-key-dialog.test.tsx
index 09edc5b9..3635d047 100644
--- a/tests/components/edit-key-dialog.test.tsx
+++ b/tests/components/edit-key-dialog.test.tsx
@@ -73,6 +73,9 @@ describe("EditKeyDialog", () => {
description: "Test description",
access_mode: "restricted",
upstream_ids: ["upstream-1", "upstream-2"],
+ spending_rules: null,
+ spending_rule_statuses: [],
+ is_quota_exceeded: false,
is_active: true,
expires_at: null,
created_at: "2024-01-01T00:00:00Z",
@@ -192,6 +195,7 @@ describe("EditKeyDialog", () => {
data: expect.objectContaining({
access_mode: "restricted",
upstream_ids: ["upstream-1", "upstream-2", "upstream-3"],
+ spending_rules: null,
}),
});
});
@@ -226,6 +230,46 @@ describe("EditKeyDialog", () => {
expect(screen.getByText("cancel")).toBeInTheDocument();
expect(screen.getByText("save")).toBeInTheDocument();
});
+
+ it("renders existing spending rules and preserves them on submit", async () => {
+ mockUpdateMutateAsync.mockResolvedValueOnce({});
+
+ const quotaApiKey: APIKeyResponse = {
+ ...mockApiKey,
+ spending_rules: [{ period_type: "rolling", limit: 15, period_hours: 6 }],
+ spending_rule_statuses: [
+ {
+ period_type: "rolling",
+ period_hours: 6,
+ current_spending: 8,
+ spending_limit: 15,
+ percent_used: 53.3,
+ is_exceeded: false,
+ resets_at: null,
+ estimated_recovery_at: null,
+ },
+ ],
+ };
+
+ render(, {
+ wrapper: Wrapper,
+ });
+
+ expect(screen.getByText("spendingRuleLabel")).toBeInTheDocument();
+ expect(screen.getByDisplayValue("15")).toBeInTheDocument();
+ expect(screen.getByDisplayValue("6")).toBeInTheDocument();
+
+ fireEvent.click(screen.getByText("save"));
+
+ await waitFor(() => {
+ expect(mockUpdateMutateAsync).toHaveBeenCalledWith({
+ id: "key-123",
+ data: expect.objectContaining({
+ spending_rules: [{ period_type: "rolling", limit: 15, period_hours: 6 }],
+ }),
+ });
+ });
+ });
});
describe("Form Validation", () => {
@@ -281,6 +325,7 @@ describe("EditKeyDialog", () => {
data: expect.objectContaining({
access_mode: "unrestricted",
upstream_ids: [],
+ spending_rules: null,
}),
});
});
@@ -311,6 +356,7 @@ describe("EditKeyDialog", () => {
is_active: true,
access_mode: "restricted",
upstream_ids: ["upstream-1", "upstream-2"],
+ spending_rules: null,
}),
});
});
diff --git a/tests/components/keys-table.test.tsx b/tests/components/keys-table.test.tsx
index c41766ab..e58276dc 100644
--- a/tests/components/keys-table.test.tsx
+++ b/tests/components/keys-table.test.tsx
@@ -57,6 +57,9 @@ describe("KeysTable", () => {
description: "Test description",
access_mode: "restricted",
upstream_ids: ["upstream-1", "upstream-2"],
+ spending_rules: null,
+ spending_rule_statuses: [],
+ is_quota_exceeded: false,
is_active: true,
expires_at: null,
created_at: "2024-06-01T10:00:00Z",
@@ -151,6 +154,44 @@ describe("KeysTable", () => {
expect(screen.getByText("unrestrictedAccess")).toBeInTheDocument();
});
+
+ it("renders rule-level quota status when spending data is present", () => {
+ render(
+
+ );
+
+ // quotaExceeded badge shown in collapsed row
+ expect(screen.getAllByText("quotaExceeded").length).toBeGreaterThan(0);
+
+ // Click to expand the quota details
+ const nameCell = screen.getByText("Test API Key").closest("td")!;
+ fireEvent.click(nameCell);
+
+ expect(screen.getByText("110.0%")).toBeInTheDocument();
+ expect(screen.getByText("quotaRecoveryTime")).toBeInTheDocument();
+ });
});
describe("Expiry Formatting", () => {
@@ -524,7 +565,7 @@ describe("KeysTable", () => {
render();
const searchInput = screen.getByPlaceholderText("searchKeys");
- expect(searchInput).toHaveClass("max-w-md");
+ expect(searchInput).toHaveClass("max-w-sm");
});
});
diff --git a/tests/unit/api/admin/keys/route.test.ts b/tests/unit/api/admin/keys/route.test.ts
new file mode 100644
index 00000000..c2332873
--- /dev/null
+++ b/tests/unit/api/admin/keys/route.test.ts
@@ -0,0 +1,196 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { NextRequest } from "next/server";
+
+vi.mock("@/lib/utils/auth", () => ({
+ validateAdminAuth: vi.fn((authHeader: string | null) => authHeader === "Bearer test-admin-token"),
+}));
+
+const mockCreateApiKey = vi.fn();
+const mockUpdateApiKey = vi.fn();
+
+class ApiKeyNotFoundError extends Error {}
+
+vi.mock("@/lib/services/key-manager", () => ({
+ createApiKey: (...args: unknown[]) => mockCreateApiKey(...args),
+ updateApiKey: (...args: unknown[]) => mockUpdateApiKey(...args),
+ listApiKeys: vi.fn(),
+ getApiKeyById: vi.fn(),
+ deleteApiKey: vi.fn(),
+ ApiKeyNotFoundError,
+}));
+
+const QUOTA_RULE = { period_type: "rolling" as const, limit: 15, period_hours: 6 };
+const KEY_ID = "11111111-1111-4111-8111-111111111111";
+const UPSTREAM_ID = "22222222-2222-4222-8222-222222222222";
+
+function buildServiceApiKey(overrides?: Record) {
+ return {
+ id: KEY_ID,
+ keyPrefix: "sk-auto-test",
+ keyValue: "sk-auto-test-full",
+ name: "Quota Key",
+ description: "quota aware",
+ accessMode: "restricted" as const,
+ upstreamIds: [UPSTREAM_ID],
+ spendingRules: [QUOTA_RULE],
+ spendingRuleStatuses: [
+ {
+ periodType: "rolling" as const,
+ periodHours: 6,
+ currentSpending: 8,
+ spendingLimit: 15,
+ percentUsed: 53.3,
+ isExceeded: false,
+ resetsAt: null,
+ estimatedRecoveryAt: new Date("2024-01-01T06:00:00.000Z"),
+ },
+ ],
+ isQuotaExceeded: false,
+ isActive: true,
+ expiresAt: null,
+ createdAt: new Date("2024-01-01T00:00:00.000Z"),
+ updatedAt: new Date("2024-01-01T00:00:00.000Z"),
+ ...overrides,
+ };
+}
+
+describe("admin keys routes spending rules", () => {
+ beforeEach(() => {
+ vi.clearAllMocks();
+ });
+
+ it("POST /api/admin/keys should pass spending_rules into createApiKey and return quota fields", async () => {
+ const { POST } = await import("@/app/api/admin/keys/route");
+ mockCreateApiKey.mockResolvedValueOnce(buildServiceApiKey());
+
+ const request = new NextRequest("http://localhost:3000/api/admin/keys", {
+ method: "POST",
+ headers: {
+ authorization: "Bearer test-admin-token",
+ "content-type": "application/json",
+ },
+ body: JSON.stringify({
+ name: "Quota Key",
+ access_mode: "restricted",
+ upstream_ids: [UPSTREAM_ID],
+ spending_rules: [QUOTA_RULE],
+ }),
+ });
+
+ const response = await POST(request);
+ const data = await response.json();
+
+ expect(response.status).toBe(201);
+ expect(mockCreateApiKey).toHaveBeenCalledWith(
+ expect.objectContaining({
+ name: "Quota Key",
+ accessMode: "restricted",
+ upstreamIds: [UPSTREAM_ID],
+ spendingRules: [QUOTA_RULE],
+ })
+ );
+ expect(data).toEqual(
+ expect.objectContaining({
+ spending_rules: [QUOTA_RULE],
+ spending_rule_statuses: [
+ expect.objectContaining({
+ period_type: "rolling",
+ period_hours: 6,
+ current_spending: 8,
+ spending_limit: 15,
+ estimated_recovery_at: "2024-01-01T06:00:00.000Z",
+ }),
+ ],
+ is_quota_exceeded: false,
+ })
+ );
+ });
+
+ it("POST /api/admin/keys should reject rolling rules without period_hours", async () => {
+ const { POST } = await import("@/app/api/admin/keys/route");
+
+ const request = new NextRequest("http://localhost:3000/api/admin/keys", {
+ method: "POST",
+ headers: {
+ authorization: "Bearer test-admin-token",
+ "content-type": "application/json",
+ },
+ body: JSON.stringify({
+ name: "Broken Quota Key",
+ access_mode: "restricted",
+ upstream_ids: [UPSTREAM_ID],
+ spending_rules: [{ period_type: "rolling", limit: 15 }],
+ }),
+ });
+
+ const response = await POST(request);
+ const data = await response.json();
+
+ expect(response.status).toBe(400);
+ expect(data.error).toContain("period_hours");
+ expect(mockCreateApiKey).not.toHaveBeenCalled();
+ });
+
+ it("PUT /api/admin/keys/[id] should pass spending_rules into updateApiKey and return quota fields", async () => {
+ const { PUT } = await import("@/app/api/admin/keys/[id]/route");
+ mockUpdateApiKey.mockResolvedValueOnce(
+ buildServiceApiKey({
+ spendingRuleStatuses: [
+ {
+ periodType: "rolling" as const,
+ periodHours: 6,
+ currentSpending: 15,
+ spendingLimit: 15,
+ percentUsed: 100,
+ isExceeded: true,
+ resetsAt: null,
+ estimatedRecoveryAt: new Date("2024-01-01T08:00:00.000Z"),
+ },
+ ],
+ isQuotaExceeded: true,
+ })
+ );
+
+ const request = new NextRequest(`http://localhost:3000/api/admin/keys/${KEY_ID}`, {
+ method: "PUT",
+ headers: {
+ authorization: "Bearer test-admin-token",
+ "content-type": "application/json",
+ },
+ body: JSON.stringify({
+ access_mode: "restricted",
+ upstream_ids: [UPSTREAM_ID],
+ spending_rules: [QUOTA_RULE],
+ }),
+ });
+
+ const response = await PUT(request, {
+ params: Promise.resolve({ id: KEY_ID }),
+ });
+ const data = await response.json();
+
+ expect(response.status).toBe(200);
+ expect(mockUpdateApiKey).toHaveBeenCalledWith(
+ KEY_ID,
+ expect.objectContaining({
+ accessMode: "restricted",
+ upstreamIds: [UPSTREAM_ID],
+ spendingRules: [QUOTA_RULE],
+ })
+ );
+ expect(data).toEqual(
+ expect.objectContaining({
+ spending_rules: [QUOTA_RULE],
+ spending_rule_statuses: [
+ expect.objectContaining({
+ current_spending: 15,
+ spending_limit: 15,
+ is_exceeded: true,
+ estimated_recovery_at: "2024-01-01T08:00:00.000Z",
+ }),
+ ],
+ is_quota_exceeded: true,
+ })
+ );
+ });
+});
diff --git a/tests/unit/api/proxy/route.test.ts b/tests/unit/api/proxy/route.test.ts
index 6058cab1..db31d2c1 100644
--- a/tests/unit/api/proxy/route.test.ts
+++ b/tests/unit/api/proxy/route.test.ts
@@ -1,6 +1,26 @@
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
import { NextRequest } from "next/server";
+const { mockApiKeyQuotaTracker, mockResolveBillingModelPrice } = vi.hoisted(() => ({
+ mockApiKeyQuotaTracker: {
+ initialize: vi.fn(async () => {}),
+ getQuotaStatus: vi.fn(() => null),
+ },
+ mockResolveBillingModelPrice: vi.fn(async () => ({
+ model: "gpt-5.2",
+ source: "manual",
+ inputPricePerMillion: 1,
+ outputPricePerMillion: 1,
+ cacheReadInputPricePerMillion: null,
+ cacheWriteInputPricePerMillion: null,
+ matchedRuleType: "flat",
+ matchedRuleDisplayLabel: null,
+ appliedTierThreshold: null,
+ modelMaxInputTokens: null,
+ modelMaxOutputTokens: null,
+ })),
+}));
+
vi.mock("@/lib/utils/auth", () => ({
extractApiKey: vi.fn((authHeader: string | null) => {
if (!authHeader) return null;
@@ -102,6 +122,14 @@ vi.mock("@/lib/services/billing-cost-service", () => ({
})),
}));
+vi.mock("@/lib/services/billing-price-service", () => ({
+ resolveBillingModelPrice: mockResolveBillingModelPrice,
+}));
+
+vi.mock("@/lib/services/api-key-quota-tracker", () => ({
+ apiKeyQuotaTracker: mockApiKeyQuotaTracker,
+}));
+
// Mock load-balancer module
vi.mock("@/lib/services/load-balancer", () => {
const selectFromUpstreamCandidates = vi.fn();
@@ -312,6 +340,21 @@ describe("proxy route upstream selection", () => {
beforeEach(async () => {
vi.resetAllMocks();
vi.resetModules();
+ mockApiKeyQuotaTracker.initialize.mockResolvedValue(undefined);
+ mockApiKeyQuotaTracker.getQuotaStatus.mockReturnValue(null);
+ mockResolveBillingModelPrice.mockResolvedValue({
+ model: "gpt-5.2",
+ source: "manual",
+ inputPricePerMillion: 1,
+ outputPricePerMillion: 1,
+ cacheReadInputPricePerMillion: null,
+ cacheWriteInputPricePerMillion: null,
+ matchedRuleType: "flat",
+ matchedRuleDisplayLabel: null,
+ appliedTierThreshold: null,
+ modelMaxInputTokens: null,
+ modelMaxOutputTokens: null,
+ });
delete process.env.RECORDER_ENABLED;
delete process.env.RECORDER_MODE;
const routeModule = await import("@/app/api/proxy/v1/[...path]/route");
@@ -530,6 +573,309 @@ describe("proxy route upstream selection", () => {
});
});
+ it("should reject streaming requests before upstream routing when API key quota is exceeded", async () => {
+ const { db } = await import("@/lib/db");
+ const { forwardRequest } = await import("@/lib/services/proxy-client");
+ const { logRequest } = await import("@/lib/services/request-logger");
+ const { routeByModel } = await import("@/lib/services/model-router");
+ mockResolveBillingModelPrice.mockResolvedValueOnce({
+ model: "gpt-5.2",
+ source: "manual",
+ inputPricePerMillion: 1,
+ outputPricePerMillion: 1,
+ cacheReadInputPricePerMillion: null,
+ cacheWriteInputPricePerMillion: null,
+ matchedRuleType: "flat",
+ matchedRuleDisplayLabel: null,
+ appliedTierThreshold: null,
+ modelMaxInputTokens: null,
+ modelMaxOutputTokens: null,
+ });
+
+ vi.mocked(db.query.apiKeys.findMany).mockResolvedValueOnce([
+ { id: "key-1", keyHash: "hash-1", expiresAt: null, isActive: true },
+ ]);
+ vi.mocked(db.query.apiKeyUpstreams.findMany).mockResolvedValueOnce([
+ { upstreamId: "up-openai" },
+ ]);
+ vi.mocked(db.query.upstreams.findMany).mockResolvedValueOnce([
+ {
+ id: "up-openai",
+ name: "openai-main",
+ providerType: "openai",
+ baseUrl: "https://api.openai.com",
+ isDefault: false,
+ isActive: true,
+ timeout: 60,
+ priority: 0,
+ weight: 1,
+ routeCapabilities: ["openai_chat_compatible"],
+ },
+ ]);
+ mockApiKeyQuotaTracker.getQuotaStatus.mockReturnValue({
+ apiKeyId: "key-1",
+ apiKeyName: "Quota Key",
+ isExceeded: true,
+ rules: [
+ {
+ periodType: "daily",
+ periodHours: null,
+ currentSpending: 25,
+ spendingLimit: 25,
+ percentUsed: 100,
+ isExceeded: true,
+ resetsAt: new Date("2024-01-02T00:00:00Z"),
+ estimatedRecoveryAt: null,
+ },
+ ],
+ });
+
+ const request = new NextRequest("http://localhost/api/proxy/v1/chat/completions", {
+ method: "POST",
+ headers: {
+ authorization: "Bearer sk-test",
+ "content-type": "application/json",
+ },
+ body: JSON.stringify({
+ model: "gpt-5.2",
+ messages: [{ role: "user", content: "hello" }],
+ stream: true,
+ }),
+ });
+
+ const response = await POST(request, {
+ params: Promise.resolve({ path: ["chat", "completions"] }),
+ });
+ const data = await response.json();
+
+ expect(response.status).toBe(429);
+ expect(data).toEqual({
+ error: expect.objectContaining({
+ code: "API_KEY_QUOTA_EXCEEDED",
+ reason: "API_KEY_QUOTA_EXCEEDED",
+ did_send_upstream: false,
+ }),
+ });
+ expect(routeByModel).not.toHaveBeenCalled();
+ expect(forwardRequest).not.toHaveBeenCalled();
+ expect(logRequest).toHaveBeenCalledWith(
+ expect.objectContaining({
+ apiKeyId: "key-1",
+ upstreamId: null,
+ path: "chat/completions",
+ model: "gpt-5.2",
+ statusCode: 429,
+ errorMessage: expect.stringContaining("API key spending quota exceeded"),
+ routingDecision: expect.objectContaining({
+ matched_route_capability: "openai_chat_compatible",
+ did_send_upstream: false,
+ actual_upstream_id: null,
+ }),
+ })
+ );
+ });
+
+ it("should allow over-quota requests to proceed when billing later resolves as usage_missing", async () => {
+ const { db } = await import("@/lib/db");
+ const { forwardRequest } = await import("@/lib/services/proxy-client");
+ const { selectFromProviderType } = await import("@/lib/services/load-balancer");
+ const { calculateAndPersistRequestBillingSnapshot } =
+ await import("@/lib/services/billing-cost-service");
+
+ vi.mocked(calculateAndPersistRequestBillingSnapshot).mockResolvedValueOnce({
+ status: "unbilled",
+ unbillableReason: "usage_missing",
+ finalCost: null,
+ source: null,
+ });
+
+ const openaiUpstream = {
+ id: "up-openai",
+ name: "openai-main",
+ providerType: "openai",
+ baseUrl: "https://api.openai.com",
+ isDefault: false,
+ isActive: true,
+ timeout: 60,
+ priority: 0,
+ weight: 1,
+ };
+
+ vi.mocked(db.query.apiKeys.findMany).mockResolvedValueOnce([
+ { id: "key-1", keyHash: "hash-1", expiresAt: null, isActive: true },
+ ]);
+ mockApiKeyQuotaTracker.getQuotaStatus.mockReturnValueOnce({
+ apiKeyId: "key-1",
+ apiKeyName: "Quota Key",
+ isExceeded: true,
+ rules: [
+ {
+ periodType: "daily",
+ periodHours: null,
+ currentSpending: 25,
+ spendingLimit: 25,
+ percentUsed: 100,
+ isExceeded: true,
+ resetsAt: new Date("2024-01-02T00:00:00Z"),
+ estimatedRecoveryAt: null,
+ },
+ ],
+ });
+ vi.mocked(db.query.apiKeyUpstreams.findMany).mockResolvedValueOnce([
+ { upstreamId: "up-openai" },
+ ]);
+ vi.mocked(db.query.upstreams.findMany).mockResolvedValueOnce([
+ {
+ ...openaiUpstream,
+ routeCapabilities: ["openai_chat_compatible"],
+ },
+ ]);
+ vi.mocked(db.query.upstreamHealth.findMany).mockResolvedValueOnce([]);
+ vi.mocked(selectFromProviderType).mockResolvedValueOnce({
+ upstream: openaiUpstream,
+ providerType: "openai",
+ selectedTier: 0,
+ circuitBreakerFiltered: 0,
+ totalCandidates: 1,
+ });
+ vi.mocked(forwardRequest).mockResolvedValueOnce({
+ statusCode: 200,
+ headers: new Headers(),
+ body: new Uint8Array(),
+ isStream: false,
+ usage: {
+ promptTokens: 0,
+ completionTokens: 0,
+ totalTokens: 0,
+ },
+ });
+
+ const request = new NextRequest("http://localhost/api/proxy/v1/chat/completions", {
+ method: "POST",
+ headers: {
+ authorization: "Bearer sk-test",
+ "content-type": "application/json",
+ },
+ body: JSON.stringify({
+ model: "gpt-5.2",
+ messages: [{ role: "user", content: "hello" }],
+ }),
+ });
+
+ const response = await POST(request, {
+ params: Promise.resolve({ path: ["chat", "completions"] }),
+ });
+
+ expect(response.status).toBe(200);
+ expect(forwardRequest).toHaveBeenCalledTimes(1);
+ expect(calculateAndPersistRequestBillingSnapshot).toHaveBeenCalledWith(
+ expect.objectContaining({
+ apiKeyId: "key-1",
+ upstreamId: "up-openai",
+ usage: {
+ promptTokens: 0,
+ completionTokens: 0,
+ totalTokens: 0,
+ cacheReadTokens: 0,
+ cacheWriteTokens: 0,
+ },
+ })
+ );
+ });
+
+ it("should allow over-quota streaming requests when the resolved model has no billing price", async () => {
+ const { db } = await import("@/lib/db");
+ const { forwardRequest } = await import("@/lib/services/proxy-client");
+ const { selectFromProviderType } = await import("@/lib/services/load-balancer");
+ const { resolveBillingModelPrice } = await import("@/lib/services/billing-price-service");
+
+ const openaiUpstream = {
+ id: "up-openai",
+ name: "openai-main",
+ providerType: "openai",
+ baseUrl: "https://api.openai.com",
+ isDefault: false,
+ isActive: true,
+ timeout: 60,
+ priority: 0,
+ weight: 1,
+ modelRedirects: {
+ "gpt-5.2": "gpt-unpriced",
+ },
+ };
+
+ vi.mocked(resolveBillingModelPrice).mockResolvedValueOnce(null);
+ vi.mocked(db.query.apiKeys.findMany).mockResolvedValueOnce([
+ { id: "key-1", keyHash: "hash-1", expiresAt: null, isActive: true },
+ ]);
+ mockApiKeyQuotaTracker.getQuotaStatus.mockReturnValueOnce({
+ apiKeyId: "key-1",
+ apiKeyName: "Quota Key",
+ isExceeded: true,
+ rules: [
+ {
+ periodType: "daily",
+ periodHours: null,
+ currentSpending: 25,
+ spendingLimit: 25,
+ percentUsed: 100,
+ isExceeded: true,
+ resetsAt: new Date("2024-01-02T00:00:00Z"),
+ estimatedRecoveryAt: null,
+ },
+ ],
+ });
+ vi.mocked(db.query.apiKeyUpstreams.findMany).mockResolvedValueOnce([
+ { upstreamId: "up-openai" },
+ ]);
+ vi.mocked(db.query.upstreams.findMany).mockResolvedValueOnce([
+ {
+ ...openaiUpstream,
+ routeCapabilities: ["openai_chat_compatible"],
+ },
+ ]);
+ vi.mocked(db.query.upstreamHealth.findMany).mockResolvedValueOnce([]);
+ vi.mocked(selectFromProviderType).mockResolvedValueOnce({
+ upstream: openaiUpstream,
+ providerType: "openai",
+ selectedTier: 0,
+ circuitBreakerFiltered: 0,
+ totalCandidates: 1,
+ });
+ vi.mocked(forwardRequest).mockResolvedValueOnce({
+ statusCode: 200,
+ headers: new Headers(),
+ body: new Uint8Array(),
+ isStream: false,
+ usage: {
+ promptTokens: 32,
+ completionTokens: 16,
+ totalTokens: 48,
+ },
+ });
+
+ const request = new NextRequest("http://localhost/api/proxy/v1/chat/completions", {
+ method: "POST",
+ headers: {
+ authorization: "Bearer sk-test",
+ "content-type": "application/json",
+ },
+ body: JSON.stringify({
+ model: "gpt-5.2",
+ messages: [{ role: "user", content: "hello" }],
+ stream: true,
+ }),
+ });
+
+ const response = await POST(request, {
+ params: Promise.resolve({ path: ["chat", "completions"] }),
+ });
+
+ expect(response.status).toBe(200);
+ expect(resolveBillingModelPrice).toHaveBeenCalledWith("gpt-unpriced");
+ expect(forwardRequest).toHaveBeenCalledTimes(1);
+ });
+
it("should route path capability request without model when route is matched", async () => {
const { db } = await import("@/lib/db");
const { forwardRequest, prepareUpstreamForProxy } = await import("@/lib/services/proxy-client");
diff --git a/tests/unit/hooks/use-api-keys.test.ts b/tests/unit/hooks/use-api-keys.test.ts
index 34c2725a..2107df23 100644
--- a/tests/unit/hooks/use-api-keys.test.ts
+++ b/tests/unit/hooks/use-api-keys.test.ts
@@ -86,6 +86,9 @@ describe("use-api-keys hooks", () => {
description: null,
access_mode: "restricted",
upstream_ids: ["upstream-1"],
+ spending_rules: null,
+ spending_rule_statuses: [],
+ is_quota_exceeded: false,
is_active: true,
expires_at: null,
created_at: "2024-01-01T00:00:00Z",
@@ -194,6 +197,64 @@ describe("use-api-keys hooks", () => {
expect(mockToastError).toHaveBeenCalledWith("createFailed: Creation failed");
});
+
+ it("preserves spending quota fields in create payload and response", async () => {
+ const spendingRules = [
+ { period_type: "daily" as const, limit: 25 },
+ { period_type: "rolling" as const, limit: 10, period_hours: 6 },
+ ];
+ const spendingRuleStatuses = [
+ {
+ period_type: "daily" as const,
+ period_hours: null,
+ current_spending: 12.5,
+ spending_limit: 25,
+ percent_used: 50,
+ is_exceeded: false,
+ resets_at: "2024-01-02T00:00:00Z",
+ estimated_recovery_at: null,
+ },
+ {
+ period_type: "rolling" as const,
+ period_hours: 6,
+ current_spending: 10,
+ spending_limit: 10,
+ percent_used: 100,
+ is_exceeded: true,
+ resets_at: null,
+ estimated_recovery_at: "2024-01-01T06:30:00Z",
+ },
+ ];
+ const mockResponse = {
+ ...makeApiKey({
+ id: "new-key-id",
+ key_prefix: "sk-auto-newk",
+ name: "New Key",
+ spending_rules: spendingRules,
+ spending_rule_statuses: spendingRuleStatuses,
+ is_quota_exceeded: true,
+ }),
+ key_value: "sk-auto-newkey123",
+ };
+ mockPost.mockResolvedValueOnce(mockResponse);
+
+ const { result } = renderHook(() => useCreateAPIKey(), { wrapper });
+
+ result.current.mutate({
+ name: "New Key",
+ upstream_ids: ["upstream-1"],
+ spending_rules: spendingRules,
+ });
+
+ await waitFor(() => expect(result.current.isSuccess).toBe(true));
+
+ expect(mockPost).toHaveBeenCalledWith("/admin/keys", {
+ name: "New Key",
+ upstream_ids: ["upstream-1"],
+ spending_rules: spendingRules,
+ });
+ expect(result.current.data).toEqual(mockResponse);
+ });
});
describe("useRevealAPIKey", () => {
@@ -401,6 +462,61 @@ describe("use-api-keys hooks", () => {
upstream_ids: ["upstream-2", "upstream-3"],
});
});
+
+ it("preserves spending quota fields in update payload and response", async () => {
+ const spendingRules = [
+ { period_type: "monthly" as const, limit: 200 },
+ { period_type: "rolling" as const, limit: 40, period_hours: 12 },
+ ];
+ const spendingRuleStatuses = [
+ {
+ period_type: "monthly" as const,
+ period_hours: null,
+ current_spending: 120,
+ spending_limit: 200,
+ percent_used: 60,
+ is_exceeded: false,
+ resets_at: "2024-02-01T00:00:00Z",
+ estimated_recovery_at: null,
+ },
+ {
+ period_type: "rolling" as const,
+ period_hours: 12,
+ current_spending: 45,
+ spending_limit: 40,
+ percent_used: 112.5,
+ is_exceeded: true,
+ resets_at: null,
+ estimated_recovery_at: "2024-01-02T03:00:00Z",
+ },
+ ];
+ const mockResponse = makeApiKey({
+ id: "key-1",
+ name: "Updated Key",
+ description: "Updated description",
+ updated_at: "2024-01-02T00:00:00Z",
+ spending_rules: spendingRules,
+ spending_rule_statuses: spendingRuleStatuses,
+ is_quota_exceeded: true,
+ });
+ mockPut.mockResolvedValueOnce(mockResponse);
+
+ const { result } = renderHook(() => useUpdateAPIKey(), { wrapper });
+
+ result.current.mutate({
+ id: "key-1",
+ data: {
+ spending_rules: spendingRules,
+ },
+ });
+
+ await waitFor(() => expect(result.current.isSuccess).toBe(true));
+
+ expect(mockPut).toHaveBeenCalledWith("/admin/keys/key-1", {
+ spending_rules: spendingRules,
+ });
+ expect(result.current.data).toEqual(mockResponse);
+ });
});
describe("useToggleAPIKeyActive", () => {
diff --git a/tests/unit/services/api-key-quota-tracker.test.ts b/tests/unit/services/api-key-quota-tracker.test.ts
new file mode 100644
index 00000000..4e15022c
--- /dev/null
+++ b/tests/unit/services/api-key-quota-tracker.test.ts
@@ -0,0 +1,194 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+const { findManyMock, selectMock, fromMock, whereMock } = vi.hoisted(() => ({
+ findManyMock: vi.fn(async () => [] as unknown[]),
+ selectMock: vi.fn(),
+ fromMock: vi.fn(),
+ whereMock: vi.fn(),
+}));
+
+vi.mock("drizzle-orm", async (importOriginal) => {
+ const actual = await importOriginal();
+ return {
+ ...actual,
+ eq: vi.fn(() => ({ __op: "eq" })),
+ gte: vi.fn(() => ({ __op: "gte" })),
+ lt: vi.fn(() => ({ __op: "lt" })),
+ and: vi.fn((...args: unknown[]) => ({ __op: "and", args })),
+ isNotNull: vi.fn(() => ({ __op: "isNotNull" })),
+ sum: vi.fn(() => ({ __op: "sum" })),
+ };
+});
+
+vi.mock("@/lib/db", () => ({
+ db: {
+ query: {
+ apiKeys: {
+ findMany: findManyMock,
+ },
+ },
+ select: selectMock,
+ },
+ requestBillingSnapshots: {
+ apiKeyId: "api_key_id",
+ billingStatus: "billing_status",
+ billedAt: "billed_at",
+ finalCost: "final_cost",
+ },
+ apiKeys: {
+ id: "id",
+ name: "name",
+ spendingRules: "spending_rules",
+ },
+}));
+
+import { apiKeyQuotaTracker } from "@/lib/services/api-key-quota-tracker";
+
+describe("api-key-quota-tracker", () => {
+ beforeEach(() => {
+ vi.clearAllMocks();
+ apiKeyQuotaTracker.reset();
+ });
+
+ describe("isWithinQuota", () => {
+ it("returns true when no rules are configured", () => {
+ expect(apiKeyQuotaTracker.isWithinQuota("missing-key")).toBe(true);
+ });
+
+ it("returns false when any rule is exceeded", () => {
+ apiKeyQuotaTracker.setRules(
+ "key-1",
+ [
+ { period_type: "daily", limit: 100 },
+ { period_type: "rolling", limit: 30, period_hours: 6 },
+ ],
+ "Quota Key"
+ );
+ apiKeyQuotaTracker.adjustSpending("key-1", 35);
+
+ expect(apiKeyQuotaTracker.isWithinQuota("key-1")).toBe(false);
+ });
+ });
+
+ describe("adjustSpending", () => {
+ it("applies deltas across all cached rules without dropping below zero", () => {
+ apiKeyQuotaTracker.setRules("key-2", [{ period_type: "daily", limit: 100 }], "Adjust Key");
+ apiKeyQuotaTracker.adjustSpending("key-2", 40);
+ apiKeyQuotaTracker.adjustSpending("key-2", -55);
+
+ const entries = apiKeyQuotaTracker.getCacheEntries("key-2");
+ expect(entries).toHaveLength(1);
+ expect(entries?.[0]?.currentSpending).toBe(0);
+ });
+ });
+
+ describe("getQuotaStatus", () => {
+ it("returns multi-rule status with AND semantics", () => {
+ apiKeyQuotaTracker.setRules(
+ "key-3",
+ [
+ { period_type: "daily", limit: 200 },
+ { period_type: "rolling", limit: 50, period_hours: 6 },
+ ],
+ "Status Key"
+ );
+ apiKeyQuotaTracker.adjustSpending("key-3", 80);
+
+ const status = apiKeyQuotaTracker.getQuotaStatus("key-3");
+
+ expect(status?.apiKeyName).toBe("Status Key");
+ expect(status?.rules).toHaveLength(2);
+ expect(status?.rules[0]).toEqual(
+ expect.objectContaining({
+ periodType: "daily",
+ currentSpending: 80,
+ spendingLimit: 200,
+ isExceeded: false,
+ })
+ );
+ expect(status?.rules[1]).toEqual(
+ expect.objectContaining({
+ periodType: "rolling",
+ periodHours: 6,
+ currentSpending: 80,
+ spendingLimit: 50,
+ isExceeded: true,
+ })
+ );
+ expect(status?.isExceeded).toBe(true);
+ });
+ });
+
+ describe("syncFromDb", () => {
+ it("loads configured API keys and aggregates billed spending per rule", async () => {
+ findManyMock.mockResolvedValue([
+ {
+ id: "key-4",
+ name: "Sync Key",
+ spendingRules: [
+ { period_type: "daily", limit: 100 },
+ { period_type: "rolling", limit: 30, period_hours: 5 },
+ ],
+ },
+ ]);
+ whereMock.mockResolvedValue([{ totalCost: "42.50" }]);
+ fromMock.mockReturnValue({ where: whereMock });
+ selectMock.mockReturnValue({ from: fromMock });
+
+ await apiKeyQuotaTracker.syncFromDb();
+
+ const status = apiKeyQuotaTracker.getQuotaStatus("key-4");
+ expect(status?.apiKeyName).toBe("Sync Key");
+ expect(status?.rules).toHaveLength(2);
+ expect(status?.rules[0]?.currentSpending).toBe(42.5);
+ expect(status?.rules[1]?.currentSpending).toBe(42.5);
+ });
+ });
+
+ describe("syncApiKeyFromDb", () => {
+ it("loads a single API key and clears cache when rules are removed", async () => {
+ whereMock.mockResolvedValue([{ totalCost: "15.25" }]);
+ fromMock.mockReturnValue({ where: whereMock });
+ selectMock.mockReturnValue({ from: fromMock });
+
+ await apiKeyQuotaTracker.syncApiKeyFromDb("key-5", "Single Key", [
+ { period_type: "daily", limit: 100 },
+ ]);
+ expect(apiKeyQuotaTracker.getQuotaStatus("key-5")?.rules[0]?.currentSpending).toBe(15.25);
+
+ await apiKeyQuotaTracker.syncApiKeyFromDb("key-5", "Single Key", null);
+ expect(apiKeyQuotaTracker.getQuotaStatus("key-5")).toBeNull();
+ expect(apiKeyQuotaTracker.getCacheEntries("key-5")).toBeUndefined();
+ });
+ });
+
+ describe("estimateRecoveryTime", () => {
+ it("estimates rolling recovery time from billed snapshot history", async () => {
+ vi.useFakeTimers();
+ vi.setSystemTime(new Date("2025-06-15T14:30:00Z"));
+
+ apiKeyQuotaTracker.setRules(
+ "key-6",
+ [{ period_type: "rolling", limit: 100, period_hours: 24 }],
+ "Rolling Key"
+ );
+ apiKeyQuotaTracker.adjustSpending("key-6", 150);
+
+ whereMock.mockResolvedValueOnce([
+ { billedAt: new Date("2025-06-14T14:40:00Z"), finalCost: 60 },
+ ]);
+ fromMock.mockReturnValue({ where: whereMock });
+ selectMock.mockReturnValue({ from: fromMock });
+
+ const recovery = await apiKeyQuotaTracker.estimateRecoveryTime("key-6", {
+ period_type: "rolling",
+ limit: 100,
+ period_hours: 24,
+ });
+
+ expect(recovery).toEqual(new Date("2025-06-15T15:30:00Z"));
+
+ vi.useRealTimers();
+ });
+ });
+});
diff --git a/tests/unit/services/key-manager.test.ts b/tests/unit/services/key-manager.test.ts
index 69303de1..875a2190 100644
--- a/tests/unit/services/key-manager.test.ts
+++ b/tests/unit/services/key-manager.test.ts
@@ -1,6 +1,16 @@
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
import { generateApiKey, ApiKeyNotFoundError, LegacyApiKeyError } from "@/lib/services/key-manager";
+const { mockApiKeyQuotaTracker } = vi.hoisted(() => ({
+ mockApiKeyQuotaTracker: {
+ initialize: vi.fn(async () => {}),
+ getQuotaStatus: vi.fn(() => null),
+ estimateRecoveryTime: vi.fn(async () => null),
+ syncApiKeyFromDb: vi.fn(async () => {}),
+ reset: vi.fn(),
+ },
+}));
+
// Mock the database module
vi.mock("@/lib/db", () => {
// Create db methods that can be reused for both db and transaction
@@ -54,6 +64,10 @@ vi.mock("@/lib/db", () => {
};
});
+vi.mock("@/lib/services/api-key-quota-tracker", () => ({
+ apiKeyQuotaTracker: mockApiKeyQuotaTracker,
+}));
+
// Mock auth utilities
vi.mock("@/lib/utils/auth", () => ({
hashApiKey: vi.fn().mockResolvedValue("hashed-key"),
@@ -246,6 +260,84 @@ describe("key-manager", () => {
expect(result.upstreamIds).toEqual(["upstream-1"]);
});
+ it("should still return the created key when quota tracker sync fails after persistence", async () => {
+ const { db } = await import("@/lib/db");
+ const { createApiKey } = await import("@/lib/services/key-manager");
+
+ const mockApiKey = {
+ id: "key-1",
+ keyHash: "hashed-key",
+ keyValueEncrypted: "encrypted:sk-auto-test123",
+ keyPrefix: "sk-auto-test",
+ name: "Test Key",
+ description: "Test description",
+ accessMode: "restricted",
+ spendingRules: null,
+ isActive: true,
+ expiresAt: null,
+ createdAt: new Date(),
+ updatedAt: new Date(),
+ };
+
+ vi.mocked(db.query.upstreams.findMany).mockResolvedValueOnce([
+ { id: "upstream-1", name: "OpenAI" },
+ ] as never);
+ vi.mocked(db.insert).mockReturnValue({
+ values: vi.fn().mockReturnValue({
+ returning: vi.fn().mockResolvedValue([mockApiKey]),
+ }),
+ } as unknown as ReturnType);
+ mockApiKeyQuotaTracker.syncApiKeyFromDb.mockRejectedValueOnce(new Error("sync failed"));
+
+ const result = await createApiKey({
+ name: "Test Key",
+ upstreamIds: ["upstream-1"],
+ description: "Test description",
+ });
+
+ expect(result.name).toBe("Test Key");
+ expect(result.upstreamIds).toEqual(["upstream-1"]);
+ expect(result.spendingRuleStatuses).toEqual([]);
+ expect(result.isQuotaExceeded).toBe(false);
+ });
+
+ it("should still return created key when quota tracker sync fails after persistence", async () => {
+ const { db } = await import("@/lib/db");
+ const { createApiKey } = await import("@/lib/services/key-manager");
+
+ const mockApiKey = {
+ id: "key-1",
+ keyHash: "hashed-key",
+ keyValueEncrypted: "encrypted:sk-auto-test123",
+ keyPrefix: "sk-auto-test",
+ name: "Test Key",
+ description: null,
+ accessMode: "unrestricted",
+ spendingRules: null,
+ isActive: true,
+ expiresAt: null,
+ createdAt: new Date(),
+ updatedAt: new Date(),
+ };
+
+ vi.mocked(db.insert).mockReturnValue({
+ values: vi.fn().mockReturnValue({
+ returning: vi.fn().mockResolvedValue([mockApiKey]),
+ }),
+ } as unknown as ReturnType);
+ mockApiKeyQuotaTracker.syncApiKeyFromDb.mockRejectedValueOnce(new Error("sync failed"));
+
+ const result = await createApiKey({
+ name: "Test Key",
+ upstreamIds: [],
+ accessMode: "unrestricted",
+ });
+
+ expect(result.id).toBe("key-1");
+ expect(result.spendingRuleStatuses).toEqual([]);
+ expect(result.isQuotaExceeded).toBe(false);
+ });
+
it("should create API key with expiration date", async () => {
const { db } = await import("@/lib/db");
const { createApiKey } = await import("@/lib/services/key-manager");
@@ -334,6 +426,26 @@ describe("key-manager", () => {
await expect(deleteApiKey("key-1")).resolves.toBeUndefined();
expect(db.delete).toHaveBeenCalled();
});
+
+ it("should still resolve delete when quota tracker sync fails after deletion", async () => {
+ const { db } = await import("@/lib/db");
+ const { deleteApiKey } = await import("@/lib/services/key-manager");
+
+ vi.mocked(db.query.apiKeys.findFirst).mockResolvedValueOnce({
+ id: "key-1",
+ keyPrefix: "sk-auto-test",
+ name: "Test Key",
+ } as never);
+
+ const mockWhere = vi.fn().mockResolvedValue(undefined);
+ vi.mocked(db.delete).mockReturnValue({
+ where: mockWhere,
+ } as unknown as ReturnType);
+ mockApiKeyQuotaTracker.syncApiKeyFromDb.mockRejectedValueOnce(new Error("sync failed"));
+
+ await expect(deleteApiKey("key-1")).resolves.toBeUndefined();
+ expect(db.delete).toHaveBeenCalled();
+ });
});
describe("listApiKeys", () => {
@@ -884,6 +996,61 @@ describe("key-manager", () => {
expect(result.description).toBe("New description");
});
+ it("should still return updated key when quota status hydration fails after update", async () => {
+ const { db } = await import("@/lib/db");
+ const { updateApiKey } = await import("@/lib/services/key-manager");
+
+ const mockExistingKey = {
+ id: "key-1",
+ keyPrefix: "sk-auto-test",
+ name: "Test Key",
+ description: null,
+ accessMode: "restricted",
+ spendingRules: [{ period_type: "daily", limit: 50 }],
+ isActive: true,
+ expiresAt: null,
+ createdAt: new Date("2024-01-01"),
+ updatedAt: new Date("2024-01-01"),
+ };
+
+ const mockUpdatedKey = {
+ ...mockExistingKey,
+ name: "Updated Key",
+ updatedAt: new Date(),
+ };
+
+ vi.mocked(db.query.apiKeys.findFirst).mockResolvedValueOnce(mockExistingKey as never);
+
+ const mockReturning = vi.fn().mockResolvedValue([mockUpdatedKey]);
+ const mockWhere = vi.fn(() => ({ returning: mockReturning }));
+ const mockSet = vi.fn(() => ({ where: mockWhere }));
+ vi.mocked(db.update).mockReturnValue({ set: mockSet } as unknown as ReturnType<
+ typeof db.update
+ >);
+
+ vi.mocked(db.query.apiKeyUpstreams.findMany).mockResolvedValueOnce([
+ { apiKeyId: "key-1", upstreamId: "upstream-1" },
+ ] as never);
+ mockApiKeyQuotaTracker.initialize.mockRejectedValueOnce(new Error("quota init failed"));
+
+ const result = await updateApiKey("key-1", { name: "Updated Key" });
+
+ expect(result.name).toBe("Updated Key");
+ expect(result.isQuotaExceeded).toBe(false);
+ expect(result.spendingRuleStatuses).toEqual([
+ {
+ periodType: "daily",
+ periodHours: null,
+ currentSpending: 0,
+ spendingLimit: 50,
+ percentUsed: 0,
+ isExceeded: false,
+ resetsAt: null,
+ estimatedRecoveryAt: null,
+ },
+ ]);
+ });
+
it("should update isActive status", async () => {
const { db } = await import("@/lib/db");
const { updateApiKey } = await import("@/lib/services/key-manager");
diff --git a/tests/unit/services/stats-service.test.ts b/tests/unit/services/stats-service.test.ts
index edb38264..3c30afec 100644
--- a/tests/unit/services/stats-service.test.ts
+++ b/tests/unit/services/stats-service.test.ts
@@ -895,6 +895,130 @@ describe("stats-service", () => {
expect(result.apiKeys[0].totalTokens).toBe(0);
});
+ it("should include quota-rejected requests for api keys but exclude them from upstreams", async () => {
+ const { db } = await import("@/lib/db");
+ const { getLeaderboardStats } = await import("@/lib/services/stats-service");
+
+ const emptyMainQuery = {
+ from: vi.fn().mockReturnValue({
+ where: vi.fn().mockReturnValue({
+ groupBy: vi.fn().mockReturnValue({
+ orderBy: vi.fn().mockReturnValue({
+ limit: vi.fn().mockResolvedValue([]),
+ }),
+ }),
+ }),
+ }),
+ } as unknown as ReturnType;
+
+ vi.mocked(db.select)
+ .mockReturnValueOnce({
+ from: vi.fn().mockReturnValue({
+ where: vi.fn().mockReturnValue({
+ groupBy: vi.fn().mockReturnValue({
+ orderBy: vi.fn().mockReturnValue({
+ limit: vi
+ .fn()
+ .mockResolvedValue([
+ { apiKeyId: "key-1", requestCount: 3, totalTokens: "100" },
+ ]),
+ }),
+ }),
+ }),
+ }),
+ } as unknown as ReturnType)
+ .mockReturnValueOnce({
+ from: vi.fn().mockReturnValue({
+ leftJoin: vi.fn().mockReturnValue({
+ where: vi.fn().mockReturnValue({
+ groupBy: vi.fn().mockResolvedValue([{ apiKeyId: "key-1", totalCost: "1.25" }]),
+ }),
+ }),
+ }),
+ } as unknown as ReturnType)
+ .mockReturnValueOnce({
+ from: vi.fn().mockReturnValue({
+ where: vi.fn().mockReturnValue({
+ groupBy: vi.fn().mockResolvedValue([{ groupKey: "key-1", name: "gpt-4o", cnt: 2 }]),
+ }),
+ }),
+ } as unknown as ReturnType)
+ .mockReturnValueOnce({
+ from: vi.fn().mockReturnValue({
+ where: vi.fn().mockReturnValue({
+ groupBy: vi.fn().mockReturnValue({
+ orderBy: vi.fn().mockReturnValue({
+ limit: vi.fn().mockResolvedValue([
+ {
+ upstreamId: "upstream-1",
+ requestCount: 2,
+ totalTokens: "80",
+ avgTtft: null,
+ totalCompletionTokens: 0,
+ totalDurationMs: 0,
+ totalCacheReadTokens: null,
+ totalEffectivePromptTokens: 0,
+ },
+ ]),
+ }),
+ }),
+ }),
+ }),
+ } as unknown as ReturnType)
+ .mockReturnValueOnce({
+ from: vi.fn().mockReturnValue({
+ leftJoin: vi.fn().mockReturnValue({
+ where: vi.fn().mockReturnValue({
+ groupBy: vi
+ .fn()
+ .mockResolvedValue([{ upstreamId: "upstream-1", totalCost: "1.25" }]),
+ }),
+ }),
+ }),
+ } as unknown as ReturnType)
+ .mockReturnValueOnce({
+ from: vi.fn().mockReturnValue({
+ where: vi.fn().mockReturnValue({
+ groupBy: vi
+ .fn()
+ .mockResolvedValue([{ groupKey: "upstream-1", name: "gpt-4o", cnt: 2 }]),
+ }),
+ }),
+ } as unknown as ReturnType)
+ .mockReturnValueOnce(emptyMainQuery);
+
+ vi.mocked(db.query.apiKeys.findMany).mockResolvedValueOnce([
+ { id: "key-1", name: "Quota Key", keyPrefix: "sk-quota-" },
+ ]);
+ vi.mocked(db.query.upstreams.findMany).mockResolvedValueOnce([
+ {
+ id: "upstream-1",
+ name: "Primary Upstream",
+ routeCapabilities: ["openai_chat_compatible"],
+ },
+ ]);
+
+ const result = await getLeaderboardStats("7d", 5);
+
+ expect(result.apiKeys).toHaveLength(1);
+ expect(result.apiKeys[0]).toMatchObject({
+ id: "key-1",
+ requestCount: 3,
+ totalTokens: 100,
+ totalCostUsd: 1.25,
+ modelDistribution: [{ name: "gpt-4o", count: 2 }],
+ });
+
+ expect(result.upstreams).toHaveLength(1);
+ expect(result.upstreams[0]).toMatchObject({
+ id: "upstream-1",
+ requestCount: 2,
+ totalTokens: 80,
+ totalCostUsd: 1.25,
+ modelDistribution: [{ name: "gpt-4o", count: 2 }],
+ });
+ });
+
it("should calculate upstream avgTps from aggregated stream fields", async () => {
const { db } = await import("@/lib/db");
const { getLeaderboardStats } = await import("@/lib/services/stats-service");
diff --git a/tests/unit/utils/api-transformers.test.ts b/tests/unit/utils/api-transformers.test.ts
index 8458c170..17d7a51a 100644
--- a/tests/unit/utils/api-transformers.test.ts
+++ b/tests/unit/utils/api-transformers.test.ts
@@ -206,6 +206,20 @@ describe("api-transformers", () => {
description: "Main production API key",
accessMode: "restricted",
upstreamIds: ["upstream-1", "upstream-2"],
+ spendingRules: [{ period_type: "daily", limit: 25 }],
+ spendingRuleStatuses: [
+ {
+ periodType: "daily",
+ periodHours: null,
+ currentSpending: 12.5,
+ spendingLimit: 25,
+ percentUsed: 50,
+ isExceeded: false,
+ resetsAt: new Date("2025-01-16T00:00:00.000Z"),
+ estimatedRecoveryAt: null,
+ },
+ ],
+ isQuotaExceeded: false,
isActive: true,
expiresAt: new Date("2025-01-15T00:00:00.000Z"),
createdAt: new Date("2024-01-15T10:00:00.000Z"),
@@ -221,6 +235,20 @@ describe("api-transformers", () => {
description: "Main production API key",
access_mode: "restricted",
upstream_ids: ["upstream-1", "upstream-2"],
+ spending_rules: [{ period_type: "daily", limit: 25 }],
+ spending_rule_statuses: [
+ {
+ period_type: "daily",
+ period_hours: null,
+ current_spending: 12.5,
+ spending_limit: 25,
+ percent_used: 50,
+ is_exceeded: false,
+ resets_at: "2025-01-16T00:00:00.000Z",
+ estimated_recovery_at: null,
+ },
+ ],
+ is_quota_exceeded: false,
is_active: true,
expires_at: "2025-01-15T00:00:00.000Z",
created_at: "2024-01-15T10:00:00.000Z",
@@ -236,6 +264,9 @@ describe("api-transformers", () => {
description: null,
accessMode: "unrestricted",
upstreamIds: [],
+ spendingRules: null,
+ spendingRuleStatuses: [],
+ isQuotaExceeded: false,
isActive: false,
expiresAt: null,
createdAt: new Date("2024-01-10T08:00:00.000Z"),
@@ -262,6 +293,20 @@ describe("api-transformers", () => {
description: "Newly created key",
accessMode: "restricted",
upstreamIds: ["upstream-1"],
+ spendingRules: [{ period_type: "rolling", limit: 10, period_hours: 6 }],
+ spendingRuleStatuses: [
+ {
+ periodType: "rolling",
+ periodHours: 6,
+ currentSpending: 10,
+ spendingLimit: 10,
+ percentUsed: 100,
+ isExceeded: true,
+ resetsAt: null,
+ estimatedRecoveryAt: new Date("2024-01-15T16:00:00.000Z"),
+ },
+ ],
+ isQuotaExceeded: true,
isActive: true,
expiresAt: new Date("2025-06-01T00:00:00.000Z"),
createdAt: new Date("2024-01-15T10:00:00.000Z"),
@@ -278,6 +323,20 @@ describe("api-transformers", () => {
description: "Newly created key",
access_mode: "restricted",
upstream_ids: ["upstream-1"],
+ spending_rules: [{ period_type: "rolling", limit: 10, period_hours: 6 }],
+ spending_rule_statuses: [
+ {
+ period_type: "rolling",
+ period_hours: 6,
+ current_spending: 10,
+ spending_limit: 10,
+ percent_used: 100,
+ is_exceeded: true,
+ resets_at: null,
+ estimated_recovery_at: "2024-01-15T16:00:00.000Z",
+ },
+ ],
+ is_quota_exceeded: true,
is_active: true,
expires_at: "2025-06-01T00:00:00.000Z",
created_at: "2024-01-15T10:00:00.000Z",
@@ -294,6 +353,9 @@ describe("api-transformers", () => {
description: null,
accessMode: "unrestricted",
upstreamIds: [],
+ spendingRules: null,
+ spendingRuleStatuses: [],
+ isQuotaExceeded: false,
isActive: true,
expiresAt: null,
createdAt: new Date("2024-01-15T10:00:00.000Z"),
@@ -337,7 +399,22 @@ describe("api-transformers", () => {
keyPrefix: "ar_live_",
name: "Key 1",
description: "First key",
+ accessMode: "restricted",
upstreamIds: ["upstream-1"],
+ spendingRules: [{ period_type: "monthly", limit: 100 }],
+ spendingRuleStatuses: [
+ {
+ periodType: "monthly",
+ periodHours: null,
+ currentSpending: 12,
+ spendingLimit: 100,
+ percentUsed: 12,
+ isExceeded: false,
+ resetsAt: new Date("2024-02-01T00:00:00.000Z"),
+ estimatedRecoveryAt: null,
+ },
+ ],
+ isQuotaExceeded: false,
isActive: true,
expiresAt: null,
createdAt: new Date("2024-01-15T10:00:00.000Z"),
@@ -358,6 +435,8 @@ describe("api-transformers", () => {
expect(result.page_size).toBe(15);
expect(result.total_pages).toBe(4);
expect(result.items[0].key_prefix).toBe("ar_live_");
+ expect(result.items[0].spending_rules).toEqual([{ period_type: "monthly", limit: 100 }]);
+ expect(result.items[0].is_quota_exceeded).toBe(false);
});
});