feat(rate-limit): 实现分层 RPM 限速,支持系统默认/用户/独立Key三级配置

- 新增用户级 rate_limit 字段,支持系统默认/用户自定义/不限制三种模式
- 独立 Key 的 rate_limit 语义调整:null=跟随系统默认,0=不限制,>0=自定义
- 实现 UserRpmLimiter 基于 Redis sliding window 的 RPM 限速引擎
- Pipeline 请求流程集成用户级 RPM 检查
- 管理后台和用户面板新增 RPM 限速配置与实时状态查看
- 系统设置新增全局默认 RPM 配置项
- 迁移脚本回填现有 API Key 的 rate_limit 默认值
- 新增用户/Key RPM 状态监控 API 和前端展示

Closes #231

Co-authored-by: LewisPen <LewisPen@nyadoo.com>
This commit is contained in:
fawney19
2026-03-15 14:22:59 +08:00
co-authored by LewisPen
parent 920a383136
commit f92b0943b5
35 changed files with 2051 additions and 238 deletions
@@ -0,0 +1,54 @@
"""add user rate_limit and backfill normal api key limits
Revision ID: b7e8f9a0c1d2
Revises: b7c8d9e0f1a2
Create Date: 2026-03-13 12:00:00.000000+00:00
"""
from __future__ import annotations
from collections.abc import Sequence
import sqlalchemy as sa
from sqlalchemy import inspect
from alembic import op
# revision identifiers, used by Alembic.
revision: str = "b7e8f9a0c1d2"
down_revision: str | None = "b7c8d9e0f1a2"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def column_exists(table_name: str, column_name: str) -> bool:
bind = op.get_bind()
inspector = inspect(bind)
columns = [c["name"] for c in inspector.get_columns(table_name)]
return column_name in columns
def upgrade() -> None:
if not column_exists("users", "rate_limit"):
op.add_column("users", sa.Column("rate_limit", sa.Integer(), nullable=True))
# 普通 Key 新语义不再允许 NULL;存量 NULL 统一回填为 0(不限制)。
op.execute(sa.text("""
UPDATE api_keys
SET rate_limit = 0
WHERE is_standalone = FALSE
AND rate_limit IS NULL
"""))
def downgrade() -> None:
# 恢复普通 Key 的 rate_limit 为 NULL(与 upgrade 中回填 0 对应)
op.execute(sa.text("""
UPDATE api_keys
SET rate_limit = NULL
WHERE is_standalone = FALSE
AND rate_limit = 0
"""))
if column_exists("users", "rate_limit"):
op.drop_column("users", "rate_limit")
+4 -3
View File
@@ -86,6 +86,7 @@ export interface UserExport {
allowed_providers?: string[] | null
allowed_api_formats?: string[] | null
allowed_models?: string[] | null
rate_limit?: number | null // null = 跟随系统默认,0 = 不限制
model_capability_settings?: Record<string, Record<string, boolean>>
unlimited?: boolean
wallet?: BillingSummary | null
@@ -102,7 +103,7 @@ export interface UserApiKeyExport {
allowed_providers?: string[] | null
allowed_api_formats?: string[] | null
allowed_models?: string[] | null
rate_limit?: number | null // null = 无限制
rate_limit?: number | null // legacy/null 兼容;1.3+ standalone null = 跟随系统默认
concurrent_limit?: number | null
force_capabilities?: Record<string, boolean>
is_active: boolean
@@ -349,7 +350,7 @@ export interface AdminApiKey {
total_requests?: number
total_tokens?: number
total_cost_usd?: number
rate_limit?: number | null // null = 限制
rate_limit?: number | null // null = 跟随系统默认,0 = 不限制
allowed_providers?: string[] | null // 允许的提供商列表
allowed_api_formats?: string[] | null // 允许的 API 格式列表
allowed_models?: string[] | null // 允许的模型列表
@@ -365,7 +366,7 @@ export interface CreateStandaloneApiKeyRequest {
allowed_providers?: string[] | null
allowed_api_formats?: string[] | null
allowed_models?: string[] | null
rate_limit?: number | null // null = 限制
rate_limit?: number | null // null = 跟随系统默认,0 = 不限制
expires_at?: string | null // ISO 日期字符串,如 "2025-12-31"null = 永不过期
initial_balance_usd: number | null // 初始余额,null = 无限制
unlimited_balance?: boolean | null // 编辑时仅切换额度模式,不调整余额数值
+14 -2
View File
@@ -142,6 +142,7 @@ export interface ApiKey {
created_at: string
total_requests?: number
total_cost_usd?: number
rate_limit?: number | null
allowed_providers?: ProviderConfig[]
force_capabilities?: Record<string, boolean> | null // 强制能力配置
}
@@ -181,8 +182,8 @@ export const meApi = {
return response.data
},
async createApiKey(name: string): Promise<ApiKey> {
const response = await apiClient.post<ApiKey>('/api/users/me/api-keys', { name })
async createApiKey(data: { name: string; rate_limit?: number }): Promise<ApiKey> {
const response = await apiClient.post<ApiKey>('/api/users/me/api-keys', data)
return response.data
},
@@ -212,6 +213,17 @@ export const meApi = {
return response.data
},
async updateApiKey(
keyId: string,
data: { name?: string; rate_limit?: number | null }
): Promise<ApiKey & { message: string }> {
const response = await apiClient.put<ApiKey & { message: string }>(
`/api/users/me/api-keys/${keyId}`,
data
)
return response.data
},
// 使用统计
async getUsage(params?: {
start_date?: string
+29 -3
View File
@@ -10,6 +10,7 @@ export interface User {
allowed_providers: string[] | null // 允许使用的提供商 ID 列表
allowed_api_formats: string[] | null // 允许使用的 API 格式列表
allowed_models: string[] | null // 允许使用的模型名称列表
rate_limit?: number | null // null = 跟随系统默认,0 = 不限制
created_at: string
updated_at?: string
last_login_at?: string | null
@@ -25,6 +26,7 @@ export interface CreateUserRequest {
allowed_providers?: string[] | null
allowed_api_formats?: string[] | null
allowed_models?: string[] | null
rate_limit?: number | null
}
export interface UpdateUserRequest {
@@ -36,6 +38,7 @@ export interface UpdateUserRequest {
allowed_providers?: string[] | null
allowed_api_formats?: string[] | null
allowed_models?: string[] | null
rate_limit?: number | null
}
export interface ApiKey {
@@ -49,11 +52,16 @@ export interface ApiKey {
is_active: boolean
is_locked: boolean // 管理员锁定标志
is_standalone: boolean // 是否为独立余额Key
rate_limit?: number // 速率限制(请求/分钟)
rate_limit?: number | null // 普通Key: 0 = 不限制,历史 null 视为跟随系统默认
total_requests?: number // 总请求数
total_cost_usd?: number // 总费用
}
export interface UpsertUserApiKeyRequest {
name?: string
rate_limit?: number | null
}
export const usersApi = {
async getAllUsers(): Promise<User[]> {
const response = await apiClient.get<User[]>('/api/admin/users')
@@ -84,8 +92,26 @@ export const usersApi = {
return response.data.api_keys
},
async createApiKey(userId: string, name?: string): Promise<ApiKey & { key: string }> {
const response = await apiClient.post<ApiKey & { key: string }>(`/api/admin/users/${userId}/api-keys`, { name })
async createApiKey(
userId: string,
data: UpsertUserApiKeyRequest
): Promise<ApiKey & { key: string }> {
const response = await apiClient.post<ApiKey & { key: string }>(
`/api/admin/users/${userId}/api-keys`,
data
)
return response.data
},
async updateApiKey(
userId: string,
keyId: string,
data: UpsertUserApiKeyRequest
): Promise<ApiKey & { message: string }> {
const response = await apiClient.put<ApiKey & { message: string }>(
`/api/admin/users/${userId}/api-keys/${keyId}`,
data
)
return response.data
},
@@ -96,23 +96,6 @@
{{ form.expires_at ? '到期后' + (form.auto_delete_on_expiry ? '自动删除' : '仅禁用') + '(当天 23:59 失效)' : '留空表示永不过期' }}
</p>
</div>
<div class="space-y-2">
<Label
for="form-rate-limit"
class="text-sm font-medium"
>速率限制 (请求/分钟)</Label>
<Input
id="form-rate-limit"
:model-value="form.rate_limit ?? ''"
type="number"
min="1"
max="10000"
placeholder="留空不限制"
class="h-10"
@update:model-value="(v) => form.rate_limit = parseNumberInput(v, { min: 1, max: 10000 })"
/>
</div>
</div>
<!-- 右侧:访问限制 -->
@@ -190,6 +173,36 @@
</div>
</div>
<div class="space-y-2">
<Label
for="form-rate-limit"
class="text-sm font-medium"
>速率限制 (请求/分钟)</Label>
<div class="flex items-center gap-3">
<div class="flex-1 min-w-0">
<Input
v-if="!form.rate_limit_inherited"
id="form-rate-limit"
:model-value="form.rate_limit ?? ''"
type="number"
min="0"
max="10000"
placeholder="0 = 不限速"
class="h-10"
@update:model-value="(v) => form.rate_limit = parseNumberInput(v, { min: 0, max: 10000 })"
/>
<span
v-else
class="flex h-10 w-full items-center rounded-lg border bg-background px-3 text-sm text-muted-foreground opacity-60"
>跟随系统</span>
</div>
<Switch
v-model="form.rate_limit_inherited"
class="shrink-0"
/>
</div>
</div>
<!-- 额度 -->
<div class="space-y-2">
<Label class="text-sm font-medium">额度</Label>
@@ -267,7 +280,7 @@ export interface StandaloneKeyFormData {
initial_balance_usd?: number
unlimited_balance?: boolean
expires_at?: string // ISO 日期字符串,如 "2025-12-31"undefined = 永不过期
rate_limit?: number
rate_limit?: number | null
auto_delete_on_expiry: boolean
allowed_providers?: string[] | null
allowed_api_formats?: string[] | null
@@ -280,6 +293,7 @@ interface StandaloneKeyFormState {
initial_balance_usd?: number
unlimited_balance?: boolean
expires_at?: string
rate_limit_inherited: boolean
rate_limit?: number
auto_delete_on_expiry: boolean
provider_unrestricted: boolean
@@ -333,6 +347,7 @@ const form = ref<StandaloneKeyFormState>({
initial_balance_usd: 10,
unlimited_balance: false,
expires_at: undefined,
rate_limit_inherited: true,
rate_limit: undefined,
auto_delete_on_expiry: false,
provider_unrestricted: true,
@@ -356,6 +371,7 @@ function resetForm() {
initial_balance_usd: 10,
unlimited_balance: false,
expires_at: undefined,
rate_limit_inherited: true,
rate_limit: undefined,
auto_delete_on_expiry: false,
provider_unrestricted: true,
@@ -375,7 +391,8 @@ function loadKeyData() {
initial_balance_usd: props.apiKey.initial_balance_usd,
unlimited_balance: props.apiKey.initial_balance_usd == null,
expires_at: props.apiKey.expires_at,
rate_limit: props.apiKey.rate_limit,
rate_limit_inherited: props.apiKey.rate_limit == null,
rate_limit: props.apiKey.rate_limit ?? undefined,
auto_delete_on_expiry: props.apiKey.auto_delete_on_expiry,
provider_unrestricted: props.apiKey.allowed_providers == null,
api_format_unrestricted: props.apiKey.allowed_api_formats == null,
@@ -425,7 +442,7 @@ function handleSubmit() {
initial_balance_usd: form.value.initial_balance_usd,
unlimited_balance: form.value.unlimited_balance,
expires_at: form.value.expires_at,
rate_limit: form.value.rate_limit,
rate_limit: form.value.rate_limit_inherited ? null : (form.value.rate_limit ?? 0),
auto_delete_on_expiry: form.value.auto_delete_on_expiry,
allowed_providers: form.value.provider_unrestricted ? null : [...form.value.allowed_providers],
allowed_api_formats: form.value.api_format_unrestricted ? null : [...form.value.allowed_api_formats],
@@ -256,6 +256,36 @@
</div>
</div>
<div class="space-y-2">
<Label
for="form-rate-limit"
class="text-sm font-medium"
>速率限制 (请求/分钟)</Label>
<div class="flex items-center gap-3">
<div class="flex-1 min-w-0">
<Input
v-if="!form.rate_limit_inherited"
id="form-rate-limit"
:model-value="form.rate_limit ?? ''"
type="number"
min="0"
max="10000"
placeholder="0 = 不限速"
class="h-10"
@update:model-value="(v) => form.rate_limit = parseNumberInput(v, { min: 0, max: 10000 })"
/>
<span
v-else
class="flex h-10 w-full items-center rounded-lg border bg-background px-3 text-sm text-muted-foreground opacity-60"
>跟随系统默认</span>
</div>
<Switch
v-model="form.rate_limit_inherited"
class="shrink-0"
/>
</div>
</div>
<!-- 额度 -->
<div class="space-y-2">
<Label class="text-sm font-medium">额度</Label>
@@ -352,6 +382,7 @@ export interface UserFormData {
allowed_providers?: string[] | null
allowed_api_formats?: string[] | null
allowed_models?: string[] | null
rate_limit?: number | null
}
const props = defineProps<{
@@ -406,9 +437,11 @@ const form = ref({
provider_unrestricted: true,
api_format_unrestricted: true,
model_unrestricted: true,
rate_limit_inherited: true,
allowed_providers: [] as string[],
allowed_api_formats: [] as string[],
allowed_models: [] as string[],
rate_limit: undefined as number | undefined,
})
function createFieldNonce(): string {
@@ -429,9 +462,11 @@ function resetForm() {
provider_unrestricted: true,
api_format_unrestricted: true,
model_unrestricted: true,
rate_limit_inherited: true,
allowed_providers: [],
allowed_api_formats: [],
allowed_models: [],
rate_limit: undefined,
}
}
@@ -451,9 +486,11 @@ function loadUserData() {
provider_unrestricted: props.user.allowed_providers == null,
api_format_unrestricted: props.user.allowed_api_formats == null,
model_unrestricted: props.user.allowed_models == null,
rate_limit_inherited: props.user.rate_limit == null,
allowed_providers: props.user.allowed_providers ? [...props.user.allowed_providers] : [],
allowed_api_formats: props.user.allowed_api_formats ? [...props.user.allowed_api_formats] : [],
allowed_models: props.user.allowed_models ? [...props.user.allowed_models] : [],
rate_limit: props.user.rate_limit ?? undefined,
}
}
@@ -543,6 +580,7 @@ async function handleSubmit() {
allowed_models: form.value.model_unrestricted
? null
: [...form.value.allowed_models],
rate_limit: form.value.rate_limit_inherited ? null : (form.value.rate_limit ?? 0),
}
if (isEditMode.value && props.user?.id) {
+24 -3
View File
@@ -1,6 +1,13 @@
import { defineStore } from 'pinia'
import { ref } from 'vue'
import { usersApi, type User, type CreateUserRequest, type UpdateUserRequest, type ApiKey } from '@/api/users'
import {
usersApi,
type User,
type CreateUserRequest,
type UpdateUserRequest,
type ApiKey,
type UpsertUserApiKeyRequest
} from '@/api/users'
import { parseApiError } from '@/utils/errorParser'
export const useUsersStore = defineStore('users', () => {
@@ -84,15 +91,28 @@ export const useUsersStore = defineStore('users', () => {
}
}
async function createApiKey(userId: string, name?: string): Promise<ApiKey> {
async function createApiKey(userId: string, data: UpsertUserApiKeyRequest): Promise<ApiKey> {
try {
return await usersApi.createApiKey(userId, name)
return await usersApi.createApiKey(userId, data)
} catch (err: unknown) {
error.value = parseApiError(err, '创建 API Key 失败')
throw err
}
}
async function updateApiKey(
userId: string,
keyId: string,
data: UpsertUserApiKeyRequest
): Promise<ApiKey> {
try {
return await usersApi.updateApiKey(userId, keyId, data)
} catch (err: unknown) {
error.value = parseApiError(err, '更新 API Key 失败')
throw err
}
}
async function deleteApiKey(userId: string, keyId: string) {
try {
await usersApi.deleteApiKey(userId, keyId)
@@ -121,6 +141,7 @@ export const useUsersStore = defineStore('users', () => {
deleteUser,
getUserApiKeys,
createApiKey,
updateApiKey,
deleteApiKey,
getFullApiKey
}
+22
View File
@@ -159,3 +159,25 @@ export function formatHitRate(rate: number | undefined): string {
if (typeof rate !== 'number' || Number.isNaN(rate)) return '-'
return `${rate.toFixed(2)}%`
}
// Rate limit formatting (supports "inherit" semantics: null = inherit system default)
export function formatRateLimitInheritable(rateLimit?: number | null): string {
if (rateLimit == null) return '跟随系统'
if (rateLimit === 0) return '不限速'
return `${rateLimit}/min`
}
// Rate limit formatting (simple: null/0 both mean unlimited)
export function formatRateLimitSimple(rateLimit?: number | null): string {
if (rateLimit == null || rateLimit === 0) return '不限速'
return `${rateLimit}/min`
}
// Rate limit state helpers
export function isRateLimitInherited(rateLimit?: number | null): boolean {
return rateLimit == null
}
export function isRateLimitUnlimited(rateLimit?: number | null): boolean {
return rateLimit === 0
}
+36 -13
View File
@@ -105,8 +105,8 @@
<TableHead class="w-[240px] h-12 font-semibold">
钱包
</TableHead>
<TableHead class="w-[130px] h-12 font-semibold">
使用统计
<TableHead class="w-[190px] h-12 font-semibold">
统计/限速
</TableHead>
<TableHead class="w-[110px] h-12 font-semibold">
有效期
@@ -221,7 +221,23 @@
请求: <span class="font-medium text-foreground">{{ (apiKey.total_requests || 0).toLocaleString() }}</span>
</div>
<div class="text-muted-foreground">
速率: <span class="font-medium text-foreground">{{ apiKey.rate_limit ? `${apiKey.rate_limit}/min` : '未设置' }}</span>
Tokens: <span class="font-medium text-foreground">{{ formatTokens(apiKey.total_tokens || 0) }}</span>
</div>
<div class="flex items-center gap-1 text-muted-foreground">
<span>限速:</span>
<Badge
v-if="isRateLimitInherited(apiKey.rate_limit) || isRateLimitUnlimited(apiKey.rate_limit)"
variant="secondary"
class="h-5 px-1.5 py-0 text-[10px] font-medium"
>
{{ formatRateLimitInheritable(apiKey.rate_limit) }}
</Badge>
<span
v-else
class="font-medium text-foreground"
>
{{ formatRateLimitInheritable(apiKey.rate_limit) }}
</span>
</div>
</div>
</TableCell>
@@ -387,6 +403,12 @@
>
{{ walletStatusLabel(getApiKeyWalletStatus(apiKey.id)) }}
</Badge>
<Badge
variant="secondary"
class="h-5 px-1.5 py-0 text-[10px] font-medium"
>
{{ formatRateLimitInheritable(apiKey.rate_limit) }}
</Badge>
<Badge
v-if="apiKey.auto_delete_on_expiry"
variant="secondary"
@@ -429,14 +451,6 @@
</div>
<div class="grid grid-cols-2 gap-2.5 text-xs">
<div class="rounded-lg border border-border/50 bg-background/70 p-2.5">
<div class="mb-1 text-muted-foreground">
速率限制
</div>
<div class="font-semibold text-foreground">
{{ apiKey.rate_limit ? `${apiKey.rate_limit}/min` : '未设置' }}
</div>
</div>
<div class="rounded-lg border border-border/50 bg-background/70 p-2.5">
<div class="mb-1 text-muted-foreground">
请求次数
@@ -445,6 +459,14 @@
{{ (apiKey.total_requests || 0).toLocaleString() }}
</div>
</div>
<div class="rounded-lg border border-border/50 bg-background/70 p-2.5">
<div class="mb-1 text-muted-foreground">
Tokens
</div>
<div class="font-semibold text-foreground">
{{ formatTokens(apiKey.total_tokens || 0) }}
</div>
</div>
<div class="col-span-2 rounded-lg border border-border/50 bg-background/70 p-2.5">
<div class="mb-1 text-muted-foreground">
有效期
@@ -663,6 +685,7 @@ import {
import { StandaloneKeyFormDialog, type StandaloneKeyFormData } from '@/features/api-keys'
import { parseApiError } from '@/utils/errorParser'
import { formatTokens, formatRateLimitInheritable, isRateLimitInherited, isRateLimitUnlimited } from '@/utils/format'
import { log } from '@/utils/logger'
const { success, error } = useToast()
@@ -1028,7 +1051,7 @@ async function handleKeyFormSubmit(data: StandaloneKeyFormData) {
const updateData: Partial<CreateStandaloneApiKeyRequest> = {
name: data.name || undefined,
unlimited_balance: Boolean(data.unlimited_balance),
rate_limit: data.rate_limit ?? null, // undefined = 无限制,显式传 null
rate_limit: data.rate_limit ?? null, // undefined = 跟随系统默认,显式传 null
expires_at: data.expires_at || null, // undefined/空 = 永不过期
auto_delete_on_expiry: data.auto_delete_on_expiry,
// 空数组表示清除限制(允许全部),后端会将空数组存为 NULL
@@ -1057,7 +1080,7 @@ async function handleKeyFormSubmit(data: StandaloneKeyFormData) {
const createData: CreateStandaloneApiKeyRequest = {
name: data.name || undefined,
initial_balance_usd: isUnlimited ? null : (data.initial_balance_usd as number),
rate_limit: data.rate_limit ?? null, // undefined = 无限制,显式传 null
rate_limit: data.rate_limit ?? null, // undefined = 跟随系统默认,显式传 null
expires_at: data.expires_at || null, // undefined/空 = 永不过期
auto_delete_on_expiry: data.auto_delete_on_expiry,
// 空数组表示不设置限制(允许全部),后端会将空数组存为 NULL
+192 -38
View File
@@ -180,7 +180,7 @@
钱包
</TableHead>
<TableHead class="w-[170px] h-12 font-semibold">
使用统计
统计/限速
</TableHead>
<TableHead class="w-[110px] h-12 font-semibold">
创建时间
@@ -258,26 +258,42 @@
</div>
</TableCell>
<TableCell class="py-4">
<div
v-if="userStats[user.id]"
class="space-y-1 text-xs"
>
<div class="flex items-center text-muted-foreground">
<span class="w-14">请求:</span>
<span class="font-medium text-foreground">{{ formatNumber(userStats[user.id]?.request_count) }}</span>
<div class="space-y-1 text-xs">
<template v-if="userStats[user.id]">
<div class="flex items-center text-muted-foreground">
<span class="w-14">请求:</span>
<span class="font-medium text-foreground">{{ formatNumber(userStats[user.id]?.request_count) }}</span>
</div>
<div class="flex items-center text-muted-foreground">
<span class="w-14">Tokens:</span>
<span class="font-medium text-foreground">{{ formatTokens(userStats[user.id]?.total_tokens ?? 0) }}</span>
</div>
</template>
<div
v-else
class="flex items-center text-muted-foreground"
>
<span class="w-14">统计:</span>
<span v-if="loadingStats">加载中...</span>
<span v-else>无数据</span>
</div>
<div class="flex items-center text-muted-foreground">
<span class="w-14">Tokens:</span>
<span class="font-medium text-foreground">{{ formatTokens(userStats[user.id]?.total_tokens ?? 0) }}</span>
<span class="w-14">限速:</span>
<Badge
v-if="isRateLimitInherited(user.rate_limit) || isRateLimitUnlimited(user.rate_limit)"
variant="secondary"
class="h-5 px-1.5 py-0 text-[10px] font-medium"
>
{{ formatRateLimitInheritable(user.rate_limit) }}
</Badge>
<span
v-else
class="font-medium text-foreground"
>
{{ formatRateLimitInheritable(user.rate_limit) }}
</span>
</div>
</div>
<div
v-else
class="text-xs text-muted-foreground"
>
<span v-if="loadingStats">加载中...</span>
<span v-else>无数据</span>
</div>
</TableCell>
<TableCell class="py-4 text-xs text-muted-foreground">
{{ formatDate(user.created_at) }}
@@ -436,6 +452,12 @@
>
{{ walletStatusLabel(getUserWalletStatus(user.id)) }}
</Badge>
<Badge
variant="secondary"
class="h-5 px-1.5 py-0 text-[10px] font-medium"
>
{{ formatRateLimitInheritable(user.rate_limit) }}
</Badge>
</div>
<div class="rounded-xl border border-border/60 bg-muted/40 p-3.5">
@@ -633,6 +655,12 @@
>
独立余额
</Badge>
<Badge
variant="secondary"
class="text-xs"
>
{{ formatRateLimitSimple(apiKey.rate_limit) }}
</Badge>
</div>
<div class="flex items-center gap-1 mt-0.5">
<code class="text-xs font-mono text-muted-foreground">
@@ -658,6 +686,15 @@
${{ (apiKey.total_cost_usd || 0).toFixed(4) }}
</div>
</div>
<Button
variant="ghost"
size="icon"
class="h-8 w-8"
title="编辑"
@click="openEditUserApiKeyDialog(apiKey)"
>
<SquarePen class="h-4 w-4" />
</Button>
<Button
variant="ghost"
size="icon"
@@ -718,13 +755,87 @@
<Button
class="h-10 px-5"
:disabled="creatingApiKey"
@click="createApiKey"
@click="openCreateUserApiKeyDialog"
>
{{ creatingApiKey ? '创建中...' : '创建' }}
</Button>
</template>
</Dialog>
<Dialog
v-model="showUserApiKeyFormDialog"
size="lg"
>
<template #header>
<div class="border-b border-border px-6 py-4">
<div class="flex items-center gap-3">
<div class="flex h-9 w-9 items-center justify-center rounded-lg bg-kraft/10 flex-shrink-0">
<Key class="h-5 w-5 text-kraft" />
</div>
<div class="flex-1 min-w-0">
<h3 class="text-lg font-semibold text-foreground leading-tight">
{{ editingUserApiKey ? '编辑 API Key' : '创建 API Key' }}
</h3>
<p class="text-xs text-muted-foreground">
{{ editingUserApiKey ? '更新用户 API Key 的名称和速率限制' : '为用户创建新的 API Key' }}
</p>
</div>
</div>
</div>
</template>
<div class="space-y-4">
<div class="space-y-2">
<Label
for="admin-user-key-name"
class="text-sm font-medium"
>密钥名称</Label>
<Input
id="admin-user-key-name"
v-model="userApiKeyForm.name"
class="h-10"
placeholder="例如:生产环境 Key"
/>
</div>
<div class="space-y-2">
<Label
for="admin-user-key-rate-limit"
class="text-sm font-medium"
>速率限制 (请求/分钟)</Label>
<Input
id="admin-user-key-rate-limit"
:model-value="userApiKeyForm.rate_limit ?? ''"
type="number"
min="0"
max="10000"
class="h-10"
placeholder="留空不限"
@update:model-value="(v) => userApiKeyForm.rate_limit = parseNumberInput(v, { min: 0, max: 10000 })"
/>
<p class="text-xs text-muted-foreground">
留空表示不限制
</p>
</div>
</div>
<template #footer>
<Button
variant="outline"
class="h-10 px-5"
@click="closeUserApiKeyFormDialog"
>
取消
</Button>
<Button
class="h-10 px-5"
:disabled="creatingApiKey"
@click="submitUserApiKeyForm"
>
{{ creatingApiKey ? (editingUserApiKey ? '保存中...' : '创建中...') : (editingUserApiKey ? '保存' : '创建') }}
</Button>
</template>
</Dialog>
<WalletOpsDrawer
:open="showWalletActionDialogState"
:wallet="walletActionTarget?.wallet || null"
@@ -849,6 +960,8 @@ import {
import UserFormDialog, { type UserFormData } from '@/features/users/components/UserFormDialog.vue'
import WalletOpsDrawer from '@/features/wallet/components/WalletOpsDrawer.vue'
import { parseApiError } from '@/utils/errorParser'
import { formatTokens, formatRateLimitInheritable, formatRateLimitSimple, isRateLimitInherited, isRateLimitUnlimited } from '@/utils/format'
import { parseNumberInput } from '@/utils/form'
import { log } from '@/utils/logger'
const { success, error } = useToast()
@@ -864,11 +977,17 @@ const userFormDialogRef = ref<InstanceType<typeof UserFormDialog>>()
// API Keys 对话框状态
const showApiKeysDialog = ref(false)
const showNewApiKeyDialog = ref(false)
const showUserApiKeyFormDialog = ref(false)
const selectedUser = ref<User | null>(null)
const userApiKeys = ref<ApiKey[]>([])
const newApiKey = ref('')
const creatingApiKey = ref(false)
const apiKeyInput = ref<HTMLInputElement>()
const editingUserApiKey = ref<ApiKey | null>(null)
const userApiKeyForm = ref({
name: '',
rate_limit: undefined as number | undefined,
})
// 用户统计
const userStats = ref<Record<string, UsageByUser>>({})
@@ -979,15 +1098,6 @@ async function loadUserWallets() {
}
}
function formatTokens(tokens: number): string {
if (tokens >= 1000000) {
return `${(tokens / 1000000).toFixed(1)}M`
} else if (tokens >= 1000) {
return `${(tokens / 1000).toFixed(1)}K`
}
return tokens.toString()
}
function formatNumber(value?: number | null): string {
const numericValue = typeof value === 'number' && Number.isFinite(value) ? value : 0
return numericValue.toLocaleString()
@@ -1071,7 +1181,8 @@ function editUser(user: User) {
is_active: user.is_active,
allowed_providers: user.allowed_providers == null ? null : [...user.allowed_providers],
allowed_api_formats: user.allowed_api_formats == null ? null : [...user.allowed_api_formats],
allowed_models: user.allowed_models == null ? null : [...user.allowed_models]
allowed_models: user.allowed_models == null ? null : [...user.allowed_models],
rate_limit: user.rate_limit ?? null
}
showUserFormDialog.value = true
}
@@ -1093,7 +1204,8 @@ async function handleUserFormSubmit(data: UserFormData & { password?: string; un
role: data.role,
allowed_providers: data.allowed_providers,
allowed_api_formats: data.allowed_api_formats,
allowed_models: data.allowed_models
allowed_models: data.allowed_models,
rate_limit: data.rate_limit ?? null
}
if (data.password) {
updateData.password = data.password
@@ -1112,7 +1224,8 @@ async function handleUserFormSubmit(data: UserFormData & { password?: string; un
role: data.role,
allowed_providers: data.allowed_providers,
allowed_api_formats: data.allowed_api_formats,
allowed_models: data.allowed_models
allowed_models: data.allowed_models,
rate_limit: data.rate_limit ?? null
})
// 如果创建时指定为禁用,则更新状态
if (data.is_active === false && newUser) {
@@ -1145,20 +1258,61 @@ async function loadUserApiKeys(userId: string) {
}
}
async function createApiKey() {
function openCreateUserApiKeyDialog() {
userApiKeyForm.value = {
name: `Key-${new Date().toISOString().split('T')[0]}`,
rate_limit: undefined,
}
editingUserApiKey.value = null
showUserApiKeyFormDialog.value = true
}
function openEditUserApiKeyDialog(apiKey: ApiKey) {
editingUserApiKey.value = apiKey
userApiKeyForm.value = {
name: apiKey.name || '',
rate_limit: apiKey.rate_limit ?? undefined,
}
showUserApiKeyFormDialog.value = true
}
function closeUserApiKeyFormDialog() {
showUserApiKeyFormDialog.value = false
editingUserApiKey.value = null
userApiKeyForm.value = {
name: '',
rate_limit: undefined,
}
}
async function submitUserApiKeyForm() {
if (!selectedUser.value) return
if (!userApiKeyForm.value.name.trim()) {
error('请输入密钥名称', editingUserApiKey.value ? '更新 API Key 失败' : '创建 API Key 失败')
return
}
creatingApiKey.value = true
try {
const response = await usersStore.createApiKey(
selectedUser.value.id,
`Key-${new Date().toISOString().split('T')[0]}`
)
newApiKey.value = response.key || ''
showNewApiKeyDialog.value = true
if (editingUserApiKey.value) {
await usersStore.updateApiKey(selectedUser.value.id, editingUserApiKey.value.id, {
name: userApiKeyForm.value.name,
rate_limit: userApiKeyForm.value.rate_limit ?? 0,
})
success('API Key已更新')
} else {
const response = await usersStore.createApiKey(selectedUser.value.id, {
name: userApiKeyForm.value.name,
rate_limit: userApiKeyForm.value.rate_limit ?? 0,
})
newApiKey.value = response.key || ''
showNewApiKeyDialog.value = true
success('API Key创建成功')
}
await loadUserApiKeys(selectedUser.value.id)
closeUserApiKeyFormDialog()
} catch (err: unknown) {
error(parseApiError(err, '未知错误'), '创建 API Key 失败')
error(parseApiError(err, '未知错误'), editingUserApiKey.value ? '更新 API Key 失败' : '创建 API Key 失败')
} finally {
creatingApiKey.value = false
}
@@ -39,7 +39,7 @@
for="rate-limit"
class="block text-sm font-medium"
>
每分钟请求限制
默认速率限制 (请求/分钟)
</Label>
<Input
id="rate-limit"
@@ -50,7 +50,7 @@
@update:model-value="$emit('update:rateLimitPerMinute', Number($event))"
/>
<p class="mt-1 text-xs text-muted-foreground">
0 表示不限制
0 表示默认不限制未单独配置的用户和独立 Key 会跟随这里
</p>
</div>
+110 -19
View File
@@ -20,7 +20,7 @@
size="icon"
class="h-8 w-8"
title="创建新 API Key"
@click="showCreateDialog = true"
@click="openCreateApiKeyDialog"
>
<Plus class="w-3.5 h-3.5" />
</Button>
@@ -56,7 +56,7 @@
<Button
size="lg"
class="shadow-lg shadow-primary/20"
@click="showCreateDialog = true"
@click="openCreateApiKeyDialog"
>
<Plus class="mr-2 h-4 w-4" />
创建新 API Key
@@ -157,17 +157,23 @@
<div class="flex flex-col items-center gap-1">
<Badge
:variant="apiKey.is_active ? 'success' : 'secondary'"
class="font-medium px-3 py-1"
class="h-5 px-2 py-0 text-[10px] font-medium"
>
{{ apiKey.is_active ? '活跃' : '禁用' }}
</Badge>
<Badge
v-if="apiKey.is_locked"
variant="warning"
class="font-medium text-[10px]"
class="h-5 px-2 py-0 text-[10px] font-medium"
>
已锁定
</Badge>
<Badge
variant="secondary"
class="h-5 px-2 py-0 text-[10px] font-medium"
>
{{ formatRateLimitSimple(apiKey.rate_limit) }}
</Badge>
</div>
</TableCell>
@@ -179,6 +185,16 @@
<!-- 操作按钮 -->
<TableCell class="py-4">
<div class="flex justify-center gap-1">
<Button
variant="ghost"
size="icon"
class="h-8 w-8"
:title="apiKey.is_locked ? '已锁定' : '编辑'"
:disabled="apiKey.is_locked"
@click="openEditApiKeyDialog(apiKey)"
>
<SquarePen class="h-4 w-4" />
</Button>
<Button
variant="ghost"
size="icon"
@@ -237,8 +253,24 @@
>
已锁定
</Badge>
<Badge
variant="secondary"
class="text-[10px] px-1.5 py-0"
>
{{ formatRateLimitSimple(apiKey.rate_limit) }}
</Badge>
</div>
<div class="flex items-center gap-0.5 flex-shrink-0">
<Button
variant="ghost"
size="icon"
class="h-7 w-7"
:title="apiKey.is_locked ? '已锁定' : '编辑'"
:disabled="apiKey.is_locked"
@click="openEditApiKeyDialog(apiKey)"
>
<SquarePen class="h-3.5 w-3.5" />
</Button>
<Button
variant="ghost"
size="icon"
@@ -288,6 +320,10 @@
<span class="text-foreground font-medium">
{{ formatNumber(apiKey.total_requests || 0) }}
</span>
<span class="text-muted-foreground"></span>
<span class="text-muted-foreground">
{{ formatRateLimitSimple(apiKey.rate_limit) }}
</span>
</div>
</div>
</div>
@@ -316,10 +352,10 @@
</div>
<div class="flex-1 min-w-0">
<h3 class="text-lg font-semibold text-foreground leading-tight">
创建 API 密钥
{{ editingApiKey ? '编辑 API 密钥' : '创建 API 密钥' }}
</h3>
<p class="text-xs text-muted-foreground">
创建一个新的密钥用于访问 API 服务
{{ editingApiKey ? '更新密钥名称和速率限制' : '创建一个新的密钥用于访问 API 服务' }}
</p>
</div>
</div>
@@ -344,26 +380,46 @@
给密钥起一个有意义的名称方便识别
</p>
</div>
<div class="space-y-2">
<Label
for="key-rate-limit"
class="text-sm font-semibold"
>速率限制 (请求/分钟)</Label>
<Input
id="key-rate-limit"
:model-value="newKeyRateLimit ?? ''"
type="number"
min="0"
max="10000"
placeholder="留空不限"
class="h-11 border-border/60"
@update:model-value="(v) => newKeyRateLimit = parseNumberInput(v, { min: 0, max: 10000 })"
/>
<p class="text-xs text-muted-foreground">
留空不限
</p>
</div>
</div>
<template #footer>
<Button
variant="outline"
class="h-11 px-6"
@click="showCreateDialog = false"
@click="closeApiKeyDialog"
>
取消
</Button>
<Button
class="h-11 px-6 shadow-lg shadow-primary/20"
:disabled="creating"
@click="createApiKey"
@click="saveApiKey"
>
<Loader2
v-if="creating"
class="animate-spin h-4 w-4 mr-2"
/>
{{ creating ? '创建中...' : '创建' }}
{{ creating ? (editingApiKey ? '保存中...' : '创建中...') : (editingApiKey ? '保存' : '创建') }}
</Button>
</template>
</Dialog>
@@ -455,10 +511,12 @@ import {
TableRow
} from '@/components/ui'
import RefreshButton from '@/components/ui/refresh-button.vue'
import { Plus, Key, Copy, Trash2, Loader2, Activity, CheckCircle, Power } from 'lucide-vue-next'
import { Plus, Key, Copy, Trash2, Loader2, Activity, CheckCircle, Power, SquarePen } from 'lucide-vue-next'
import { useToast } from '@/composables/useToast'
import { log } from '@/utils/logger'
import { parseApiError } from '@/utils/errorParser'
import { formatRateLimitSimple } from '@/utils/format'
import { parseNumberInput } from '@/utils/form'
import { getErrorStatus } from '@/types/api-error'
import { computed } from 'vue'
@@ -483,8 +541,10 @@ const showKeyDialog = ref(false)
const showDeleteDialog = ref(false)
const newKeyName = ref('')
const newKeyRateLimit = ref<number | undefined>(undefined)
const newKeyValue = ref('')
const keyToDelete = ref<ApiKey | null>(null)
const editingApiKey = ref<ApiKey | null>(null)
onMounted(() => {
loadApiKeys()
@@ -509,7 +569,28 @@ async function loadApiKeys() {
}
}
async function createApiKey() {
function openEditApiKeyDialog(apiKey: ApiKey) {
editingApiKey.value = apiKey
newKeyName.value = apiKey.name || ''
newKeyRateLimit.value = apiKey.rate_limit ?? undefined
showCreateDialog.value = true
}
function openCreateApiKeyDialog() {
editingApiKey.value = null
newKeyName.value = ''
newKeyRateLimit.value = undefined
showCreateDialog.value = true
}
function closeApiKeyDialog() {
showCreateDialog.value = false
editingApiKey.value = null
newKeyName.value = ''
newKeyRateLimit.value = undefined
}
async function saveApiKey() {
if (!newKeyName.value.trim()) {
showError('请输入密钥名称')
return
@@ -517,16 +598,26 @@ async function createApiKey() {
creating.value = true
try {
const newKey = await meApi.createApiKey(newKeyName.value)
newKeyValue.value = newKey.key || ''
showCreateDialog.value = false
showKeyDialog.value = true
newKeyName.value = ''
if (editingApiKey.value) {
await meApi.updateApiKey(editingApiKey.value.id, {
name: newKeyName.value,
rate_limit: newKeyRateLimit.value ?? 0,
})
success('API 密钥更新成功')
} else {
const newKey = await meApi.createApiKey({
name: newKeyName.value,
rate_limit: newKeyRateLimit.value ?? 0,
})
newKeyValue.value = newKey.key || ''
showKeyDialog.value = true
success('API 密钥创建成功')
}
closeApiKeyDialog()
await loadApiKeys()
success('API 密钥创建成功')
} catch (error) {
log.error('创建 API 密钥失败:', error)
showError('创建 API 密钥失败')
log.error(editingApiKey.value ? '更新 API 密钥失败:' : '创建 API 密钥失败:', error)
showError(editingApiKey.value ? '更新 API 密钥失败' : '创建 API 密钥失败')
} finally {
creating.value = false
}
+35 -5
View File
@@ -22,7 +22,7 @@ from src.core.exceptions import InvalidRequestException, NotFoundException
from src.core.logger import logger
from src.database import get_db, get_db_context
from src.models.api import CreateApiKeyRequest
from src.models.database import ApiKey, Wallet
from src.models.database import ApiKey, Usage, Wallet
from src.services.user.apikey import ApiKeyService
from src.services.user.bulk_cleanup import pre_clean_api_key
from src.services.wallet import WalletService
@@ -70,7 +70,9 @@ router = APIRouter(prefix="/api/admin/api-keys", tags=["Admin - API Keys (Standa
pipeline = get_pipeline()
def _serialize_standalone_key_item(api_key: ApiKey) -> dict[str, Any]:
def _serialize_standalone_key_item(
api_key: ApiKey, *, total_tokens: int | None = None
) -> dict[str, Any]:
return {
"id": api_key.id,
"user_id": api_key.user_id,
@@ -79,6 +81,7 @@ def _serialize_standalone_key_item(api_key: ApiKey) -> dict[str, Any]:
"is_active": api_key.is_active,
"is_standalone": api_key.is_standalone,
"total_requests": api_key.total_requests,
"total_tokens": int(total_tokens or 0),
"total_cost_usd": float(api_key.total_cost_usd or 0),
"rate_limit": api_key.rate_limit,
"allowed_providers": api_key.allowed_providers,
@@ -116,8 +119,27 @@ def _list_standalone_api_keys_sync(
for api_key in api_keys:
db.refresh(api_key)
token_map: dict[str, int] = {}
if api_keys:
stats_rows = (
db.query(
Usage.api_key_id,
func.sum(Usage.total_tokens).label("total_tokens"),
)
.filter(Usage.api_key_id.in_([api_key.id for api_key in api_keys]))
.group_by(Usage.api_key_id)
.all()
)
token_map = {row.api_key_id: int(row.total_tokens or 0) for row in stats_rows}
return {
"api_keys": [_serialize_standalone_key_item(api_key) for api_key in api_keys],
"api_keys": [
_serialize_standalone_key_item(
api_key,
total_tokens=token_map.get(api_key.id, 0),
)
for api_key in api_keys
],
"total": total,
"limit": limit,
"skip": skip,
@@ -385,7 +407,7 @@ async def create_standalone_api_key(
- `allowed_providers`: 可选,允许使用的提供商列表
- `allowed_api_formats`: 可选,允许使用的 API 格式列表
- `allowed_models`: 可选,允许使用的模型列表
- `rate_limit`: 可选,速率限制配置(请求数/秒
- `rate_limit`: 可选,每分钟请求限制(null 表示跟随系统默认,0 表示不限制
- `expire_days`: 可选,过期天数(与 expires_at 二选一)
- `expires_at`: 可选,过期时间(ISO 格式或 YYYY-MM-DD 格式,优先级高于 expire_days
- `auto_delete_on_expiry`: 可选,过期后是否自动删除
@@ -421,7 +443,7 @@ async def update_api_key(
**请求体字段**:
- `name`: 可选,API Key 的名称
- `unlimited_balance`: 可选,是否无限余额(true=无限,false=有限,不修改余额数值)
- `rate_limit`: 可选,速率限制配置null 表示限制)
- `rate_limit`: 可选,每分钟请求限制null 表示跟随系统默认,0 表示不限制)
- `allowed_providers`: 可选,允许使用的提供商列表
- `allowed_api_formats`: 可选,允许使用的 API 格式列表
- `allowed_models`: 可选,允许使用的模型列表
@@ -683,6 +705,14 @@ class AdminGetKeyDetailAdapter(AdminApiAdapter):
"is_active": api_key.is_active,
"is_standalone": api_key.is_standalone,
"total_requests": api_key.total_requests,
"total_tokens": int(
(
db.query(func.sum(Usage.total_tokens))
.filter(Usage.api_key_id == api_key.id)
.scalar()
)
or 0
),
"total_cost_usd": float(api_key.total_cost_usd or 0),
"rate_limit": api_key.rate_limit,
"allowed_providers": api_key.allowed_providers,
+46 -2
View File
@@ -2252,6 +2252,7 @@ class AdminExportUsersAdapter(AdminApiAdapter):
"allowed_providers": user.allowed_providers,
"allowed_api_formats": user.allowed_api_formats,
"allowed_models": user.allowed_models,
"rate_limit": user.rate_limit,
"model_capability_settings": user.model_capability_settings,
"unlimited": wallet_service.is_unlimited_wallet(wallet),
"wallet": (wallet_service.serialize_wallet_summary(wallet) if wallet else None),
@@ -2265,7 +2266,7 @@ class AdminExportUsersAdapter(AdminApiAdapter):
standalone_keys_data = [self._serialize_api_key(key, db=db) for key in standalone_keys]
return {
"version": "1.2",
"version": "1.3",
"exported_at": datetime.now(timezone.utc).isoformat(),
"users": users_data,
"standalone_keys": standalone_keys_data,
@@ -2273,6 +2274,17 @@ class AdminExportUsersAdapter(AdminApiAdapter):
class AdminImportUsersAdapter(AdminApiAdapter):
@staticmethod
def _is_legacy_users_export(version: object) -> bool:
if version is None:
return True
normalized = str(version).strip()
try:
parts = normalized.split(".")
return (int(parts[0]), int(parts[1])) < (1, 3)
except Exception:
return True
@staticmethod
def _resolve_api_key_material(key_data: dict[str, Any]) -> tuple[str | None, str | None]:
"""解析用户 API Key 导入材料,优先使用明文 key。"""
@@ -2289,6 +2301,30 @@ class AdminImportUsersAdapter(AdminApiAdapter):
key_encrypted = key_data.get("key_encrypted")
return key_hash, key_encrypted
@staticmethod
def _normalize_imported_user_rate_limit(user_data: dict[str, Any]) -> int | None:
if "rate_limit" not in user_data:
return None
value = user_data.get("rate_limit")
return int(value) if value is not None else None
@staticmethod
def _normalize_imported_api_key_rate_limit(
key_data: dict[str, Any],
*,
is_standalone: bool,
legacy_export: bool,
) -> int | None:
if "rate_limit" not in key_data:
return None if is_standalone and not legacy_export else 0
value = key_data.get("rate_limit")
if value is None:
if is_standalone and not legacy_export:
return None
return 0
return int(value)
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
"""导入用户数据"""
import uuid
@@ -2306,6 +2342,7 @@ class AdminImportUsersAdapter(AdminApiAdapter):
# 获取导入选项
merge_mode = payload.get("merge_mode", "skip") # skip, overwrite, error
legacy_export = self._is_legacy_users_export(payload.get("version"))
users_data = payload.get("users", [])
standalone_keys_data = payload.get("standalone_keys", [])
@@ -2358,7 +2395,11 @@ class AdminImportUsersAdapter(AdminApiAdapter):
allowed_providers=key_data.get("allowed_providers"),
allowed_api_formats=key_data.get("allowed_api_formats"),
allowed_models=key_data.get("allowed_models"),
rate_limit=key_data.get("rate_limit"),
rate_limit=self._normalize_imported_api_key_rate_limit(
key_data,
is_standalone=is_standalone or key_data.get("is_standalone", False),
legacy_export=legacy_export,
),
concurrent_limit=key_data.get("concurrent_limit", 5),
force_capabilities=key_data.get("force_capabilities"),
is_active=key_data.get("is_active", True),
@@ -2396,6 +2437,7 @@ class AdminImportUsersAdapter(AdminApiAdapter):
and wallet_payload.get("limit_mode") in {"finite", "unlimited"}
else ("unlimited" if user_data.get("unlimited") else "finite")
)
imported_user_rate_limit = self._normalize_imported_user_rate_limit(user_data)
if existing_user:
user_id = existing_user.id
@@ -2413,6 +2455,7 @@ class AdminImportUsersAdapter(AdminApiAdapter):
existing_user.allowed_providers = user_data.get("allowed_providers")
existing_user.allowed_api_formats = user_data.get("allowed_api_formats")
existing_user.allowed_models = user_data.get("allowed_models")
existing_user.rate_limit = imported_user_rate_limit
existing_user.model_capability_settings = user_data.get(
"model_capability_settings"
)
@@ -2449,6 +2492,7 @@ class AdminImportUsersAdapter(AdminApiAdapter):
allowed_providers=user_data.get("allowed_providers"),
allowed_api_formats=user_data.get("allowed_api_formats"),
allowed_models=user_data.get("allowed_models"),
rate_limit=imported_user_rate_limit,
model_capability_settings=user_data.get("model_capability_settings"),
is_active=user_data.get("is_active", True),
)
+105 -1
View File
@@ -18,7 +18,7 @@ from src.core.exceptions import InvalidRequestException, NotFoundException, tran
from src.core.logger import logger
from src.database import get_db, get_db_context
from src.models.admin_requests import UpdateUserRequest
from src.models.api import CreateApiKeyRequest, CreateUserRequest
from src.models.api import CreateApiKeyRequest, CreateUserRequest, UpdateMyApiKeyRequest
from src.models.database import ApiKey, User, UserRole, Wallet
from src.services.cache.user_cache import UserCacheService
from src.services.system.config import SystemConfigService
@@ -57,6 +57,7 @@ def _serialize_user(
"allowed_providers": user.allowed_providers,
"allowed_api_formats": user.allowed_api_formats,
"allowed_models": user.allowed_models,
"rate_limit": user.rate_limit,
"unlimited": WalletService.is_unlimited_wallet(resolved_wallet),
"is_active": user.is_active,
"created_at": user.created_at.isoformat(),
@@ -89,6 +90,7 @@ def _create_user_sync(
allowed_providers=request.allowed_providers,
allowed_api_formats=request.allowed_api_formats,
allowed_models=request.allowed_models,
rate_limit=request.rate_limit,
)
return _serialize_user(db, user), {
"action": "create_user",
@@ -255,6 +257,61 @@ def _delete_user_key_sync(user_id: str, key_id: str) -> tuple[dict[str, Any], di
}
def _update_user_key_sync(
user_id: str,
key_id: str,
request: UpdateMyApiKeyRequest,
) -> tuple[dict[str, Any], dict[str, Any]]:
with get_db_context() as db:
api_key = (
db.query(ApiKey)
.filter(
ApiKey.id == key_id,
ApiKey.user_id == user_id,
ApiKey.is_standalone == False,
)
.first()
)
if not api_key:
raise NotFoundException("API Key不存在或不属于该用户", "api_key")
update_data = request.model_dump(exclude_unset=True)
if "rate_limit" in update_data and update_data["rate_limit"] is None:
update_data["rate_limit"] = 0
updated_key = ApiKeyService.update_api_key(db, key_id, **update_data)
if not updated_key:
raise NotFoundException("API Key不存在或不属于该用户", "api_key")
return (
{
"id": updated_key.id,
"name": updated_key.name,
"key_display": updated_key.get_display_key(),
"is_active": updated_key.is_active,
"is_locked": updated_key.is_locked,
"total_requests": updated_key.total_requests,
"total_cost_usd": float(updated_key.total_cost_usd or 0),
"rate_limit": updated_key.rate_limit,
"expires_at": (
updated_key.expires_at.isoformat() if updated_key.expires_at else None
),
"last_used_at": (
updated_key.last_used_at.isoformat() if updated_key.last_used_at else None
),
"created_at": updated_key.created_at.isoformat(),
"message": "API Key更新成功",
},
{
"action": "update_user_api_key",
"target_user_id": user_id,
"key_id": key_id,
"updated_fields": list(update_data.keys()),
},
)
def _toggle_user_key_lock_sync(
user_id: str,
key_id: str,
@@ -452,6 +509,26 @@ async def delete_user_api_key(
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
@router.put("/{user_id}/api-keys/{key_id}")
async def update_user_api_key(
user_id: str,
key_id: str,
request: Request,
db: Session = Depends(get_db),
) -> Any:
"""
更新用户的 API 密钥
更新指定用户的普通 API 密钥基础配置。
**路径参数**:
- `user_id`: 用户 ID (UUID)
- `key_id`: 密钥 ID
"""
adapter = AdminUpdateUserKeyAdapter(user_id=user_id, key_id=key_id)
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
@router.patch("/{user_id}/api-keys/{key_id}/lock")
async def toggle_user_api_key_lock(
user_id: str,
@@ -700,6 +777,33 @@ class AdminDeleteUserKeyAdapter(AdminApiAdapter):
return response
class AdminUpdateUserKeyAdapter(AdminApiAdapter):
"""更新用户的普通 API Key"""
def __init__(self, user_id: str, key_id: str):
self.user_id = user_id
self.key_id = key_id
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
payload = context.ensure_json_body()
try:
request = UpdateMyApiKeyRequest.model_validate(payload)
except ValidationError as e:
errors = e.errors()
if errors:
raise InvalidRequestException(translate_pydantic_error(errors[0]))
raise InvalidRequestException("请求数据验证失败")
response, audit_meta = await run_in_threadpool(
_update_user_key_sync,
self.user_id,
self.key_id,
request,
)
context.add_audit_metadata(**audit_meta)
return response
class AdminToggleUserKeyLockAdapter(AdminApiAdapter):
"""切换用户普通 API Key 的锁定状态"""
+54
View File
@@ -19,7 +19,9 @@ from src.core.logger import logger
from src.database.database import create_session
from src.models.database import ApiKey, AuditEventType, User
from src.services.auth.service import AuthService
from src.services.rate_limit.user_rpm_limiter import SYSTEM_RPM_CONFIG_KEY, get_user_rpm_limiter
from src.services.system.audit import AuditService
from src.services.system.config import SystemConfigService
from src.services.usage.service import UsageService
from src.services.wallet import WalletService
from src.utils.perf import PerfRecorder
@@ -123,6 +125,9 @@ class ApiRequestPipeline:
auth_duration = PerfRecorder.stop(auth_start, "pipeline_auth", labels=perf_labels)
_record_perf_metric("auth_ms", auth_duration)
if mode in {ApiMode.STANDARD, ApiMode.PROXY} and api_key and user:
await self._check_user_rate_limit(http_request, db, user, api_key)
raw_body = None
should_eager_read_body = http_request.method in {"POST", "PUT", "PATCH"} and getattr(
adapter, "eager_request_body", True
@@ -267,6 +272,55 @@ class ApiRequestPipeline:
# Internal helpers
# --------------------------------------------------------------------- #
async def _check_user_rate_limit(
self,
request: Request,
db: Session,
user: User,
api_key: ApiKey,
) -> None:
limiter = await get_user_rpm_limiter()
system_default_raw = SystemConfigService.get_config(db, SYSTEM_RPM_CONFIG_KEY, default=0)
system_default = max(int(system_default_raw or 0), 0)
if api_key.is_standalone:
effective_user_limit = (
max(int(api_key.rate_limit or 0), 0)
if api_key.rate_limit is not None
else system_default
)
user_rpm_key = limiter.get_standalone_rpm_key(api_key.id)
key_rpm_limit = 0
else:
effective_user_limit = (
max(int(user.rate_limit or 0), 0) if user.rate_limit is not None else system_default
)
user_rpm_key = limiter.get_user_rpm_key(user.id)
key_rpm_limit = max(int(api_key.rate_limit or 0), 0)
result = await limiter.check_and_consume(
user_rpm_key=user_rpm_key,
user_rpm_limit=effective_user_limit,
key_rpm_key=limiter.get_key_rpm_key(api_key.id),
key_rpm_limit=key_rpm_limit,
)
if result.allowed:
return
scope = result.scope or "user"
limit = result.limit or (effective_user_limit if scope == "user" else key_rpm_limit)
retry_after = result.retry_after or limiter.get_retry_after()
headers = {
"Retry-After": str(retry_after),
"X-RateLimit-Limit": str(limit),
"X-RateLimit-Remaining": "0",
"X-RateLimit-Scope": scope,
}
request.state.rate_limit_scope = scope
raise HTTPException(status_code=429, detail="请求过于频繁,请稍后重试", headers=headers)
async def _authenticate_client(
self, request: Request, db: Session, adapter: ApiAdapter, **_kw: object
) -> tuple[User, ApiKey]:
+73 -23
View File
@@ -16,7 +16,8 @@ from src.api.base.pipeline import get_pipeline
from src.core.logger import logger
from src.database import get_db
from src.models.database import ApiKey, AuditLog
from src.plugins.manager import get_plugin_manager
from src.services.rate_limit.user_rpm_limiter import SYSTEM_RPM_CONFIG_KEY, get_user_rpm_limiter
from src.services.system.config import SystemConfigService
router = APIRouter(prefix="/api/monitoring", tags=["Monitoring"])
pipeline = get_pipeline()
@@ -146,10 +147,6 @@ class UserRateLimitStatusAdapter(AuthenticatedApiAdapter):
if not user:
raise HTTPException(status_code=401, detail="未登录")
rate_limiter = _get_rate_limit_plugin()
if not rate_limiter or not hasattr(rate_limiter, "get_rate_limit_headers"):
raise HTTPException(status_code=503, detail="速率限制插件未启用或不支持状态查询")
api_keys = (
db.query(ApiKey)
.filter(ApiKey.user_id == user.id, ApiKey.is_active.is_(True))
@@ -157,31 +154,84 @@ class UserRateLimitStatusAdapter(AuthenticatedApiAdapter):
.all()
)
try:
limiter = await get_user_rpm_limiter()
system_default_raw = SystemConfigService.get_config(
db, SYSTEM_RPM_CONFIG_KEY, default=0
)
system_default = max(int(system_default_raw or 0), 0)
reset_at = limiter.get_reset_at()
window = f"{limiter.bucket_seconds}s"
except Exception as exc:
logger.warning("读取新 RPM 限流状态失败,回退插件状态接口: {}", exc)
limiter = None
system_default = 0
reset_at = None
window = None
rate_limit_info = []
for key in api_keys:
try:
headers = rate_limiter.get_rate_limit_headers(key)
except Exception as exc:
logger.warning(f"无法获取Key {key.id} 的限流信息: {exc}")
headers = {}
if limiter is not None:
if key.is_standalone:
user_limit = key.rate_limit if key.rate_limit is not None else system_default
user_scope_key = limiter.get_standalone_rpm_key(key.id)
key_limit = 0
else:
user_limit = user.rate_limit if user.rate_limit is not None else system_default
user_scope_key = limiter.get_user_rpm_key(user.id)
key_limit = max(int(key.rate_limit or 0), 0)
user_count = (
await limiter.get_scope_count(user_scope_key)
if user_limit and user_limit > 0
else 0
)
key_count = (
await limiter.get_scope_count(limiter.get_key_rpm_key(key.id))
if key_limit > 0
else 0
)
user_remaining = max(user_limit - user_count, 0) if user_limit > 0 else None
key_remaining = max(key_limit - key_count, 0) if key_limit > 0 else None
scoped_statuses: list[tuple[str, int, int]] = []
if user_limit > 0 and user_remaining is not None:
scoped_statuses.append(("user", user_limit, user_remaining))
if key_limit > 0 and key_remaining is not None:
scoped_statuses.append(("key", key_limit, key_remaining))
primary_scope = (
min(scoped_statuses, key=lambda item: item[2]) if scoped_statuses else None
)
rate_limit_info.append(
{
"api_key_name": key.name or f"Key-{key.id}",
"limit": primary_scope[1] if primary_scope else None,
"remaining": primary_scope[2] if primary_scope else None,
"scope": primary_scope[0] if primary_scope else None,
"reset_time": reset_at.isoformat() if reset_at else None,
"window": window,
"user_limit": user_limit,
"user_remaining": user_remaining,
"key_limit": key_limit if key_limit > 0 else None,
"key_remaining": key_remaining,
}
)
continue
rate_limit_info.append(
{
"api_key_name": key.name or f"Key-{key.id}",
"limit": headers.get("X-RateLimit-Limit"),
"remaining": headers.get("X-RateLimit-Remaining"),
"reset_time": headers.get("X-RateLimit-Reset"),
"window": headers.get("X-RateLimit-Window"),
"limit": None,
"remaining": None,
"scope": None,
"reset_time": None,
"window": None,
"user_limit": None,
"user_remaining": None,
"key_limit": None,
"key_remaining": None,
}
)
return {"user_id": user.id, "api_keys": rate_limit_info}
def _get_rate_limit_plugin() -> Any:
try:
plugin_manager = get_plugin_manager()
return plugin_manager.get_plugin("rate_limit")
except Exception as exc:
logger.warning(f"获取速率限制插件失败: {exc}")
return None
+83
View File
@@ -33,6 +33,7 @@ from src.models.api import (
PublicGlobalModelListResponse,
PublicGlobalModelResponse,
UpdateApiKeyProvidersRequest,
UpdateMyApiKeyRequest,
UpdatePreferencesRequest,
UpdateProfileRequest,
)
@@ -146,6 +147,7 @@ def _create_my_api_key_sync(user_id: str, request: CreateMyApiKeyRequest) -> dic
db=db,
user_id=user_id,
name=request.name,
rate_limit=request.rate_limit,
)
except ValueError as exc:
raise InvalidRequestException(str(exc)) from exc
@@ -154,6 +156,7 @@ def _create_my_api_key_sync(user_id: str, request: CreateMyApiKeyRequest) -> dic
"name": api_key.name,
"key": plain_key,
"key_display": api_key.get_display_key(),
"rate_limit": api_key.rate_limit,
"message": "API密钥创建成功",
}
@@ -189,6 +192,50 @@ def _toggle_my_api_key_sync(user_id: str, key_id: str) -> dict[str, Any]:
}
def _update_my_api_key_sync(
user_id: str,
key_id: str,
request: UpdateMyApiKeyRequest,
) -> dict[str, Any]:
with get_db_context() as db:
api_key = (
db.query(ApiKey)
.filter(
ApiKey.id == key_id,
ApiKey.user_id == user_id,
ApiKey.is_standalone == False,
)
.first()
)
if not api_key:
raise NotFoundException("API密钥不存在", "api_key")
if api_key.is_locked:
raise ForbiddenException("该密钥已被管理员锁定,无法修改")
update_data = request.model_dump(exclude_unset=True)
if "rate_limit" in update_data and update_data["rate_limit"] is None:
update_data["rate_limit"] = 0
updated = ApiKeyService.update_api_key(db, key_id, **update_data)
if not updated:
raise NotFoundException("API密钥不存在", "api_key")
return {
"id": updated.id,
"name": updated.name,
"key_display": updated.get_display_key(),
"is_active": updated.is_active,
"is_locked": updated.is_locked,
"allowed_providers": updated.allowed_providers,
"force_capabilities": updated.force_capabilities,
"rate_limit": updated.rate_limit,
"last_used_at": updated.last_used_at.isoformat() if updated.last_used_at else None,
"expires_at": updated.expires_at.isoformat() if updated.expires_at else None,
"created_at": updated.created_at.isoformat(),
"message": "API密钥已更新",
}
def _update_api_key_providers_sync(
user_id: str,
api_key_id: str,
@@ -484,6 +531,20 @@ async def delete_my_api_key(key_id: str, request: Request, db: Session = Depends
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
@router.put("/api-keys/{key_id}")
async def update_my_api_key(key_id: str, request: Request, db: Session = Depends(get_db)) -> Any:
"""
更新 API 密钥
更新指定 API 密钥的基础配置。
**路径参数**:
- `key_id`: 密钥 ID
"""
adapter = UpdateMyApiKeyAdapter(key_id=key_id)
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
@router.patch("/api-keys/{key_id}")
async def toggle_my_api_key(key_id: str, request: Request, db: Session = Depends(get_db)) -> Any:
"""
@@ -877,6 +938,7 @@ class ListMyApiKeysAdapter(AuthenticatedApiAdapter):
"created_at": key.created_at.isoformat(),
"total_requests": real_stats["total_requests"],
"total_cost_usd": real_stats["total_cost_usd"],
"rate_limit": key.rate_limit,
"allowed_providers": key.allowed_providers,
"force_capabilities": key.force_capabilities,
}
@@ -965,6 +1027,27 @@ class GetMyApiKeyDetailAdapter(AuthenticatedApiAdapter):
}
@dataclass
class UpdateMyApiKeyAdapter(AuthenticatedApiAdapter):
"""更新 API 密钥基础配置的适配器"""
key_id: str
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
payload = context.ensure_json_body()
try:
request = UpdateMyApiKeyRequest.model_validate(payload)
except ValidationError as e:
errors = e.errors()
if errors:
raise InvalidRequestException(translate_pydantic_error(errors[0]))
raise InvalidRequestException("请求数据验证失败")
return await run_in_threadpool(
_update_my_api_key_sync, context.user.id, self.key_id, request
)
@dataclass
class DeleteMyApiKeyAdapter(AuthenticatedApiAdapter):
"""删除 API 密钥的适配器"""
-2
View File
@@ -109,8 +109,6 @@ class Config:
# 支付回调安全配置(公开回调入口必须携带该共享密钥)
self.payment_callback_secret = os.getenv("PAYMENT_CALLBACK_SECRET", "").strip()
# LLM API 速率限制配置(每分钟请求数)
self.llm_api_rate_limit = int(os.getenv("LLM_API_RATE_LIMIT", "100"))
self.public_api_rate_limit = int(os.getenv("PUBLIC_API_RATE_LIMIT", "60"))
# 异常处理配置
+11
View File
@@ -45,6 +45,7 @@ if TYPE_CHECKING:
from src.services.model.fetch_scheduler import ModelFetchScheduler
from src.services.provider_keys.pool_quota_probe_scheduler import PoolQuotaProbeScheduler
from src.services.rate_limit.concurrency_manager import ConcurrencyManager
from src.services.rate_limit.user_rpm_limiter import UserRpmLimiter
from src.services.system.maintenance_scheduler import MaintenanceScheduler
from src.services.system.scheduler import TaskScheduler
from src.services.task.polling.task_poller import TaskPollerService
@@ -99,6 +100,7 @@ class LifecycleState:
redis_client: Redis | None = None
concurrency_manager: ConcurrencyManager | None = None
user_rpm_limiter: UserRpmLimiter | None = None
plugin_manager: PluginManager | None = None
available_modules: list[ModuleDefinition] = field(default_factory=list)
task_coordinator: StartupTaskCoordinator | None = None
@@ -181,6 +183,11 @@ async def _initialize_core_infrastructure(state: LifecycleState) -> None:
state.concurrency_manager = await get_concurrency_manager()
logger.info("初始化用户/API Key RPM 限流器...")
from src.services.rate_limit.user_rpm_limiter import get_user_rpm_limiter
state.user_rpm_limiter = await get_user_rpm_limiter()
# 初始化批量提交器(提升数据库并发能力)
logger.info("初始化批量提交器...")
from src.core.batch_committer import init_batch_committer
@@ -541,6 +548,10 @@ async def _run_shutdown(state: LifecycleState) -> None:
if state.concurrency_manager:
await state.concurrency_manager.close()
logger.info("关闭用户/API Key RPM 限流器...")
if state.user_rpm_limiter:
await state.user_rpm_limiter.close()
# 关闭全局Redis客户端
logger.info("关闭全局Redis客户端...")
from src.clients.redis_client import close_redis_client
-42
View File
@@ -8,7 +8,6 @@
from __future__ import annotations
import hashlib
import time
from typing import TYPE_CHECKING
@@ -50,7 +49,6 @@ class PluginMiddleware:
self._notification_cache_expires: float = 0.0
# 从配置读取速率限制值
self.llm_api_rate_limit = config.llm_api_rate_limit
self.public_api_rate_limit = config.public_api_rate_limit
# 完全跳过限流的路径(静态资源、文档等)
@@ -69,14 +67,6 @@ class PluginMiddleware:
"/api/monitoring/", # 监控端点
]
# LLM API 端点(需要特殊的速率限制策略)
self.llm_api_paths = [
"/v1/messages",
"/v1/chat/completions",
"/v1/responses",
"/v1/completions",
]
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
"""ASGI 入口点"""
if scope["type"] != "http":
@@ -291,13 +281,6 @@ class PluginMiddleware:
return "unknown"
def _is_llm_api_path(self, path: str) -> bool:
"""检查是否为 LLM API 端点"""
for llm_path in self.llm_api_paths:
if path.startswith(llm_path):
return True
return False
async def _get_rate_limit_key_and_config(
self, request: Request
) -> tuple[str | None, int | None]:
@@ -305,7 +288,6 @@ class PluginMiddleware:
获取速率限制的key和配置
策略说明:
- /v1/messages, /v1/chat/completions LLM API: API Key 限流
- /api/public/* 端点: 使用服务器级别 IP 限制
- /api/admin/* 端点: 跳过 skip_rate_limit_paths 中跳过
- /api/auth/* 端点: 跳过由路由层的 IPRateLimiter 处理
@@ -315,30 +297,6 @@ class PluginMiddleware:
"""
path = request.url.path
# LLM API 端点: 按 API Key 或 IP 限流
if self._is_llm_api_path(path):
# 尝试从请求头获取 API Key
auth_header = request.headers.get("authorization", "")
api_key = request.headers.get("x-api-key", "")
if auth_header.lower().startswith("bearer "):
api_key = auth_header[7:]
if api_key:
# 使用 API Key 的哈希作为限制 key(避免日志泄露完整 key)
key_hash = hashlib.sha256(api_key.encode()).hexdigest()[:16]
key = f"llm_api_key:{key_hash}"
request.state.rate_limit_key_type = "api_key"
else:
# 无 API Key 时使用 IP 限制(更严格)
client_ip = self._get_client_ip(request)
key = f"llm_ip:{client_ip}"
request.state.rate_limit_key_type = "ip"
rate_limit = self.llm_api_rate_limit
request.state.rate_limit_value = rate_limit
return key, rate_limit
# /api/public/* 端点: 使用服务器级别 IP 地址作为限制 key
if path.startswith("/api/public/"):
client_ip = self._get_client_ip(request)
+5
View File
@@ -692,6 +692,11 @@ class UpdateUserRequest(BaseModel):
allowed_providers: list[str] | None = Field(None, description="允许使用的提供商 ID 列表")
allowed_api_formats: list[str] | None = Field(None, description="允许使用的 API 格式列表")
allowed_models: list[str] | None = Field(None, description="允许使用的模型名称列表")
rate_limit: int | None = Field(
None,
ge=0,
description="每分钟请求限制;null 表示继承系统默认,0 表示不限制",
)
@field_validator("username")
@classmethod
+29 -2
View File
@@ -252,6 +252,11 @@ class CreateUserRequest(BaseModel):
allowed_models: list[str] | None = Field(
default=None, description="允许使用的模型名称列表,null表示无限制"
)
rate_limit: int | None = Field(
default=None,
ge=0,
description="每分钟请求限制;null 表示继承系统默认,0 表示不限制",
)
@field_validator("initial_gift_usd", mode="before")
@classmethod
@@ -335,6 +340,11 @@ class UpdateUserRequest(BaseModel):
allowed_providers: list[str] | None = None # 允许使用的提供商 ID 列表
allowed_api_formats: list[str] | None = None # 允许使用的 API 格式列表
allowed_models: list[str] | None = None # 允许使用的模型名称列表
rate_limit: int | None = Field(
default=None,
ge=0,
description="每分钟请求限制;null 表示继承系统默认,0 表示不限制",
)
is_active: bool | None = None
@field_validator("allowed_api_formats")
@@ -351,7 +361,11 @@ class CreateApiKeyRequest(BaseModel):
allowed_providers: list[str] | None = None # 允许使用的提供商 ID 列表
allowed_api_formats: list[str] | None = None # 允许使用的 API 格式列表
allowed_models: list[str] | None = None # 允许使用的模型名称列表
rate_limit: int | None = None # None = 无限制
rate_limit: int | None = Field(
None,
ge=0,
description="每分钟请求限制;独立Key: null=继承系统默认,0=不限制;普通Key: 0=不限制",
)
expire_days: int | None = None # None = 永不过期,数字 = 多少天后过期
expires_at: str | None = None # ISO 日期字符串,如 "2025-12-31",优先于 expire_days
initial_balance_usd: float | None = Field(
@@ -382,6 +396,7 @@ class UserResponse(BaseModel):
allowed_providers: list[str] | None = None # 允许使用的提供商 ID 列表
allowed_api_formats: list[str] | None = None # 允许使用的 API 格式列表
allowed_models: list[str] | None = None # 允许使用的模型名称列表
rate_limit: int | None = None
unlimited: bool = False
is_active: bool
created_at: datetime
@@ -402,7 +417,7 @@ class ApiKeyResponse(BaseModel):
total_cost_usd: float
allowed_providers: list[str] | None
allowed_models: list[str] | None
rate_limit: int
rate_limit: int | None
is_active: bool
expires_at: datetime | None = None
is_standalone: bool = False
@@ -772,6 +787,18 @@ class CreateMyApiKeyRequest(BaseModel):
"""创建我的API密钥请求"""
name: str
rate_limit: int = Field(0, ge=0, description="该 Key 的每分钟请求限制,0 表示不限制")
class UpdateMyApiKeyRequest(BaseModel):
"""更新我的 API 密钥请求"""
name: str | None = None
rate_limit: int | None = Field(
None,
ge=0,
description="该 Key 的每分钟请求限制;0 表示不限制,null 表示不修改",
)
class ProviderConfig(BaseModel):
+6 -1
View File
@@ -110,6 +110,9 @@ class User(Base):
allowed_providers = Column(JSON, nullable=True) # 允许使用的提供商 ID 列表
allowed_api_formats = Column(JSON, nullable=True) # 允许使用的 API 格式列表
allowed_models = Column(JSON, nullable=True) # 允许使用的模型名称列表
rate_limit = Column(
Integer, nullable=True, default=None
) # 每分钟请求限制,NULL=继承系统默认,0=不限制,N=N RPM
# Key 能力配置
model_capability_settings = Column(JSON, nullable=True) # 用户针对特定模型的能力配置
@@ -209,7 +212,9 @@ class ApiKey(Base):
allowed_providers = Column(JSON, nullable=True) # 允许使用的提供商 ID 列表
allowed_api_formats = Column(JSON, nullable=True) # 允许使用的 API 格式列表
allowed_models = Column(JSON, nullable=True) # 允许使用的模型名称列表
rate_limit = Column(Integer, default=None, nullable=True) # 每分钟请求限制,None = 无限制
rate_limit = Column(
Integer, default=None, nullable=True
) # 每分钟请求限制;独立Key: NULL=继承系统默认,普通Key: 0=不限制
concurrent_limit = Column(Integer, default=5, nullable=True) # 并发请求限制
# Key 能力配置
+3
View File
@@ -12,6 +12,7 @@ from src.services.rate_limit.adaptive_rpm import (
from src.services.rate_limit.concurrency_manager import ConcurrencyManager
from src.services.rate_limit.detector import RateLimitDetector
from src.services.rate_limit.ip_limiter import IPRateLimiter
from src.services.rate_limit.user_rpm_limiter import UserRpmLimiter, get_user_rpm_limiter
__all__ = [
"AdaptiveConcurrencyManager", # 向后兼容
@@ -19,5 +20,7 @@ __all__ = [
"ConcurrencyManager",
"IPRateLimiter",
"RateLimitDetector",
"UserRpmLimiter",
"get_adaptive_rpm_manager",
"get_user_rpm_limiter",
]
+354
View File
@@ -0,0 +1,354 @@
"""
用户/API Key RPM 限制器
支持两层叠加限流
1. 用户级或独立 Key RPM
2. 普通 Key 子限制 RPM
实现策略
- Redis 可用时使用分钟桶 + Lua 脚本原子检查/消费
- Redis 不可用时降级为内存计数仅适用于单实例
"""
from __future__ import annotations
import asyncio
import time
from dataclasses import dataclass
from datetime import datetime, timezone
import redis.asyncio as aioredis
from src.config.settings import config
from src.core.logger import logger
SYSTEM_RPM_CONFIG_KEY = "rate_limit_per_minute"
@dataclass(slots=True)
class RpmCheckResult:
"""RPM 检查结果。"""
allowed: bool
scope: str | None = None
limit: int | None = None
remaining: int | None = None
retry_after: int | None = None
class UserRpmLimiter:
"""用户/API Key 双层 RPM 限制器。
通过模块级 ``get_user_rpm_limiter()`` 工厂函数获取唯一实例
不要直接调用构造函数
"""
_CHECK_AND_CONSUME_SCRIPT = """
local user_key = KEYS[1]
local key_key = KEYS[2]
local user_limit = tonumber(ARGV[1])
local key_limit = tonumber(ARGV[2])
local ttl = tonumber(ARGV[3])
local retry_after = tonumber(ARGV[4])
local user_count = 0
if user_limit > 0 then
user_count = tonumber(redis.call('GET', user_key) or '0')
if user_count >= user_limit then
return {0, 1, user_limit, 0, retry_after}
end
end
local key_count = 0
if key_limit > 0 then
key_count = tonumber(redis.call('GET', key_key) or '0')
if key_count >= key_limit then
return {0, 2, key_limit, 0, retry_after}
end
end
local remaining = -1
if user_limit > 0 then
user_count = redis.call('INCR', user_key)
redis.call('EXPIRE', user_key, ttl)
remaining = user_limit - user_count
end
if key_limit > 0 then
key_count = redis.call('INCR', key_key)
redis.call('EXPIRE', key_key, ttl)
local key_remaining = key_limit - key_count
if remaining == -1 or key_remaining < remaining then
remaining = key_remaining
end
end
return {1, 0, 0, remaining, 0}
"""
def __init__(self) -> None:
self._redis: aioredis.Redis | None = None
self._bucket_seconds = int(config.rpm_bucket_seconds)
self._key_ttl_seconds = int(config.rpm_key_ttl_seconds)
self._cleanup_interval_seconds = int(config.rpm_cleanup_interval_seconds)
self._memory_lock: asyncio.Lock = asyncio.Lock()
self._memory_counts: dict[str, tuple[int, int]] = {}
self._cleanup_task: asyncio.Task | None = None
async def initialize(self) -> None:
if self._redis is not None:
return
try:
from src.clients.redis_client import get_redis_client
self._redis = await get_redis_client(require_redis=False)
if self._redis:
logger.info("[OK] UserRpmLimiter 已复用全局 Redis 客户端")
return
except Exception as exc:
logger.warning("初始化 UserRpmLimiter Redis 客户端失败,降级为内存模式: {}", exc)
self._redis = None
self._start_background_cleanup()
async def close(self) -> None:
if self._cleanup_task is not None:
self._cleanup_task.cancel()
try:
await self._cleanup_task
except asyncio.CancelledError:
pass
self._cleanup_task = None
@property
def bucket_seconds(self) -> int:
return self._bucket_seconds
def get_user_rpm_key(self, user_id: str, bucket: int | None = None) -> str:
b = bucket if bucket is not None else self._get_rpm_bucket()
return f"rpm:user:{user_id}:{b}"
def get_standalone_rpm_key(self, api_key_id: str, bucket: int | None = None) -> str:
b = bucket if bucket is not None else self._get_rpm_bucket()
return f"rpm:ukey:{api_key_id}:{b}"
def get_key_rpm_key(self, api_key_id: str, bucket: int | None = None) -> str:
b = bucket if bucket is not None else self._get_rpm_bucket()
return f"rpm:key:{api_key_id}:{b}"
def get_retry_after(self, now_ts: float | None = None) -> int:
ts = now_ts if now_ts is not None else time.time()
elapsed = int(ts % self._bucket_seconds)
return max(1, self._bucket_seconds - elapsed)
def get_reset_at(self, now_ts: float | None = None) -> datetime:
ts = now_ts if now_ts is not None else time.time()
bucket = self._get_rpm_bucket(ts)
reset_ts = (bucket + 1) * self._bucket_seconds
return datetime.fromtimestamp(reset_ts, tz=timezone.utc)
async def get_scope_count(self, scope_key: str) -> int:
await self.initialize()
if self._redis is None:
async with self._memory_lock:
bucket = self._get_rpm_bucket()
self._cleanup_expired_memory_counts(bucket)
return self._get_memory_count(scope_key)
try:
result = await self._redis.get(scope_key)
return int(result) if result else 0
except Exception as exc:
logger.warning("读取 RPM 计数失败,回退内存模式: {}", exc)
if config.rate_limit_fail_open:
return 0
async with self._memory_lock:
bucket = self._get_rpm_bucket()
self._cleanup_expired_memory_counts(bucket)
return self._get_memory_count(scope_key)
async def check_and_consume(
self,
*,
user_rpm_key: str,
user_rpm_limit: int,
key_rpm_key: str,
key_rpm_limit: int,
) -> RpmCheckResult:
"""原子检查并消费两层 RPM 配额。"""
await self.initialize()
normalized_user_limit = max(int(user_rpm_limit or 0), 0)
normalized_key_limit = max(int(key_rpm_limit or 0), 0)
if normalized_user_limit <= 0 and normalized_key_limit <= 0:
return RpmCheckResult(allowed=True)
if self._redis is None:
return await self._check_and_consume_memory(
user_rpm_key=user_rpm_key,
user_rpm_limit=normalized_user_limit,
key_rpm_key=key_rpm_key,
key_rpm_limit=normalized_key_limit,
)
retry_after = self.get_retry_after()
try:
raw_result = await self._redis.eval(
self._CHECK_AND_CONSUME_SCRIPT,
2,
user_rpm_key,
key_rpm_key,
normalized_user_limit,
normalized_key_limit,
self._key_ttl_seconds,
retry_after,
)
return self._parse_redis_result(raw_result)
except Exception as exc:
logger.warning("Redis RPM 检查失败: {}", exc)
if config.rate_limit_fail_open:
return RpmCheckResult(allowed=True)
return await self._check_and_consume_memory(
user_rpm_key=user_rpm_key,
user_rpm_limit=normalized_user_limit,
key_rpm_key=key_rpm_key,
key_rpm_limit=normalized_key_limit,
)
def _parse_redis_result(self, raw_result: object) -> RpmCheckResult:
values = list(raw_result) if isinstance(raw_result, (list, tuple)) else [raw_result]
allowed = int(values[0]) == 1
scope_code = int(values[1]) if len(values) > 1 else 0
limit = int(values[2]) if len(values) > 2 and values[2] is not None else None
remaining = int(values[3]) if len(values) > 3 and values[3] is not None else None
retry_after = int(values[4]) if len(values) > 4 and values[4] is not None else None
scope = {1: "user", 2: "key"}.get(scope_code)
return RpmCheckResult(
allowed=allowed,
scope=scope,
limit=limit,
remaining=remaining,
retry_after=retry_after,
)
async def _check_and_consume_memory(
self,
*,
user_rpm_key: str,
user_rpm_limit: int,
key_rpm_key: str,
key_rpm_limit: int,
) -> RpmCheckResult:
async with self._memory_lock:
bucket = self._get_rpm_bucket()
self._cleanup_expired_memory_counts(bucket)
user_count = self._get_memory_count(user_rpm_key)
if user_rpm_limit > 0 and user_count >= user_rpm_limit:
return RpmCheckResult(
allowed=False,
scope="user",
limit=user_rpm_limit,
remaining=0,
retry_after=self.get_retry_after(),
)
key_count = self._get_memory_count(key_rpm_key)
if key_rpm_limit > 0 and key_count >= key_rpm_limit:
return RpmCheckResult(
allowed=False,
scope="key",
limit=key_rpm_limit,
remaining=0,
retry_after=self.get_retry_after(),
)
remaining_candidates: list[int] = []
if user_rpm_limit > 0:
user_count += 1
self._set_memory_count(user_rpm_key, user_count)
remaining_candidates.append(user_rpm_limit - user_count)
if key_rpm_limit > 0:
key_count += 1
self._set_memory_count(key_rpm_key, key_count)
remaining_candidates.append(key_rpm_limit - key_count)
remaining = min(remaining_candidates) if remaining_candidates else None
return RpmCheckResult(allowed=True, remaining=remaining)
def _start_background_cleanup(self) -> None:
if self._cleanup_task is not None:
return
async def cleanup_loop() -> None:
while True:
try:
await asyncio.sleep(self._bucket_seconds)
async with self._memory_lock:
self._cleanup_expired_memory_counts(self._get_rpm_bucket())
except asyncio.CancelledError:
break
except Exception as exc:
logger.debug("UserRpmLimiter 后台清理异常: {}", exc)
try:
self._cleanup_task = asyncio.create_task(cleanup_loop())
except RuntimeError:
self._cleanup_task = None
def _get_rpm_bucket(self, now_ts: float | None = None) -> int:
ts = now_ts if now_ts is not None else time.time()
return int(ts // self._bucket_seconds)
def _split_scope_key(self, scope_key: str) -> tuple[str, int]:
base_key, bucket_str = scope_key.rsplit(":", 1)
return base_key, int(bucket_str)
def _get_memory_count(self, scope_key: str) -> int:
base_key, bucket = self._split_scope_key(scope_key)
stored = self._memory_counts.get(base_key)
if not stored:
return 0
stored_bucket, count = stored
if stored_bucket != bucket:
self._memory_counts.pop(base_key, None)
return 0
return count
def _set_memory_count(self, scope_key: str, count: int) -> None:
base_key, bucket = self._split_scope_key(scope_key)
self._memory_counts[base_key] = (bucket, count)
def _cleanup_expired_memory_counts(self, current_bucket: int) -> None:
expired_keys = [
base_key
for base_key, (bucket, _count) in self._memory_counts.items()
if bucket < current_bucket
]
for base_key in expired_keys:
self._memory_counts.pop(base_key, None)
if expired_keys:
logger.debug(
"[CLEANUP] 清理了 {} 个过期的用户/API Key RPM 计数(interval={}s",
len(expired_keys),
self._cleanup_interval_seconds,
)
_user_rpm_limiter: UserRpmLimiter | None = None
async def get_user_rpm_limiter() -> UserRpmLimiter:
global _user_rpm_limiter
if _user_rpm_limiter is None:
_user_rpm_limiter = UserRpmLimiter()
await _user_rpm_limiter.initialize()
return _user_rpm_limiter
+10 -36
View File
@@ -61,6 +61,10 @@ class ApiKeyService:
if final_expires_at is None and expire_days:
final_expires_at = datetime.now(timezone.utc) + timedelta(days=expire_days)
normalized_rate_limit = rate_limit
if not is_standalone and normalized_rate_limit is None:
normalized_rate_limit = 0
api_key = ApiKey(
user_id=user_id,
key_hash=key_hash,
@@ -69,7 +73,7 @@ class ApiKeyService:
allowed_providers=allowed_providers,
allowed_api_formats=allowed_api_formats,
allowed_models=allowed_models,
rate_limit=rate_limit,
rate_limit=normalized_rate_limit,
concurrent_limit=concurrent_limit,
expires_at=final_expires_at,
is_standalone=is_standalone,
@@ -144,7 +148,8 @@ class ApiKeyService:
# 允许显式设置为空数组/None 的字段(NULL=不限制,[]=全部禁用)
nullable_list_fields = {"allowed_providers", "allowed_api_formats", "allowed_models"}
# 允许显式设置为 None 的字段(如 expires_at=None 表示永不过期rate_limit=None 表示无限制)
# 允许显式设置为 None 的字段(如 expires_at=None 表示永不过期
# standalone rate_limit=None 表示继承系统默认)
nullable_fields = {"expires_at", "rate_limit"}
for field, value in kwargs.items():
@@ -154,7 +159,9 @@ class ApiKeyService:
if field in nullable_list_fields:
setattr(api_key, field, value)
elif field in nullable_fields:
# 这些字段允许显式设置为 None
if field == "rate_limit" and not api_key.is_standalone and value is None:
setattr(api_key, field, 0)
continue
setattr(api_key, field, value)
elif value is not None:
setattr(api_key, field, value)
@@ -180,39 +187,6 @@ class ApiKeyService:
logger.info(f"删除API密钥: ID {key_id}")
return True
@staticmethod
def check_rate_limit(db: Session, api_key: ApiKey, window_minutes: int = 1) -> tuple[bool, int]:
"""检查速率限制
Returns:
(is_allowed, remaining): 是否允许请求剩余可用次数
rate_limit None 时表示不限制返回 (True, -1)
"""
# 如果 rate_limit 为 None,表示不限制
if api_key.rate_limit is None:
return True, -1 # -1 表示无限制
# 计算时间窗口
window_start = datetime.now(timezone.utc) - timedelta(minutes=window_minutes)
# 统计窗口内的请求数
request_count = (
db.query(func.count(Usage.id))
.filter(Usage.api_key_id == api_key.id, Usage.created_at >= window_start)
.scalar()
or 0
)
# 检查是否超限
is_allowed = request_count < api_key.rate_limit
if not is_allowed:
logger.warning(
f"API密钥速率限制: Key ID {api_key.id}, 请求数 {request_count}/{api_key.rate_limit}"
)
return is_allowed, api_key.rate_limit - request_count
@staticmethod
def cleanup_expired_keys(db: Session, auto_delete: bool = False) -> int:
"""清理过期的API密钥
+4
View File
@@ -38,6 +38,7 @@ class UserService:
allowed_providers: list[str] | None = None,
allowed_api_formats: list[str] | None = None,
allowed_models: list[str] | None = None,
rate_limit: int | None = None,
) -> User:
"""创建新用户。"""
@@ -74,6 +75,7 @@ class UserService:
allowed_providers=allowed_providers,
allowed_api_formats=allowed_api_formats,
allowed_models=allowed_models,
rate_limit=rate_limit,
)
user.set_password(password)
@@ -218,6 +220,7 @@ class UserService:
"allowed_providers",
"allowed_api_formats",
"allowed_models",
"rate_limit",
]
# 允许设置为 None 的字段(表示无限制)
@@ -225,6 +228,7 @@ class UserService:
"allowed_providers",
"allowed_api_formats",
"allowed_models",
"rate_limit",
]
for field, value in kwargs.items():
+97 -4
View File
@@ -1,7 +1,9 @@
from __future__ import annotations
from contextlib import contextmanager
from datetime import datetime, timezone
from types import SimpleNamespace
from typing import Generator
from unittest.mock import MagicMock
import pytest
@@ -18,6 +20,7 @@ from src.api.admin.api_keys.routes import router as admin_api_keys_router
from src.api.admin.users.routes import (
AdminGetUserKeyFullKeyAdapter,
AdminToggleUserKeyLockAdapter,
AdminUpdateUserKeyAdapter,
)
from src.api.admin.users.routes import router as admin_users_router
from src.core.exceptions import InvalidRequestException, NotFoundException
@@ -25,6 +28,15 @@ from src.database import get_db
from src.models.api import CreateApiKeyRequest
def _patch_get_db_context(monkeypatch: pytest.MonkeyPatch, db: MagicMock) -> None:
@contextmanager
def _fake_ctx() -> Generator[MagicMock, None, None]:
yield db
monkeypatch.setattr("src.api.admin.users.routes.get_db_context", _fake_ctx)
monkeypatch.setattr("src.api.admin.api_keys.routes.get_db_context", _fake_ctx)
def _build_context(db: MagicMock) -> SimpleNamespace:
return SimpleNamespace(
db=db,
@@ -46,11 +58,15 @@ def _build_admin_users_app(db: MagicMock, monkeypatch: pytest.MonkeyPatch) -> Te
*, adapter: object, http_request: object, db: MagicMock, mode: object
) -> object:
_ = http_request, mode
try:
payload = await http_request.json()
except Exception:
payload = {}
context = SimpleNamespace(
db=db,
request=SimpleNamespace(state=SimpleNamespace()),
user=SimpleNamespace(id="admin-1"),
ensure_json_body=lambda: {},
ensure_json_body=lambda: payload,
add_audit_metadata=lambda **_: None,
)
return await adapter.handle(context)
@@ -82,10 +98,11 @@ def _build_admin_api_keys_app(db: MagicMock, monkeypatch: pytest.MonkeyPatch) ->
@pytest.mark.asyncio
async def test_toggle_user_key_lock_adapter_success() -> None:
async def test_toggle_user_key_lock_adapter_success(monkeypatch: pytest.MonkeyPatch) -> None:
db = MagicMock()
api_key = SimpleNamespace(id="key-1", user_id="user-1", is_standalone=False, is_locked=False)
_mock_query_first(db, api_key)
_patch_get_db_context(monkeypatch, db)
adapter = AdminToggleUserKeyLockAdapter(user_id="user-1", key_id="key-1")
result = await adapter.handle(_build_context(db))
@@ -98,9 +115,12 @@ async def test_toggle_user_key_lock_adapter_success() -> None:
@pytest.mark.asyncio
async def test_toggle_user_key_lock_adapter_not_found_for_standalone_or_wrong_owner() -> None:
async def test_toggle_user_key_lock_adapter_not_found_for_standalone_or_wrong_owner(
monkeypatch: pytest.MonkeyPatch,
) -> None:
db = MagicMock()
_mock_query_first(db, None)
_patch_get_db_context(monkeypatch, db)
adapter = AdminToggleUserKeyLockAdapter(user_id="user-1", key_id="key-standalone")
with pytest.raises(NotFoundException):
@@ -170,7 +190,9 @@ async def test_get_user_key_full_key_adapter_returns_500_on_decrypt_error(
@pytest.mark.asyncio
async def test_standalone_toggle_adapters_reject_normal_user_key() -> None:
async def test_standalone_toggle_adapters_reject_normal_user_key(
monkeypatch: pytest.MonkeyPatch,
) -> None:
db = MagicMock()
normal_key = SimpleNamespace(
id="key-user",
@@ -182,6 +204,7 @@ async def test_standalone_toggle_adapters_reject_normal_user_key() -> None:
updated_at=datetime.now(timezone.utc),
)
_mock_query_first(db, normal_key)
_patch_get_db_context(monkeypatch, db)
context = _build_context(db)
with pytest.raises(InvalidRequestException):
@@ -195,6 +218,7 @@ def test_user_key_lock_route_path_smoke(monkeypatch: pytest.MonkeyPatch) -> None
db = MagicMock()
api_key = SimpleNamespace(id="key-5", user_id="user-2", is_standalone=False, is_locked=False)
_mock_query_first(db, api_key)
_patch_get_db_context(monkeypatch, db)
client = _build_admin_users_app(db, monkeypatch)
response = client.patch("/api/admin/users/user-2/api-keys/key-5/lock")
@@ -220,6 +244,72 @@ def test_user_key_full_key_route_path_smoke(monkeypatch: pytest.MonkeyPatch) ->
assert response.json() == {"key": "sk-user-route-key"}
@pytest.mark.asyncio
async def test_update_user_key_adapter_passes_rate_limit_and_name(
monkeypatch: pytest.MonkeyPatch,
) -> None:
captured: dict[str, object] = {}
def _update_user_key_sync(
user_id: str, key_id: str, request: object
) -> tuple[dict[str, object], dict[str, object]]:
captured["user_id"] = user_id
captured["key_id"] = key_id
captured["name"] = getattr(request, "name", None)
captured["rate_limit"] = getattr(request, "rate_limit", None)
return {"id": key_id, "name": captured["name"], "rate_limit": captured["rate_limit"]}, {}
monkeypatch.setattr("src.api.admin.users.routes._update_user_key_sync", _update_user_key_sync)
adapter = AdminUpdateUserKeyAdapter(user_id="user-1", key_id="key-7")
context = SimpleNamespace(
db=MagicMock(),
request=SimpleNamespace(state=SimpleNamespace()),
ensure_json_body=lambda: {"name": "Renamed Key", "rate_limit": 12},
add_audit_metadata=lambda **_: None,
)
result = await adapter.handle(context)
assert result["id"] == "key-7"
assert captured == {
"user_id": "user-1",
"key_id": "key-7",
"name": "Renamed Key",
"rate_limit": 12,
}
def test_update_user_key_route_path_smoke(monkeypatch: pytest.MonkeyPatch) -> None:
captured: dict[str, object] = {}
def _update_user_key_sync(
user_id: str, key_id: str, request: object
) -> tuple[dict[str, object], dict[str, object]]:
captured["user_id"] = user_id
captured["key_id"] = key_id
captured["name"] = getattr(request, "name", None)
captured["rate_limit"] = getattr(request, "rate_limit", None)
return {"id": key_id, "name": captured["name"], "rate_limit": captured["rate_limit"]}, {}
monkeypatch.setattr("src.api.admin.users.routes._update_user_key_sync", _update_user_key_sync)
client = _build_admin_users_app(MagicMock(), monkeypatch)
response = client.put(
"/api/admin/users/user-2/api-keys/key-8",
json={"name": "Updated", "rate_limit": 9},
)
assert response.status_code == 200
assert response.json()["rate_limit"] == 9
assert captured == {
"user_id": "user-2",
"key_id": "key-8",
"name": "Updated",
"rate_limit": 9,
}
def test_standalone_lock_route_removed(monkeypatch: pytest.MonkeyPatch) -> None:
client = _build_admin_api_keys_app(MagicMock(), monkeypatch)
response = client.patch("/api/admin/api-keys/key-1/lock")
@@ -228,6 +318,7 @@ def test_standalone_lock_route_removed(monkeypatch: pytest.MonkeyPatch) -> None:
def test_standalone_list_route_does_not_expose_is_locked(monkeypatch: pytest.MonkeyPatch) -> None:
db = MagicMock()
_patch_get_db_context(monkeypatch, db)
api_key = SimpleNamespace(
id="sa-key-1",
user_id="admin-1",
@@ -306,6 +397,7 @@ async def test_create_standalone_key_adapter_preserves_empty_restriction_lists(
monkeypatch: pytest.MonkeyPatch,
) -> None:
db = MagicMock()
_patch_get_db_context(monkeypatch, db)
captured: dict[str, object] = {}
created_key = SimpleNamespace(
id="sa-key-3",
@@ -365,6 +457,7 @@ async def test_update_standalone_key_adapter_preserves_empty_restriction_lists(
monkeypatch: pytest.MonkeyPatch,
) -> None:
db = MagicMock()
_patch_get_db_context(monkeypatch, db)
existing_key = SimpleNamespace(id="sa-key-4", is_standalone=True)
_mock_query_first(db, existing_key)
+9 -17
View File
@@ -49,6 +49,7 @@ def test_list_users_uses_wallet_batch_lookup(monkeypatch: pytest.MonkeyPatch) ->
allowed_providers=None,
allowed_api_formats=None,
allowed_models=None,
rate_limit=None,
is_active=True,
created_at=now,
updated_at=now,
@@ -62,6 +63,7 @@ def test_list_users_uses_wallet_batch_lookup(monkeypatch: pytest.MonkeyPatch) ->
allowed_providers=None,
allowed_api_formats=None,
allowed_models=None,
rate_limit=None,
is_active=True,
created_at=now,
updated_at=None,
@@ -101,24 +103,14 @@ async def test_create_user_adapter_preserves_empty_restriction_lists(
db = MagicMock()
captured: dict[str, Any] = {}
def _create_user(**kwargs: Any) -> SimpleNamespace:
captured.update(kwargs)
return SimpleNamespace(
id="user-3",
email="u3@example.com",
username="user3",
role=SimpleNamespace(value="user"),
is_active=True,
allowed_providers=[],
allowed_api_formats=[],
allowed_models=[],
)
def _create_user_sync(request: Any, role: Any) -> tuple[dict[str, Any], dict[str, Any]]:
captured["allowed_providers"] = request.allowed_providers
captured["allowed_api_formats"] = request.allowed_api_formats
captured["allowed_models"] = request.allowed_models
captured["role"] = role
return {"id": "user-3"}, {}
monkeypatch.setattr("src.api.admin.users.routes.UserService.create_user", _create_user)
monkeypatch.setattr(
"src.api.admin.users.routes._serialize_user",
lambda _db, user: {"id": user.id},
)
monkeypatch.setattr("src.api.admin.users.routes._create_user_sync", _create_user_sync)
context = SimpleNamespace(
db=db,
@@ -0,0 +1,120 @@
from __future__ import annotations
from datetime import datetime, timezone
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
import pytest
from src.api.monitoring.user import UserRateLimitStatusAdapter
def _build_query_with_keys(keys: list[object]) -> MagicMock:
query = MagicMock()
query.filter.return_value.order_by.return_value.all.return_value = keys
return query
@pytest.mark.asyncio
async def test_rate_limit_status_adapter_reports_user_and_key_layers(
monkeypatch: pytest.MonkeyPatch,
) -> None:
now = datetime(2026, 3, 13, 12, 0, tzinfo=timezone.utc)
db = MagicMock()
user = SimpleNamespace(id="user-1", rate_limit=None)
key = SimpleNamespace(id="key-1", name="Primary", is_standalone=False, rate_limit=10)
standalone = SimpleNamespace(
id="skey-1", name="Standalone", is_standalone=True, rate_limit=None
)
db.query.return_value = _build_query_with_keys([key, standalone])
limiter = MagicMock()
limiter.bucket_seconds = 60
limiter.get_reset_at.return_value = now
limiter.get_user_rpm_key.return_value = "rpm:user:user-1:bucket"
limiter.get_standalone_rpm_key.return_value = "rpm:ukey:skey-1:bucket"
limiter.get_key_rpm_key.side_effect = lambda key_id: f"rpm:key:{key_id}:bucket"
async def _get_scope_count(scope_key: str) -> int:
counts = {
"rpm:user:user-1:bucket": 55,
"rpm:key:key-1:bucket": 7,
"rpm:ukey:skey-1:bucket": 12,
}
return counts[scope_key]
limiter.get_scope_count = AsyncMock(side_effect=_get_scope_count)
monkeypatch.setattr(
"src.api.monitoring.user.get_user_rpm_limiter",
AsyncMock(return_value=limiter),
)
monkeypatch.setattr(
"src.api.monitoring.user.SystemConfigService.get_config",
lambda *_a, **_k: 60,
)
context = SimpleNamespace(db=db, user=user)
result = await UserRateLimitStatusAdapter().handle(context)
assert result["user_id"] == "user-1"
assert result["api_keys"][0] == {
"api_key_name": "Primary",
"limit": 10,
"remaining": 3,
"scope": "key",
"reset_time": now.isoformat(),
"window": "60s",
"user_limit": 60,
"user_remaining": 5,
"key_limit": 10,
"key_remaining": 3,
}
assert result["api_keys"][1] == {
"api_key_name": "Standalone",
"limit": 60,
"remaining": 48,
"scope": "user",
"reset_time": now.isoformat(),
"window": "60s",
"user_limit": 60,
"user_remaining": 48,
"key_limit": None,
"key_remaining": None,
}
@pytest.mark.asyncio
async def test_rate_limit_status_adapter_reports_unlimited_key_without_counts(
monkeypatch: pytest.MonkeyPatch,
) -> None:
db = MagicMock()
user = SimpleNamespace(id="user-1", rate_limit=0)
key = SimpleNamespace(id="key-1", name="Unlimited", is_standalone=False, rate_limit=0)
db.query.return_value = _build_query_with_keys([key])
limiter = MagicMock()
limiter.bucket_seconds = 60
limiter.get_reset_at.return_value = datetime.now(timezone.utc)
limiter.get_user_rpm_key.return_value = "rpm:user:user-1:bucket"
limiter.get_key_rpm_key.return_value = "rpm:key:key-1:bucket"
limiter.get_scope_count = AsyncMock()
monkeypatch.setattr(
"src.api.monitoring.user.get_user_rpm_limiter",
AsyncMock(return_value=limiter),
)
monkeypatch.setattr(
"src.api.monitoring.user.SystemConfigService.get_config",
lambda *_a, **_k: 60,
)
context = SimpleNamespace(db=db, user=user)
result = await UserRateLimitStatusAdapter().handle(context)
assert result["api_keys"][0]["limit"] is None
assert result["api_keys"][0]["remaining"] is None
assert result["api_keys"][0]["scope"] is None
limiter.get_scope_count.assert_not_awaited()
+161
View File
@@ -18,6 +18,7 @@ from src.api.base.adapter import ApiMode
from src.api.base.pipeline import ApiRequestPipeline
from src.core.enums import UserRole
from src.core.modules.hooks import AUTH_TOKEN_PREFIX_AUTHENTICATORS
from src.services.rate_limit.user_rpm_limiter import RpmCheckResult
class TestPipelineBalanceCalculation:
@@ -499,6 +500,166 @@ class TestPipelineAuthentication:
assert "锁定" in str(exc_info.value.detail)
class TestPipelineUserRateLimit:
@pytest.fixture
def pipeline(self) -> ApiRequestPipeline:
return ApiRequestPipeline()
@pytest.mark.asyncio
async def test_check_user_rate_limit_uses_system_default_for_user_scope(
self, pipeline: ApiRequestPipeline, monkeypatch: pytest.MonkeyPatch
) -> None:
request = MagicMock()
request.state = MagicMock()
db = MagicMock()
user = MagicMock(id="user-1", rate_limit=None)
api_key = MagicMock(id="key-1", is_standalone=False, rate_limit=0)
limiter = MagicMock()
limiter.get_user_rpm_key.return_value = "rpm:user:user-1:1"
limiter.get_key_rpm_key.return_value = "rpm:key:key-1:1"
limiter.check_and_consume = AsyncMock(return_value=RpmCheckResult(allowed=True))
monkeypatch.setattr(
"src.api.base.pipeline.get_user_rpm_limiter",
AsyncMock(return_value=limiter),
)
monkeypatch.setattr(
"src.api.base.pipeline.SystemConfigService.get_config",
lambda *_a, **_k: 60,
)
await pipeline._check_user_rate_limit(request, db, user, api_key)
limiter.check_and_consume.assert_awaited_once_with(
user_rpm_key="rpm:user:user-1:1",
user_rpm_limit=60,
key_rpm_key="rpm:key:key-1:1",
key_rpm_limit=0,
)
@pytest.mark.asyncio
async def test_check_user_rate_limit_returns_429_with_scope_header(
self, pipeline: ApiRequestPipeline, monkeypatch: pytest.MonkeyPatch
) -> None:
request = MagicMock()
request.state = MagicMock()
db = MagicMock()
user = MagicMock(id="user-1", rate_limit=100)
api_key = MagicMock(id="key-1", is_standalone=False, rate_limit=10)
limiter = MagicMock()
limiter.get_user_rpm_key.return_value = "rpm:user:user-1:1"
limiter.get_key_rpm_key.return_value = "rpm:key:key-1:1"
limiter.get_retry_after.return_value = 17
limiter.check_and_consume = AsyncMock(
return_value=RpmCheckResult(
allowed=False,
scope="key",
limit=10,
remaining=0,
retry_after=17,
)
)
monkeypatch.setattr(
"src.api.base.pipeline.get_user_rpm_limiter",
AsyncMock(return_value=limiter),
)
monkeypatch.setattr(
"src.api.base.pipeline.SystemConfigService.get_config",
lambda *_a, **_k: 60,
)
with pytest.raises(HTTPException) as exc_info:
await pipeline._check_user_rate_limit(request, db, user, api_key)
assert exc_info.value.status_code == 429
assert exc_info.value.headers == {
"Retry-After": "17",
"X-RateLimit-Limit": "10",
"X-RateLimit-Remaining": "0",
"X-RateLimit-Scope": "key",
}
@pytest.mark.asyncio
async def test_check_user_rate_limit_uses_system_default_for_standalone_key(
self, pipeline: ApiRequestPipeline, monkeypatch: pytest.MonkeyPatch
) -> None:
request = MagicMock()
request.state = MagicMock()
db = MagicMock()
user = MagicMock(id="user-1", rate_limit=999)
api_key = MagicMock(id="standalone-1", is_standalone=True, rate_limit=None)
limiter = MagicMock()
limiter.get_standalone_rpm_key.return_value = "rpm:ukey:standalone-1:1"
limiter.get_key_rpm_key.return_value = "rpm:key:standalone-1:1"
limiter.check_and_consume = AsyncMock(return_value=RpmCheckResult(allowed=True))
monkeypatch.setattr(
"src.api.base.pipeline.get_user_rpm_limiter",
AsyncMock(return_value=limiter),
)
monkeypatch.setattr(
"src.api.base.pipeline.SystemConfigService.get_config",
lambda *_a, **_k: 60,
)
await pipeline._check_user_rate_limit(request, db, user, api_key)
limiter.check_and_consume.assert_awaited_once_with(
user_rpm_key="rpm:ukey:standalone-1:1",
user_rpm_limit=60,
key_rpm_key="rpm:key:standalone-1:1",
key_rpm_limit=0,
)
@pytest.mark.asyncio
async def test_check_user_rate_limit_returns_429_with_user_scope_header(
self, pipeline: ApiRequestPipeline, monkeypatch: pytest.MonkeyPatch
) -> None:
request = MagicMock()
request.state = MagicMock()
db = MagicMock()
user = MagicMock(id="user-1", rate_limit=3)
api_key = MagicMock(id="key-1", is_standalone=False, rate_limit=10)
limiter = MagicMock()
limiter.get_user_rpm_key.return_value = "rpm:user:user-1:1"
limiter.get_key_rpm_key.return_value = "rpm:key:key-1:1"
limiter.get_retry_after.return_value = 23
limiter.check_and_consume = AsyncMock(
return_value=RpmCheckResult(
allowed=False,
scope="user",
limit=3,
remaining=0,
retry_after=23,
)
)
monkeypatch.setattr(
"src.api.base.pipeline.get_user_rpm_limiter",
AsyncMock(return_value=limiter),
)
monkeypatch.setattr(
"src.api.base.pipeline.SystemConfigService.get_config",
lambda *_a, **_k: 60,
)
with pytest.raises(HTTPException) as exc_info:
await pipeline._check_user_rate_limit(request, db, user, api_key)
assert exc_info.value.status_code == 429
assert exc_info.value.headers == {
"Retry-After": "23",
"X-RateLimit-Limit": "3",
"X-RateLimit-Remaining": "0",
"X-RateLimit-Scope": "user",
}
class TestPipelineTokenPrefixAuth:
"""Tests token-prefix auth isolation."""
+110
View File
@@ -0,0 +1,110 @@
from __future__ import annotations
from types import SimpleNamespace
from typing import Any
from unittest.mock import MagicMock
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from src.api.user_me.routes import UpdateMyApiKeyAdapter
from src.api.user_me.routes import router as me_router
from src.database import get_db
def _build_me_app(db: MagicMock, monkeypatch: Any) -> TestClient:
app = FastAPI()
app.include_router(me_router)
app.dependency_overrides[get_db] = lambda: db
async def _fake_pipeline_run(
*, adapter: object, http_request: object, db: MagicMock, mode: object
) -> object:
_ = http_request, mode
try:
payload = await http_request.json()
except Exception:
payload = {}
context = SimpleNamespace(
db=db,
user=SimpleNamespace(id="user-1", email="u@example.com"),
request=SimpleNamespace(state=SimpleNamespace()),
ensure_json_body=lambda: payload,
add_audit_metadata=lambda **_: None,
)
return await adapter.handle(context)
monkeypatch.setattr("src.api.user_me.routes.pipeline.run", _fake_pipeline_run)
return TestClient(app)
async def _fake_update_my_api_key_sync(
user_id: str,
key_id: str,
request: object,
captured: dict[str, object],
) -> dict[str, object]:
captured["user_id"] = user_id
captured["key_id"] = key_id
captured["name"] = getattr(request, "name", None)
captured["rate_limit"] = getattr(request, "rate_limit", None)
return {"id": key_id, "name": captured["name"], "rate_limit": captured["rate_limit"]}
def test_update_my_api_key_route_path_smoke(monkeypatch: Any) -> None:
captured: dict[str, object] = {}
def _sync(user_id: str, key_id: str, request: object) -> dict[str, object]:
captured["user_id"] = user_id
captured["key_id"] = key_id
captured["name"] = getattr(request, "name", None)
captured["rate_limit"] = getattr(request, "rate_limit", None)
return {"id": key_id, "name": captured["name"], "rate_limit": captured["rate_limit"]}
monkeypatch.setattr("src.api.user_me.routes._update_my_api_key_sync", _sync)
client = _build_me_app(MagicMock(), monkeypatch)
response = client.put("/api/users/me/api-keys/key-1", json={"name": "Edited", "rate_limit": 6})
assert response.status_code == 200
assert response.json()["rate_limit"] == 6
assert captured == {
"user_id": "user-1",
"key_id": "key-1",
"name": "Edited",
"rate_limit": 6,
}
@pytest.mark.asyncio
async def test_update_my_api_key_adapter_passes_rate_limit_and_name(monkeypatch: Any) -> None:
captured: dict[str, object] = {}
def _sync(user_id: str, key_id: str, request: object) -> dict[str, object]:
captured["user_id"] = user_id
captured["key_id"] = key_id
captured["name"] = getattr(request, "name", None)
captured["rate_limit"] = getattr(request, "rate_limit", None)
return {"id": key_id, "name": captured["name"], "rate_limit": captured["rate_limit"]}
monkeypatch.setattr("src.api.user_me.routes._update_my_api_key_sync", _sync)
adapter = UpdateMyApiKeyAdapter(key_id="key-2")
context = SimpleNamespace(
db=MagicMock(),
user=SimpleNamespace(id="user-1"),
request=SimpleNamespace(state=SimpleNamespace()),
ensure_json_body=lambda: {"name": "Edited Again", "rate_limit": 15},
add_audit_metadata=lambda **_: None,
)
result = await adapter.handle(context)
assert result["id"] == "key-2"
assert captured == {
"user_id": "user-1",
"key_id": "key-2",
"name": "Edited Again",
"rate_limit": 15,
}
+148
View File
@@ -0,0 +1,148 @@
from __future__ import annotations
from collections.abc import AsyncGenerator
from unittest.mock import AsyncMock, MagicMock
import pytest
from src.services.rate_limit.user_rpm_limiter import RpmCheckResult, UserRpmLimiter
@pytest.fixture
async def limiter(monkeypatch: pytest.MonkeyPatch) -> AsyncGenerator[UserRpmLimiter]:
limiter = UserRpmLimiter()
limiter._redis = None
limiter._memory_counts.clear()
await limiter.close()
# 阻止 initialize() 连接真实 Redis,确保纯内存模式
monkeypatch.setattr(limiter, "initialize", AsyncMock())
yield limiter
limiter._memory_counts.clear()
limiter._redis = None
await limiter.close()
@pytest.mark.asyncio
async def test_check_and_consume_enforces_user_scope_in_memory(
limiter: UserRpmLimiter,
) -> None:
user_key = limiter.get_user_rpm_key("user-1")
key_key = limiter.get_key_rpm_key("key-1")
first = await limiter.check_and_consume(
user_rpm_key=user_key,
user_rpm_limit=1,
key_rpm_key=key_key,
key_rpm_limit=0,
)
second = await limiter.check_and_consume(
user_rpm_key=user_key,
user_rpm_limit=1,
key_rpm_key=key_key,
key_rpm_limit=0,
)
assert first == RpmCheckResult(allowed=True, remaining=0)
assert second.allowed is False
assert second.scope == "user"
assert second.limit == 1
assert second.remaining == 0
@pytest.mark.asyncio
async def test_check_and_consume_enforces_key_scope_in_memory(
limiter: UserRpmLimiter,
) -> None:
user_key = limiter.get_user_rpm_key("user-1")
key_key = limiter.get_key_rpm_key("key-1")
first = await limiter.check_and_consume(
user_rpm_key=user_key,
user_rpm_limit=3,
key_rpm_key=key_key,
key_rpm_limit=1,
)
second = await limiter.check_and_consume(
user_rpm_key=user_key,
user_rpm_limit=3,
key_rpm_key=key_key,
key_rpm_limit=1,
)
assert first.allowed is True
assert second.allowed is False
assert second.scope == "key"
assert second.limit == 1
@pytest.mark.asyncio
async def test_check_and_consume_redis_failure_falls_back_to_memory_when_fail_close(
monkeypatch: pytest.MonkeyPatch,
limiter: UserRpmLimiter,
) -> None:
fake_redis = MagicMock()
fake_redis.eval = AsyncMock(side_effect=RuntimeError("redis down"))
limiter._redis = fake_redis
monkeypatch.setattr(
"src.services.rate_limit.user_rpm_limiter.config.rate_limit_fail_open", False
)
user_key = limiter.get_user_rpm_key("user-1")
key_key = limiter.get_key_rpm_key("key-1")
first = await limiter.check_and_consume(
user_rpm_key=user_key,
user_rpm_limit=1,
key_rpm_key=key_key,
key_rpm_limit=0,
)
second = await limiter.check_and_consume(
user_rpm_key=user_key,
user_rpm_limit=1,
key_rpm_key=key_key,
key_rpm_limit=0,
)
assert first.allowed is True
assert second.allowed is False
assert second.scope == "user"
@pytest.mark.asyncio
async def test_check_and_consume_redis_failure_fail_open_allows_request(
monkeypatch: pytest.MonkeyPatch,
limiter: UserRpmLimiter,
) -> None:
fake_redis = MagicMock()
fake_redis.eval = AsyncMock(side_effect=RuntimeError("redis down"))
limiter._redis = fake_redis
monkeypatch.setattr(
"src.services.rate_limit.user_rpm_limiter.config.rate_limit_fail_open", True
)
result = await limiter.check_and_consume(
user_rpm_key=limiter.get_user_rpm_key("user-1"),
user_rpm_limit=1,
key_rpm_key=limiter.get_key_rpm_key("key-1"),
key_rpm_limit=0,
)
assert result.allowed is True
@pytest.mark.asyncio
async def test_check_and_consume_skips_when_all_limits_are_zero(
limiter: UserRpmLimiter,
) -> None:
result = await limiter.check_and_consume(
user_rpm_key=limiter.get_user_rpm_key("user-1"),
user_rpm_limit=0,
key_rpm_key=limiter.get_key_rpm_key("key-1"),
key_rpm_limit=0,
)
assert result.allowed is True
assert result.scope is None
assert limiter._memory_counts == {}
@@ -53,3 +53,29 @@ def test_import_user_api_key_material_keeps_legacy_encrypted_payload() -> None:
assert key_hash == legacy_hash
assert key_encrypted == legacy_encrypted
def test_import_user_rate_limit_defaults_to_none_when_missing() -> None:
assert AdminImportUsersAdapter._normalize_imported_user_rate_limit({}) is None
def test_import_legacy_standalone_key_null_rate_limit_becomes_unlimited() -> None:
assert (
AdminImportUsersAdapter._normalize_imported_api_key_rate_limit(
{"rate_limit": None},
is_standalone=True,
legacy_export=True,
)
== 0
)
def test_import_new_standalone_key_null_rate_limit_keeps_inherit_semantics() -> None:
assert (
AdminImportUsersAdapter._normalize_imported_api_key_rate_limit(
{"rate_limit": None},
is_standalone=True,
legacy_export=False,
)
is None
)