mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
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:
@@ -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")
|
||||||
@@ -86,6 +86,7 @@ export interface UserExport {
|
|||||||
allowed_providers?: string[] | null
|
allowed_providers?: string[] | null
|
||||||
allowed_api_formats?: string[] | null
|
allowed_api_formats?: string[] | null
|
||||||
allowed_models?: string[] | null
|
allowed_models?: string[] | null
|
||||||
|
rate_limit?: number | null // null = 跟随系统默认,0 = 不限制
|
||||||
model_capability_settings?: Record<string, Record<string, boolean>>
|
model_capability_settings?: Record<string, Record<string, boolean>>
|
||||||
unlimited?: boolean
|
unlimited?: boolean
|
||||||
wallet?: BillingSummary | null
|
wallet?: BillingSummary | null
|
||||||
@@ -102,7 +103,7 @@ export interface UserApiKeyExport {
|
|||||||
allowed_providers?: string[] | null
|
allowed_providers?: string[] | null
|
||||||
allowed_api_formats?: string[] | null
|
allowed_api_formats?: string[] | null
|
||||||
allowed_models?: 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
|
concurrent_limit?: number | null
|
||||||
force_capabilities?: Record<string, boolean>
|
force_capabilities?: Record<string, boolean>
|
||||||
is_active: boolean
|
is_active: boolean
|
||||||
@@ -349,7 +350,7 @@ export interface AdminApiKey {
|
|||||||
total_requests?: number
|
total_requests?: number
|
||||||
total_tokens?: number
|
total_tokens?: number
|
||||||
total_cost_usd?: number
|
total_cost_usd?: number
|
||||||
rate_limit?: number | null // null = 无限制
|
rate_limit?: number | null // null = 跟随系统默认,0 = 不限制
|
||||||
allowed_providers?: string[] | null // 允许的提供商列表
|
allowed_providers?: string[] | null // 允许的提供商列表
|
||||||
allowed_api_formats?: string[] | null // 允许的 API 格式列表
|
allowed_api_formats?: string[] | null // 允许的 API 格式列表
|
||||||
allowed_models?: string[] | null // 允许的模型列表
|
allowed_models?: string[] | null // 允许的模型列表
|
||||||
@@ -365,7 +366,7 @@ export interface CreateStandaloneApiKeyRequest {
|
|||||||
allowed_providers?: string[] | null
|
allowed_providers?: string[] | null
|
||||||
allowed_api_formats?: string[] | null
|
allowed_api_formats?: string[] | null
|
||||||
allowed_models?: 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 = 永不过期
|
expires_at?: string | null // ISO 日期字符串,如 "2025-12-31",null = 永不过期
|
||||||
initial_balance_usd: number | null // 初始余额,null = 无限制
|
initial_balance_usd: number | null // 初始余额,null = 无限制
|
||||||
unlimited_balance?: boolean | null // 编辑时仅切换额度模式,不调整余额数值
|
unlimited_balance?: boolean | null // 编辑时仅切换额度模式,不调整余额数值
|
||||||
|
|||||||
@@ -142,6 +142,7 @@ export interface ApiKey {
|
|||||||
created_at: string
|
created_at: string
|
||||||
total_requests?: number
|
total_requests?: number
|
||||||
total_cost_usd?: number
|
total_cost_usd?: number
|
||||||
|
rate_limit?: number | null
|
||||||
allowed_providers?: ProviderConfig[]
|
allowed_providers?: ProviderConfig[]
|
||||||
force_capabilities?: Record<string, boolean> | null // 强制能力配置
|
force_capabilities?: Record<string, boolean> | null // 强制能力配置
|
||||||
}
|
}
|
||||||
@@ -181,8 +182,8 @@ export const meApi = {
|
|||||||
return response.data
|
return response.data
|
||||||
},
|
},
|
||||||
|
|
||||||
async createApiKey(name: string): Promise<ApiKey> {
|
async createApiKey(data: { name: string; rate_limit?: number }): Promise<ApiKey> {
|
||||||
const response = await apiClient.post<ApiKey>('/api/users/me/api-keys', { name })
|
const response = await apiClient.post<ApiKey>('/api/users/me/api-keys', data)
|
||||||
return response.data
|
return response.data
|
||||||
},
|
},
|
||||||
|
|
||||||
@@ -212,6 +213,17 @@ export const meApi = {
|
|||||||
return response.data
|
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?: {
|
async getUsage(params?: {
|
||||||
start_date?: string
|
start_date?: string
|
||||||
|
|||||||
@@ -10,6 +10,7 @@ export interface User {
|
|||||||
allowed_providers: string[] | null // 允许使用的提供商 ID 列表
|
allowed_providers: string[] | null // 允许使用的提供商 ID 列表
|
||||||
allowed_api_formats: string[] | null // 允许使用的 API 格式列表
|
allowed_api_formats: string[] | null // 允许使用的 API 格式列表
|
||||||
allowed_models: string[] | null // 允许使用的模型名称列表
|
allowed_models: string[] | null // 允许使用的模型名称列表
|
||||||
|
rate_limit?: number | null // null = 跟随系统默认,0 = 不限制
|
||||||
created_at: string
|
created_at: string
|
||||||
updated_at?: string
|
updated_at?: string
|
||||||
last_login_at?: string | null
|
last_login_at?: string | null
|
||||||
@@ -25,6 +26,7 @@ export interface CreateUserRequest {
|
|||||||
allowed_providers?: string[] | null
|
allowed_providers?: string[] | null
|
||||||
allowed_api_formats?: string[] | null
|
allowed_api_formats?: string[] | null
|
||||||
allowed_models?: string[] | null
|
allowed_models?: string[] | null
|
||||||
|
rate_limit?: number | null
|
||||||
}
|
}
|
||||||
|
|
||||||
export interface UpdateUserRequest {
|
export interface UpdateUserRequest {
|
||||||
@@ -36,6 +38,7 @@ export interface UpdateUserRequest {
|
|||||||
allowed_providers?: string[] | null
|
allowed_providers?: string[] | null
|
||||||
allowed_api_formats?: string[] | null
|
allowed_api_formats?: string[] | null
|
||||||
allowed_models?: string[] | null
|
allowed_models?: string[] | null
|
||||||
|
rate_limit?: number | null
|
||||||
}
|
}
|
||||||
|
|
||||||
export interface ApiKey {
|
export interface ApiKey {
|
||||||
@@ -49,11 +52,16 @@ export interface ApiKey {
|
|||||||
is_active: boolean
|
is_active: boolean
|
||||||
is_locked: boolean // 管理员锁定标志
|
is_locked: boolean // 管理员锁定标志
|
||||||
is_standalone: boolean // 是否为独立余额Key
|
is_standalone: boolean // 是否为独立余额Key
|
||||||
rate_limit?: number // 速率限制(请求/分钟)
|
rate_limit?: number | null // 普通Key: 0 = 不限制,历史 null 视为跟随系统默认
|
||||||
total_requests?: number // 总请求数
|
total_requests?: number // 总请求数
|
||||||
total_cost_usd?: number // 总费用
|
total_cost_usd?: number // 总费用
|
||||||
}
|
}
|
||||||
|
|
||||||
|
export interface UpsertUserApiKeyRequest {
|
||||||
|
name?: string
|
||||||
|
rate_limit?: number | null
|
||||||
|
}
|
||||||
|
|
||||||
export const usersApi = {
|
export const usersApi = {
|
||||||
async getAllUsers(): Promise<User[]> {
|
async getAllUsers(): Promise<User[]> {
|
||||||
const response = await apiClient.get<User[]>('/api/admin/users')
|
const response = await apiClient.get<User[]>('/api/admin/users')
|
||||||
@@ -84,8 +92,26 @@ export const usersApi = {
|
|||||||
return response.data.api_keys
|
return response.data.api_keys
|
||||||
},
|
},
|
||||||
|
|
||||||
async createApiKey(userId: string, name?: string): Promise<ApiKey & { key: string }> {
|
async createApiKey(
|
||||||
const response = await apiClient.post<ApiKey & { key: string }>(`/api/admin/users/${userId}/api-keys`, { name })
|
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
|
return response.data
|
||||||
},
|
},
|
||||||
|
|
||||||
|
|||||||
@@ -96,23 +96,6 @@
|
|||||||
{{ form.expires_at ? '到期后' + (form.auto_delete_on_expiry ? '自动删除' : '仅禁用') + '(当天 23:59 失效)' : '留空表示永不过期' }}
|
{{ form.expires_at ? '到期后' + (form.auto_delete_on_expiry ? '自动删除' : '仅禁用') + '(当天 23:59 失效)' : '留空表示永不过期' }}
|
||||||
</p>
|
</p>
|
||||||
</div>
|
</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>
|
</div>
|
||||||
|
|
||||||
<!-- 右侧:访问限制 -->
|
<!-- 右侧:访问限制 -->
|
||||||
@@ -190,6 +173,36 @@
|
|||||||
</div>
|
</div>
|
||||||
</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">
|
<div class="space-y-2">
|
||||||
<Label class="text-sm font-medium">额度</Label>
|
<Label class="text-sm font-medium">额度</Label>
|
||||||
@@ -267,7 +280,7 @@ export interface StandaloneKeyFormData {
|
|||||||
initial_balance_usd?: number
|
initial_balance_usd?: number
|
||||||
unlimited_balance?: boolean
|
unlimited_balance?: boolean
|
||||||
expires_at?: string // ISO 日期字符串,如 "2025-12-31",undefined = 永不过期
|
expires_at?: string // ISO 日期字符串,如 "2025-12-31",undefined = 永不过期
|
||||||
rate_limit?: number
|
rate_limit?: number | null
|
||||||
auto_delete_on_expiry: boolean
|
auto_delete_on_expiry: boolean
|
||||||
allowed_providers?: string[] | null
|
allowed_providers?: string[] | null
|
||||||
allowed_api_formats?: string[] | null
|
allowed_api_formats?: string[] | null
|
||||||
@@ -280,6 +293,7 @@ interface StandaloneKeyFormState {
|
|||||||
initial_balance_usd?: number
|
initial_balance_usd?: number
|
||||||
unlimited_balance?: boolean
|
unlimited_balance?: boolean
|
||||||
expires_at?: string
|
expires_at?: string
|
||||||
|
rate_limit_inherited: boolean
|
||||||
rate_limit?: number
|
rate_limit?: number
|
||||||
auto_delete_on_expiry: boolean
|
auto_delete_on_expiry: boolean
|
||||||
provider_unrestricted: boolean
|
provider_unrestricted: boolean
|
||||||
@@ -333,6 +347,7 @@ const form = ref<StandaloneKeyFormState>({
|
|||||||
initial_balance_usd: 10,
|
initial_balance_usd: 10,
|
||||||
unlimited_balance: false,
|
unlimited_balance: false,
|
||||||
expires_at: undefined,
|
expires_at: undefined,
|
||||||
|
rate_limit_inherited: true,
|
||||||
rate_limit: undefined,
|
rate_limit: undefined,
|
||||||
auto_delete_on_expiry: false,
|
auto_delete_on_expiry: false,
|
||||||
provider_unrestricted: true,
|
provider_unrestricted: true,
|
||||||
@@ -356,6 +371,7 @@ function resetForm() {
|
|||||||
initial_balance_usd: 10,
|
initial_balance_usd: 10,
|
||||||
unlimited_balance: false,
|
unlimited_balance: false,
|
||||||
expires_at: undefined,
|
expires_at: undefined,
|
||||||
|
rate_limit_inherited: true,
|
||||||
rate_limit: undefined,
|
rate_limit: undefined,
|
||||||
auto_delete_on_expiry: false,
|
auto_delete_on_expiry: false,
|
||||||
provider_unrestricted: true,
|
provider_unrestricted: true,
|
||||||
@@ -375,7 +391,8 @@ function loadKeyData() {
|
|||||||
initial_balance_usd: props.apiKey.initial_balance_usd,
|
initial_balance_usd: props.apiKey.initial_balance_usd,
|
||||||
unlimited_balance: props.apiKey.initial_balance_usd == null,
|
unlimited_balance: props.apiKey.initial_balance_usd == null,
|
||||||
expires_at: props.apiKey.expires_at,
|
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,
|
auto_delete_on_expiry: props.apiKey.auto_delete_on_expiry,
|
||||||
provider_unrestricted: props.apiKey.allowed_providers == null,
|
provider_unrestricted: props.apiKey.allowed_providers == null,
|
||||||
api_format_unrestricted: props.apiKey.allowed_api_formats == null,
|
api_format_unrestricted: props.apiKey.allowed_api_formats == null,
|
||||||
@@ -425,7 +442,7 @@ function handleSubmit() {
|
|||||||
initial_balance_usd: form.value.initial_balance_usd,
|
initial_balance_usd: form.value.initial_balance_usd,
|
||||||
unlimited_balance: form.value.unlimited_balance,
|
unlimited_balance: form.value.unlimited_balance,
|
||||||
expires_at: form.value.expires_at,
|
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,
|
auto_delete_on_expiry: form.value.auto_delete_on_expiry,
|
||||||
allowed_providers: form.value.provider_unrestricted ? null : [...form.value.allowed_providers],
|
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],
|
allowed_api_formats: form.value.api_format_unrestricted ? null : [...form.value.allowed_api_formats],
|
||||||
|
|||||||
@@ -256,6 +256,36 @@
|
|||||||
</div>
|
</div>
|
||||||
</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">
|
<div class="space-y-2">
|
||||||
<Label class="text-sm font-medium">额度</Label>
|
<Label class="text-sm font-medium">额度</Label>
|
||||||
@@ -352,6 +382,7 @@ export interface UserFormData {
|
|||||||
allowed_providers?: string[] | null
|
allowed_providers?: string[] | null
|
||||||
allowed_api_formats?: string[] | null
|
allowed_api_formats?: string[] | null
|
||||||
allowed_models?: string[] | null
|
allowed_models?: string[] | null
|
||||||
|
rate_limit?: number | null
|
||||||
}
|
}
|
||||||
|
|
||||||
const props = defineProps<{
|
const props = defineProps<{
|
||||||
@@ -406,9 +437,11 @@ const form = ref({
|
|||||||
provider_unrestricted: true,
|
provider_unrestricted: true,
|
||||||
api_format_unrestricted: true,
|
api_format_unrestricted: true,
|
||||||
model_unrestricted: true,
|
model_unrestricted: true,
|
||||||
|
rate_limit_inherited: true,
|
||||||
allowed_providers: [] as string[],
|
allowed_providers: [] as string[],
|
||||||
allowed_api_formats: [] as string[],
|
allowed_api_formats: [] as string[],
|
||||||
allowed_models: [] as string[],
|
allowed_models: [] as string[],
|
||||||
|
rate_limit: undefined as number | undefined,
|
||||||
})
|
})
|
||||||
|
|
||||||
function createFieldNonce(): string {
|
function createFieldNonce(): string {
|
||||||
@@ -429,9 +462,11 @@ function resetForm() {
|
|||||||
provider_unrestricted: true,
|
provider_unrestricted: true,
|
||||||
api_format_unrestricted: true,
|
api_format_unrestricted: true,
|
||||||
model_unrestricted: true,
|
model_unrestricted: true,
|
||||||
|
rate_limit_inherited: true,
|
||||||
allowed_providers: [],
|
allowed_providers: [],
|
||||||
allowed_api_formats: [],
|
allowed_api_formats: [],
|
||||||
allowed_models: [],
|
allowed_models: [],
|
||||||
|
rate_limit: undefined,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -451,9 +486,11 @@ function loadUserData() {
|
|||||||
provider_unrestricted: props.user.allowed_providers == null,
|
provider_unrestricted: props.user.allowed_providers == null,
|
||||||
api_format_unrestricted: props.user.allowed_api_formats == null,
|
api_format_unrestricted: props.user.allowed_api_formats == null,
|
||||||
model_unrestricted: props.user.allowed_models == 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_providers: props.user.allowed_providers ? [...props.user.allowed_providers] : [],
|
||||||
allowed_api_formats: props.user.allowed_api_formats ? [...props.user.allowed_api_formats] : [],
|
allowed_api_formats: props.user.allowed_api_formats ? [...props.user.allowed_api_formats] : [],
|
||||||
allowed_models: props.user.allowed_models ? [...props.user.allowed_models] : [],
|
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
|
allowed_models: form.value.model_unrestricted
|
||||||
? null
|
? null
|
||||||
: [...form.value.allowed_models],
|
: [...form.value.allowed_models],
|
||||||
|
rate_limit: form.value.rate_limit_inherited ? null : (form.value.rate_limit ?? 0),
|
||||||
}
|
}
|
||||||
|
|
||||||
if (isEditMode.value && props.user?.id) {
|
if (isEditMode.value && props.user?.id) {
|
||||||
|
|||||||
@@ -1,6 +1,13 @@
|
|||||||
import { defineStore } from 'pinia'
|
import { defineStore } from 'pinia'
|
||||||
import { ref } from 'vue'
|
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'
|
import { parseApiError } from '@/utils/errorParser'
|
||||||
|
|
||||||
export const useUsersStore = defineStore('users', () => {
|
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 {
|
try {
|
||||||
return await usersApi.createApiKey(userId, name)
|
return await usersApi.createApiKey(userId, data)
|
||||||
} catch (err: unknown) {
|
} catch (err: unknown) {
|
||||||
error.value = parseApiError(err, '创建 API Key 失败')
|
error.value = parseApiError(err, '创建 API Key 失败')
|
||||||
throw err
|
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) {
|
async function deleteApiKey(userId: string, keyId: string) {
|
||||||
try {
|
try {
|
||||||
await usersApi.deleteApiKey(userId, keyId)
|
await usersApi.deleteApiKey(userId, keyId)
|
||||||
@@ -121,6 +141,7 @@ export const useUsersStore = defineStore('users', () => {
|
|||||||
deleteUser,
|
deleteUser,
|
||||||
getUserApiKeys,
|
getUserApiKeys,
|
||||||
createApiKey,
|
createApiKey,
|
||||||
|
updateApiKey,
|
||||||
deleteApiKey,
|
deleteApiKey,
|
||||||
getFullApiKey
|
getFullApiKey
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -159,3 +159,25 @@ export function formatHitRate(rate: number | undefined): string {
|
|||||||
if (typeof rate !== 'number' || Number.isNaN(rate)) return '-'
|
if (typeof rate !== 'number' || Number.isNaN(rate)) return '-'
|
||||||
return `${rate.toFixed(2)}%`
|
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
|
||||||
|
}
|
||||||
|
|||||||
@@ -105,8 +105,8 @@
|
|||||||
<TableHead class="w-[240px] h-12 font-semibold">
|
<TableHead class="w-[240px] h-12 font-semibold">
|
||||||
钱包
|
钱包
|
||||||
</TableHead>
|
</TableHead>
|
||||||
<TableHead class="w-[130px] h-12 font-semibold">
|
<TableHead class="w-[190px] h-12 font-semibold">
|
||||||
使用统计
|
统计/限速
|
||||||
</TableHead>
|
</TableHead>
|
||||||
<TableHead class="w-[110px] h-12 font-semibold">
|
<TableHead class="w-[110px] h-12 font-semibold">
|
||||||
有效期
|
有效期
|
||||||
@@ -221,7 +221,23 @@
|
|||||||
请求: <span class="font-medium text-foreground">{{ (apiKey.total_requests || 0).toLocaleString() }}</span>
|
请求: <span class="font-medium text-foreground">{{ (apiKey.total_requests || 0).toLocaleString() }}</span>
|
||||||
</div>
|
</div>
|
||||||
<div class="text-muted-foreground">
|
<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>
|
||||||
</div>
|
</div>
|
||||||
</TableCell>
|
</TableCell>
|
||||||
@@ -387,6 +403,12 @@
|
|||||||
>
|
>
|
||||||
{{ walletStatusLabel(getApiKeyWalletStatus(apiKey.id)) }}
|
{{ walletStatusLabel(getApiKeyWalletStatus(apiKey.id)) }}
|
||||||
</Badge>
|
</Badge>
|
||||||
|
<Badge
|
||||||
|
variant="secondary"
|
||||||
|
class="h-5 px-1.5 py-0 text-[10px] font-medium"
|
||||||
|
>
|
||||||
|
{{ formatRateLimitInheritable(apiKey.rate_limit) }}
|
||||||
|
</Badge>
|
||||||
<Badge
|
<Badge
|
||||||
v-if="apiKey.auto_delete_on_expiry"
|
v-if="apiKey.auto_delete_on_expiry"
|
||||||
variant="secondary"
|
variant="secondary"
|
||||||
@@ -429,14 +451,6 @@
|
|||||||
</div>
|
</div>
|
||||||
|
|
||||||
<div class="grid grid-cols-2 gap-2.5 text-xs">
|
<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="rounded-lg border border-border/50 bg-background/70 p-2.5">
|
||||||
<div class="mb-1 text-muted-foreground">
|
<div class="mb-1 text-muted-foreground">
|
||||||
请求次数
|
请求次数
|
||||||
@@ -445,6 +459,14 @@
|
|||||||
{{ (apiKey.total_requests || 0).toLocaleString() }}
|
{{ (apiKey.total_requests || 0).toLocaleString() }}
|
||||||
</div>
|
</div>
|
||||||
</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="col-span-2 rounded-lg border border-border/50 bg-background/70 p-2.5">
|
||||||
<div class="mb-1 text-muted-foreground">
|
<div class="mb-1 text-muted-foreground">
|
||||||
有效期
|
有效期
|
||||||
@@ -663,6 +685,7 @@ import {
|
|||||||
|
|
||||||
import { StandaloneKeyFormDialog, type StandaloneKeyFormData } from '@/features/api-keys'
|
import { StandaloneKeyFormDialog, type StandaloneKeyFormData } from '@/features/api-keys'
|
||||||
import { parseApiError } from '@/utils/errorParser'
|
import { parseApiError } from '@/utils/errorParser'
|
||||||
|
import { formatTokens, formatRateLimitInheritable, isRateLimitInherited, isRateLimitUnlimited } from '@/utils/format'
|
||||||
import { log } from '@/utils/logger'
|
import { log } from '@/utils/logger'
|
||||||
|
|
||||||
const { success, error } = useToast()
|
const { success, error } = useToast()
|
||||||
@@ -1028,7 +1051,7 @@ async function handleKeyFormSubmit(data: StandaloneKeyFormData) {
|
|||||||
const updateData: Partial<CreateStandaloneApiKeyRequest> = {
|
const updateData: Partial<CreateStandaloneApiKeyRequest> = {
|
||||||
name: data.name || undefined,
|
name: data.name || undefined,
|
||||||
unlimited_balance: Boolean(data.unlimited_balance),
|
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/空 = 永不过期
|
expires_at: data.expires_at || null, // undefined/空 = 永不过期
|
||||||
auto_delete_on_expiry: data.auto_delete_on_expiry,
|
auto_delete_on_expiry: data.auto_delete_on_expiry,
|
||||||
// 空数组表示清除限制(允许全部),后端会将空数组存为 NULL
|
// 空数组表示清除限制(允许全部),后端会将空数组存为 NULL
|
||||||
@@ -1057,7 +1080,7 @@ async function handleKeyFormSubmit(data: StandaloneKeyFormData) {
|
|||||||
const createData: CreateStandaloneApiKeyRequest = {
|
const createData: CreateStandaloneApiKeyRequest = {
|
||||||
name: data.name || undefined,
|
name: data.name || undefined,
|
||||||
initial_balance_usd: isUnlimited ? null : (data.initial_balance_usd as number),
|
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/空 = 永不过期
|
expires_at: data.expires_at || null, // undefined/空 = 永不过期
|
||||||
auto_delete_on_expiry: data.auto_delete_on_expiry,
|
auto_delete_on_expiry: data.auto_delete_on_expiry,
|
||||||
// 空数组表示不设置限制(允许全部),后端会将空数组存为 NULL
|
// 空数组表示不设置限制(允许全部),后端会将空数组存为 NULL
|
||||||
|
|||||||
@@ -180,7 +180,7 @@
|
|||||||
钱包
|
钱包
|
||||||
</TableHead>
|
</TableHead>
|
||||||
<TableHead class="w-[170px] h-12 font-semibold">
|
<TableHead class="w-[170px] h-12 font-semibold">
|
||||||
使用统计
|
统计/限速
|
||||||
</TableHead>
|
</TableHead>
|
||||||
<TableHead class="w-[110px] h-12 font-semibold">
|
<TableHead class="w-[110px] h-12 font-semibold">
|
||||||
创建时间
|
创建时间
|
||||||
@@ -258,26 +258,42 @@
|
|||||||
</div>
|
</div>
|
||||||
</TableCell>
|
</TableCell>
|
||||||
<TableCell class="py-4">
|
<TableCell class="py-4">
|
||||||
<div
|
<div class="space-y-1 text-xs">
|
||||||
v-if="userStats[user.id]"
|
<template v-if="userStats[user.id]">
|
||||||
class="space-y-1 text-xs"
|
<div class="flex items-center text-muted-foreground">
|
||||||
>
|
<span class="w-14">请求:</span>
|
||||||
<div class="flex items-center text-muted-foreground">
|
<span class="font-medium text-foreground">{{ formatNumber(userStats[user.id]?.request_count) }}</span>
|
||||||
<span class="w-14">请求:</span>
|
</div>
|
||||||
<span class="font-medium text-foreground">{{ formatNumber(userStats[user.id]?.request_count) }}</span>
|
<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>
|
||||||
<div class="flex items-center text-muted-foreground">
|
<div class="flex items-center text-muted-foreground">
|
||||||
<span class="w-14">Tokens:</span>
|
<span class="w-14">限速:</span>
|
||||||
<span class="font-medium text-foreground">{{ formatTokens(userStats[user.id]?.total_tokens ?? 0) }}</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>
|
</div>
|
||||||
<div
|
|
||||||
v-else
|
|
||||||
class="text-xs text-muted-foreground"
|
|
||||||
>
|
|
||||||
<span v-if="loadingStats">加载中...</span>
|
|
||||||
<span v-else>无数据</span>
|
|
||||||
</div>
|
|
||||||
</TableCell>
|
</TableCell>
|
||||||
<TableCell class="py-4 text-xs text-muted-foreground">
|
<TableCell class="py-4 text-xs text-muted-foreground">
|
||||||
{{ formatDate(user.created_at) }}
|
{{ formatDate(user.created_at) }}
|
||||||
@@ -436,6 +452,12 @@
|
|||||||
>
|
>
|
||||||
{{ walletStatusLabel(getUserWalletStatus(user.id)) }}
|
{{ walletStatusLabel(getUserWalletStatus(user.id)) }}
|
||||||
</Badge>
|
</Badge>
|
||||||
|
<Badge
|
||||||
|
variant="secondary"
|
||||||
|
class="h-5 px-1.5 py-0 text-[10px] font-medium"
|
||||||
|
>
|
||||||
|
{{ formatRateLimitInheritable(user.rate_limit) }}
|
||||||
|
</Badge>
|
||||||
</div>
|
</div>
|
||||||
|
|
||||||
<div class="rounded-xl border border-border/60 bg-muted/40 p-3.5">
|
<div class="rounded-xl border border-border/60 bg-muted/40 p-3.5">
|
||||||
@@ -633,6 +655,12 @@
|
|||||||
>
|
>
|
||||||
独立余额
|
独立余额
|
||||||
</Badge>
|
</Badge>
|
||||||
|
<Badge
|
||||||
|
variant="secondary"
|
||||||
|
class="text-xs"
|
||||||
|
>
|
||||||
|
{{ formatRateLimitSimple(apiKey.rate_limit) }}
|
||||||
|
</Badge>
|
||||||
</div>
|
</div>
|
||||||
<div class="flex items-center gap-1 mt-0.5">
|
<div class="flex items-center gap-1 mt-0.5">
|
||||||
<code class="text-xs font-mono text-muted-foreground">
|
<code class="text-xs font-mono text-muted-foreground">
|
||||||
@@ -658,6 +686,15 @@
|
|||||||
${{ (apiKey.total_cost_usd || 0).toFixed(4) }}
|
${{ (apiKey.total_cost_usd || 0).toFixed(4) }}
|
||||||
</div>
|
</div>
|
||||||
</div>
|
</div>
|
||||||
|
<Button
|
||||||
|
variant="ghost"
|
||||||
|
size="icon"
|
||||||
|
class="h-8 w-8"
|
||||||
|
title="编辑"
|
||||||
|
@click="openEditUserApiKeyDialog(apiKey)"
|
||||||
|
>
|
||||||
|
<SquarePen class="h-4 w-4" />
|
||||||
|
</Button>
|
||||||
<Button
|
<Button
|
||||||
variant="ghost"
|
variant="ghost"
|
||||||
size="icon"
|
size="icon"
|
||||||
@@ -718,13 +755,87 @@
|
|||||||
<Button
|
<Button
|
||||||
class="h-10 px-5"
|
class="h-10 px-5"
|
||||||
:disabled="creatingApiKey"
|
:disabled="creatingApiKey"
|
||||||
@click="createApiKey"
|
@click="openCreateUserApiKeyDialog"
|
||||||
>
|
>
|
||||||
{{ creatingApiKey ? '创建中...' : '创建' }}
|
{{ creatingApiKey ? '创建中...' : '创建' }}
|
||||||
</Button>
|
</Button>
|
||||||
</template>
|
</template>
|
||||||
</Dialog>
|
</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
|
<WalletOpsDrawer
|
||||||
:open="showWalletActionDialogState"
|
:open="showWalletActionDialogState"
|
||||||
:wallet="walletActionTarget?.wallet || null"
|
:wallet="walletActionTarget?.wallet || null"
|
||||||
@@ -849,6 +960,8 @@ import {
|
|||||||
import UserFormDialog, { type UserFormData } from '@/features/users/components/UserFormDialog.vue'
|
import UserFormDialog, { type UserFormData } from '@/features/users/components/UserFormDialog.vue'
|
||||||
import WalletOpsDrawer from '@/features/wallet/components/WalletOpsDrawer.vue'
|
import WalletOpsDrawer from '@/features/wallet/components/WalletOpsDrawer.vue'
|
||||||
import { parseApiError } from '@/utils/errorParser'
|
import { parseApiError } from '@/utils/errorParser'
|
||||||
|
import { formatTokens, formatRateLimitInheritable, formatRateLimitSimple, isRateLimitInherited, isRateLimitUnlimited } from '@/utils/format'
|
||||||
|
import { parseNumberInput } from '@/utils/form'
|
||||||
import { log } from '@/utils/logger'
|
import { log } from '@/utils/logger'
|
||||||
|
|
||||||
const { success, error } = useToast()
|
const { success, error } = useToast()
|
||||||
@@ -864,11 +977,17 @@ const userFormDialogRef = ref<InstanceType<typeof UserFormDialog>>()
|
|||||||
// API Keys 对话框状态
|
// API Keys 对话框状态
|
||||||
const showApiKeysDialog = ref(false)
|
const showApiKeysDialog = ref(false)
|
||||||
const showNewApiKeyDialog = ref(false)
|
const showNewApiKeyDialog = ref(false)
|
||||||
|
const showUserApiKeyFormDialog = ref(false)
|
||||||
const selectedUser = ref<User | null>(null)
|
const selectedUser = ref<User | null>(null)
|
||||||
const userApiKeys = ref<ApiKey[]>([])
|
const userApiKeys = ref<ApiKey[]>([])
|
||||||
const newApiKey = ref('')
|
const newApiKey = ref('')
|
||||||
const creatingApiKey = ref(false)
|
const creatingApiKey = ref(false)
|
||||||
const apiKeyInput = ref<HTMLInputElement>()
|
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>>({})
|
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 {
|
function formatNumber(value?: number | null): string {
|
||||||
const numericValue = typeof value === 'number' && Number.isFinite(value) ? value : 0
|
const numericValue = typeof value === 'number' && Number.isFinite(value) ? value : 0
|
||||||
return numericValue.toLocaleString()
|
return numericValue.toLocaleString()
|
||||||
@@ -1071,7 +1181,8 @@ function editUser(user: User) {
|
|||||||
is_active: user.is_active,
|
is_active: user.is_active,
|
||||||
allowed_providers: user.allowed_providers == null ? null : [...user.allowed_providers],
|
allowed_providers: user.allowed_providers == null ? null : [...user.allowed_providers],
|
||||||
allowed_api_formats: user.allowed_api_formats == null ? null : [...user.allowed_api_formats],
|
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
|
showUserFormDialog.value = true
|
||||||
}
|
}
|
||||||
@@ -1093,7 +1204,8 @@ async function handleUserFormSubmit(data: UserFormData & { password?: string; un
|
|||||||
role: data.role,
|
role: data.role,
|
||||||
allowed_providers: data.allowed_providers,
|
allowed_providers: data.allowed_providers,
|
||||||
allowed_api_formats: data.allowed_api_formats,
|
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) {
|
if (data.password) {
|
||||||
updateData.password = data.password
|
updateData.password = data.password
|
||||||
@@ -1112,7 +1224,8 @@ async function handleUserFormSubmit(data: UserFormData & { password?: string; un
|
|||||||
role: data.role,
|
role: data.role,
|
||||||
allowed_providers: data.allowed_providers,
|
allowed_providers: data.allowed_providers,
|
||||||
allowed_api_formats: data.allowed_api_formats,
|
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) {
|
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 (!selectedUser.value) return
|
||||||
|
if (!userApiKeyForm.value.name.trim()) {
|
||||||
|
error('请输入密钥名称', editingUserApiKey.value ? '更新 API Key 失败' : '创建 API Key 失败')
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
creatingApiKey.value = true
|
creatingApiKey.value = true
|
||||||
try {
|
try {
|
||||||
const response = await usersStore.createApiKey(
|
if (editingUserApiKey.value) {
|
||||||
selectedUser.value.id,
|
await usersStore.updateApiKey(selectedUser.value.id, editingUserApiKey.value.id, {
|
||||||
`Key-${new Date().toISOString().split('T')[0]}`
|
name: userApiKeyForm.value.name,
|
||||||
)
|
rate_limit: userApiKeyForm.value.rate_limit ?? 0,
|
||||||
newApiKey.value = response.key || ''
|
})
|
||||||
showNewApiKeyDialog.value = true
|
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)
|
await loadUserApiKeys(selectedUser.value.id)
|
||||||
|
closeUserApiKeyFormDialog()
|
||||||
} catch (err: unknown) {
|
} catch (err: unknown) {
|
||||||
error(parseApiError(err, '未知错误'), '创建 API Key 失败')
|
error(parseApiError(err, '未知错误'), editingUserApiKey.value ? '更新 API Key 失败' : '创建 API Key 失败')
|
||||||
} finally {
|
} finally {
|
||||||
creatingApiKey.value = false
|
creatingApiKey.value = false
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -39,7 +39,7 @@
|
|||||||
for="rate-limit"
|
for="rate-limit"
|
||||||
class="block text-sm font-medium"
|
class="block text-sm font-medium"
|
||||||
>
|
>
|
||||||
每分钟请求限制
|
默认速率限制 (请求/分钟)
|
||||||
</Label>
|
</Label>
|
||||||
<Input
|
<Input
|
||||||
id="rate-limit"
|
id="rate-limit"
|
||||||
@@ -50,7 +50,7 @@
|
|||||||
@update:model-value="$emit('update:rateLimitPerMinute', Number($event))"
|
@update:model-value="$emit('update:rateLimitPerMinute', Number($event))"
|
||||||
/>
|
/>
|
||||||
<p class="mt-1 text-xs text-muted-foreground">
|
<p class="mt-1 text-xs text-muted-foreground">
|
||||||
0 表示不限制
|
0 表示默认不限制;未单独配置的用户和独立 Key 会跟随这里
|
||||||
</p>
|
</p>
|
||||||
</div>
|
</div>
|
||||||
|
|
||||||
|
|||||||
@@ -20,7 +20,7 @@
|
|||||||
size="icon"
|
size="icon"
|
||||||
class="h-8 w-8"
|
class="h-8 w-8"
|
||||||
title="创建新 API Key"
|
title="创建新 API Key"
|
||||||
@click="showCreateDialog = true"
|
@click="openCreateApiKeyDialog"
|
||||||
>
|
>
|
||||||
<Plus class="w-3.5 h-3.5" />
|
<Plus class="w-3.5 h-3.5" />
|
||||||
</Button>
|
</Button>
|
||||||
@@ -56,7 +56,7 @@
|
|||||||
<Button
|
<Button
|
||||||
size="lg"
|
size="lg"
|
||||||
class="shadow-lg shadow-primary/20"
|
class="shadow-lg shadow-primary/20"
|
||||||
@click="showCreateDialog = true"
|
@click="openCreateApiKeyDialog"
|
||||||
>
|
>
|
||||||
<Plus class="mr-2 h-4 w-4" />
|
<Plus class="mr-2 h-4 w-4" />
|
||||||
创建新 API Key
|
创建新 API Key
|
||||||
@@ -157,17 +157,23 @@
|
|||||||
<div class="flex flex-col items-center gap-1">
|
<div class="flex flex-col items-center gap-1">
|
||||||
<Badge
|
<Badge
|
||||||
:variant="apiKey.is_active ? 'success' : 'secondary'"
|
: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 ? '活跃' : '禁用' }}
|
{{ apiKey.is_active ? '活跃' : '禁用' }}
|
||||||
</Badge>
|
</Badge>
|
||||||
<Badge
|
<Badge
|
||||||
v-if="apiKey.is_locked"
|
v-if="apiKey.is_locked"
|
||||||
variant="warning"
|
variant="warning"
|
||||||
class="font-medium text-[10px]"
|
class="h-5 px-2 py-0 text-[10px] font-medium"
|
||||||
>
|
>
|
||||||
已锁定
|
已锁定
|
||||||
</Badge>
|
</Badge>
|
||||||
|
<Badge
|
||||||
|
variant="secondary"
|
||||||
|
class="h-5 px-2 py-0 text-[10px] font-medium"
|
||||||
|
>
|
||||||
|
{{ formatRateLimitSimple(apiKey.rate_limit) }}
|
||||||
|
</Badge>
|
||||||
</div>
|
</div>
|
||||||
</TableCell>
|
</TableCell>
|
||||||
|
|
||||||
@@ -179,6 +185,16 @@
|
|||||||
<!-- 操作按钮 -->
|
<!-- 操作按钮 -->
|
||||||
<TableCell class="py-4">
|
<TableCell class="py-4">
|
||||||
<div class="flex justify-center gap-1">
|
<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
|
<Button
|
||||||
variant="ghost"
|
variant="ghost"
|
||||||
size="icon"
|
size="icon"
|
||||||
@@ -237,8 +253,24 @@
|
|||||||
>
|
>
|
||||||
已锁定
|
已锁定
|
||||||
</Badge>
|
</Badge>
|
||||||
|
<Badge
|
||||||
|
variant="secondary"
|
||||||
|
class="text-[10px] px-1.5 py-0"
|
||||||
|
>
|
||||||
|
{{ formatRateLimitSimple(apiKey.rate_limit) }}
|
||||||
|
</Badge>
|
||||||
</div>
|
</div>
|
||||||
<div class="flex items-center gap-0.5 flex-shrink-0">
|
<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
|
<Button
|
||||||
variant="ghost"
|
variant="ghost"
|
||||||
size="icon"
|
size="icon"
|
||||||
@@ -288,6 +320,10 @@
|
|||||||
<span class="text-foreground font-medium">
|
<span class="text-foreground font-medium">
|
||||||
{{ formatNumber(apiKey.total_requests || 0) }} 次
|
{{ formatNumber(apiKey.total_requests || 0) }} 次
|
||||||
</span>
|
</span>
|
||||||
|
<span class="text-muted-foreground">•</span>
|
||||||
|
<span class="text-muted-foreground">
|
||||||
|
{{ formatRateLimitSimple(apiKey.rate_limit) }}
|
||||||
|
</span>
|
||||||
</div>
|
</div>
|
||||||
</div>
|
</div>
|
||||||
</div>
|
</div>
|
||||||
@@ -316,10 +352,10 @@
|
|||||||
</div>
|
</div>
|
||||||
<div class="flex-1 min-w-0">
|
<div class="flex-1 min-w-0">
|
||||||
<h3 class="text-lg font-semibold text-foreground leading-tight">
|
<h3 class="text-lg font-semibold text-foreground leading-tight">
|
||||||
创建 API 密钥
|
{{ editingApiKey ? '编辑 API 密钥' : '创建 API 密钥' }}
|
||||||
</h3>
|
</h3>
|
||||||
<p class="text-xs text-muted-foreground">
|
<p class="text-xs text-muted-foreground">
|
||||||
创建一个新的密钥用于访问 API 服务
|
{{ editingApiKey ? '更新密钥名称和速率限制' : '创建一个新的密钥用于访问 API 服务' }}
|
||||||
</p>
|
</p>
|
||||||
</div>
|
</div>
|
||||||
</div>
|
</div>
|
||||||
@@ -344,26 +380,46 @@
|
|||||||
给密钥起一个有意义的名称方便识别
|
给密钥起一个有意义的名称方便识别
|
||||||
</p>
|
</p>
|
||||||
</div>
|
</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>
|
</div>
|
||||||
|
|
||||||
<template #footer>
|
<template #footer>
|
||||||
<Button
|
<Button
|
||||||
variant="outline"
|
variant="outline"
|
||||||
class="h-11 px-6"
|
class="h-11 px-6"
|
||||||
@click="showCreateDialog = false"
|
@click="closeApiKeyDialog"
|
||||||
>
|
>
|
||||||
取消
|
取消
|
||||||
</Button>
|
</Button>
|
||||||
<Button
|
<Button
|
||||||
class="h-11 px-6 shadow-lg shadow-primary/20"
|
class="h-11 px-6 shadow-lg shadow-primary/20"
|
||||||
:disabled="creating"
|
:disabled="creating"
|
||||||
@click="createApiKey"
|
@click="saveApiKey"
|
||||||
>
|
>
|
||||||
<Loader2
|
<Loader2
|
||||||
v-if="creating"
|
v-if="creating"
|
||||||
class="animate-spin h-4 w-4 mr-2"
|
class="animate-spin h-4 w-4 mr-2"
|
||||||
/>
|
/>
|
||||||
{{ creating ? '创建中...' : '创建' }}
|
{{ creating ? (editingApiKey ? '保存中...' : '创建中...') : (editingApiKey ? '保存' : '创建') }}
|
||||||
</Button>
|
</Button>
|
||||||
</template>
|
</template>
|
||||||
</Dialog>
|
</Dialog>
|
||||||
@@ -455,10 +511,12 @@ import {
|
|||||||
TableRow
|
TableRow
|
||||||
} from '@/components/ui'
|
} from '@/components/ui'
|
||||||
import RefreshButton from '@/components/ui/refresh-button.vue'
|
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 { useToast } from '@/composables/useToast'
|
||||||
import { log } from '@/utils/logger'
|
import { log } from '@/utils/logger'
|
||||||
import { parseApiError } from '@/utils/errorParser'
|
import { parseApiError } from '@/utils/errorParser'
|
||||||
|
import { formatRateLimitSimple } from '@/utils/format'
|
||||||
|
import { parseNumberInput } from '@/utils/form'
|
||||||
import { getErrorStatus } from '@/types/api-error'
|
import { getErrorStatus } from '@/types/api-error'
|
||||||
import { computed } from 'vue'
|
import { computed } from 'vue'
|
||||||
|
|
||||||
@@ -483,8 +541,10 @@ const showKeyDialog = ref(false)
|
|||||||
const showDeleteDialog = ref(false)
|
const showDeleteDialog = ref(false)
|
||||||
|
|
||||||
const newKeyName = ref('')
|
const newKeyName = ref('')
|
||||||
|
const newKeyRateLimit = ref<number | undefined>(undefined)
|
||||||
const newKeyValue = ref('')
|
const newKeyValue = ref('')
|
||||||
const keyToDelete = ref<ApiKey | null>(null)
|
const keyToDelete = ref<ApiKey | null>(null)
|
||||||
|
const editingApiKey = ref<ApiKey | null>(null)
|
||||||
|
|
||||||
onMounted(() => {
|
onMounted(() => {
|
||||||
loadApiKeys()
|
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()) {
|
if (!newKeyName.value.trim()) {
|
||||||
showError('请输入密钥名称')
|
showError('请输入密钥名称')
|
||||||
return
|
return
|
||||||
@@ -517,16 +598,26 @@ async function createApiKey() {
|
|||||||
|
|
||||||
creating.value = true
|
creating.value = true
|
||||||
try {
|
try {
|
||||||
const newKey = await meApi.createApiKey(newKeyName.value)
|
if (editingApiKey.value) {
|
||||||
newKeyValue.value = newKey.key || ''
|
await meApi.updateApiKey(editingApiKey.value.id, {
|
||||||
showCreateDialog.value = false
|
name: newKeyName.value,
|
||||||
showKeyDialog.value = true
|
rate_limit: newKeyRateLimit.value ?? 0,
|
||||||
newKeyName.value = ''
|
})
|
||||||
|
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()
|
await loadApiKeys()
|
||||||
success('API 密钥创建成功')
|
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
log.error('创建 API 密钥失败:', error)
|
log.error(editingApiKey.value ? '更新 API 密钥失败:' : '创建 API 密钥失败:', error)
|
||||||
showError('创建 API 密钥失败')
|
showError(editingApiKey.value ? '更新 API 密钥失败' : '创建 API 密钥失败')
|
||||||
} finally {
|
} finally {
|
||||||
creating.value = false
|
creating.value = false
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -22,7 +22,7 @@ from src.core.exceptions import InvalidRequestException, NotFoundException
|
|||||||
from src.core.logger import logger
|
from src.core.logger import logger
|
||||||
from src.database import get_db, get_db_context
|
from src.database import get_db, get_db_context
|
||||||
from src.models.api import CreateApiKeyRequest
|
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.apikey import ApiKeyService
|
||||||
from src.services.user.bulk_cleanup import pre_clean_api_key
|
from src.services.user.bulk_cleanup import pre_clean_api_key
|
||||||
from src.services.wallet import WalletService
|
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()
|
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 {
|
return {
|
||||||
"id": api_key.id,
|
"id": api_key.id,
|
||||||
"user_id": api_key.user_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_active": api_key.is_active,
|
||||||
"is_standalone": api_key.is_standalone,
|
"is_standalone": api_key.is_standalone,
|
||||||
"total_requests": api_key.total_requests,
|
"total_requests": api_key.total_requests,
|
||||||
|
"total_tokens": int(total_tokens or 0),
|
||||||
"total_cost_usd": float(api_key.total_cost_usd or 0),
|
"total_cost_usd": float(api_key.total_cost_usd or 0),
|
||||||
"rate_limit": api_key.rate_limit,
|
"rate_limit": api_key.rate_limit,
|
||||||
"allowed_providers": api_key.allowed_providers,
|
"allowed_providers": api_key.allowed_providers,
|
||||||
@@ -116,8 +119,27 @@ def _list_standalone_api_keys_sync(
|
|||||||
for api_key in api_keys:
|
for api_key in api_keys:
|
||||||
db.refresh(api_key)
|
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 {
|
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,
|
"total": total,
|
||||||
"limit": limit,
|
"limit": limit,
|
||||||
"skip": skip,
|
"skip": skip,
|
||||||
@@ -385,7 +407,7 @@ async def create_standalone_api_key(
|
|||||||
- `allowed_providers`: 可选,允许使用的提供商列表
|
- `allowed_providers`: 可选,允许使用的提供商列表
|
||||||
- `allowed_api_formats`: 可选,允许使用的 API 格式列表
|
- `allowed_api_formats`: 可选,允许使用的 API 格式列表
|
||||||
- `allowed_models`: 可选,允许使用的模型列表
|
- `allowed_models`: 可选,允许使用的模型列表
|
||||||
- `rate_limit`: 可选,速率限制配置(请求数/秒)
|
- `rate_limit`: 可选,每分钟请求限制(null 表示跟随系统默认,0 表示不限制)
|
||||||
- `expire_days`: 可选,过期天数(与 expires_at 二选一)
|
- `expire_days`: 可选,过期天数(与 expires_at 二选一)
|
||||||
- `expires_at`: 可选,过期时间(ISO 格式或 YYYY-MM-DD 格式,优先级高于 expire_days)
|
- `expires_at`: 可选,过期时间(ISO 格式或 YYYY-MM-DD 格式,优先级高于 expire_days)
|
||||||
- `auto_delete_on_expiry`: 可选,过期后是否自动删除
|
- `auto_delete_on_expiry`: 可选,过期后是否自动删除
|
||||||
@@ -421,7 +443,7 @@ async def update_api_key(
|
|||||||
**请求体字段**:
|
**请求体字段**:
|
||||||
- `name`: 可选,API Key 的名称
|
- `name`: 可选,API Key 的名称
|
||||||
- `unlimited_balance`: 可选,是否无限余额(true=无限,false=有限,不修改余额数值)
|
- `unlimited_balance`: 可选,是否无限余额(true=无限,false=有限,不修改余额数值)
|
||||||
- `rate_limit`: 可选,速率限制配置(null 表示无限制)
|
- `rate_limit`: 可选,每分钟请求限制(null 表示跟随系统默认,0 表示不限制)
|
||||||
- `allowed_providers`: 可选,允许使用的提供商列表
|
- `allowed_providers`: 可选,允许使用的提供商列表
|
||||||
- `allowed_api_formats`: 可选,允许使用的 API 格式列表
|
- `allowed_api_formats`: 可选,允许使用的 API 格式列表
|
||||||
- `allowed_models`: 可选,允许使用的模型列表
|
- `allowed_models`: 可选,允许使用的模型列表
|
||||||
@@ -683,6 +705,14 @@ class AdminGetKeyDetailAdapter(AdminApiAdapter):
|
|||||||
"is_active": api_key.is_active,
|
"is_active": api_key.is_active,
|
||||||
"is_standalone": api_key.is_standalone,
|
"is_standalone": api_key.is_standalone,
|
||||||
"total_requests": api_key.total_requests,
|
"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),
|
"total_cost_usd": float(api_key.total_cost_usd or 0),
|
||||||
"rate_limit": api_key.rate_limit,
|
"rate_limit": api_key.rate_limit,
|
||||||
"allowed_providers": api_key.allowed_providers,
|
"allowed_providers": api_key.allowed_providers,
|
||||||
|
|||||||
@@ -2252,6 +2252,7 @@ class AdminExportUsersAdapter(AdminApiAdapter):
|
|||||||
"allowed_providers": user.allowed_providers,
|
"allowed_providers": user.allowed_providers,
|
||||||
"allowed_api_formats": user.allowed_api_formats,
|
"allowed_api_formats": user.allowed_api_formats,
|
||||||
"allowed_models": user.allowed_models,
|
"allowed_models": user.allowed_models,
|
||||||
|
"rate_limit": user.rate_limit,
|
||||||
"model_capability_settings": user.model_capability_settings,
|
"model_capability_settings": user.model_capability_settings,
|
||||||
"unlimited": wallet_service.is_unlimited_wallet(wallet),
|
"unlimited": wallet_service.is_unlimited_wallet(wallet),
|
||||||
"wallet": (wallet_service.serialize_wallet_summary(wallet) if wallet else None),
|
"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]
|
standalone_keys_data = [self._serialize_api_key(key, db=db) for key in standalone_keys]
|
||||||
|
|
||||||
return {
|
return {
|
||||||
"version": "1.2",
|
"version": "1.3",
|
||||||
"exported_at": datetime.now(timezone.utc).isoformat(),
|
"exported_at": datetime.now(timezone.utc).isoformat(),
|
||||||
"users": users_data,
|
"users": users_data,
|
||||||
"standalone_keys": standalone_keys_data,
|
"standalone_keys": standalone_keys_data,
|
||||||
@@ -2273,6 +2274,17 @@ class AdminExportUsersAdapter(AdminApiAdapter):
|
|||||||
|
|
||||||
|
|
||||||
class AdminImportUsersAdapter(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
|
@staticmethod
|
||||||
def _resolve_api_key_material(key_data: dict[str, Any]) -> tuple[str | None, str | None]:
|
def _resolve_api_key_material(key_data: dict[str, Any]) -> tuple[str | None, str | None]:
|
||||||
"""解析用户 API Key 导入材料,优先使用明文 key。"""
|
"""解析用户 API Key 导入材料,优先使用明文 key。"""
|
||||||
@@ -2289,6 +2301,30 @@ class AdminImportUsersAdapter(AdminApiAdapter):
|
|||||||
key_encrypted = key_data.get("key_encrypted")
|
key_encrypted = key_data.get("key_encrypted")
|
||||||
return key_hash, 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]
|
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
|
||||||
"""导入用户数据"""
|
"""导入用户数据"""
|
||||||
import uuid
|
import uuid
|
||||||
@@ -2306,6 +2342,7 @@ class AdminImportUsersAdapter(AdminApiAdapter):
|
|||||||
|
|
||||||
# 获取导入选项
|
# 获取导入选项
|
||||||
merge_mode = payload.get("merge_mode", "skip") # skip, overwrite, error
|
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", [])
|
users_data = payload.get("users", [])
|
||||||
standalone_keys_data = payload.get("standalone_keys", [])
|
standalone_keys_data = payload.get("standalone_keys", [])
|
||||||
|
|
||||||
@@ -2358,7 +2395,11 @@ class AdminImportUsersAdapter(AdminApiAdapter):
|
|||||||
allowed_providers=key_data.get("allowed_providers"),
|
allowed_providers=key_data.get("allowed_providers"),
|
||||||
allowed_api_formats=key_data.get("allowed_api_formats"),
|
allowed_api_formats=key_data.get("allowed_api_formats"),
|
||||||
allowed_models=key_data.get("allowed_models"),
|
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),
|
concurrent_limit=key_data.get("concurrent_limit", 5),
|
||||||
force_capabilities=key_data.get("force_capabilities"),
|
force_capabilities=key_data.get("force_capabilities"),
|
||||||
is_active=key_data.get("is_active", True),
|
is_active=key_data.get("is_active", True),
|
||||||
@@ -2396,6 +2437,7 @@ class AdminImportUsersAdapter(AdminApiAdapter):
|
|||||||
and wallet_payload.get("limit_mode") in {"finite", "unlimited"}
|
and wallet_payload.get("limit_mode") in {"finite", "unlimited"}
|
||||||
else ("unlimited" if user_data.get("unlimited") else "finite")
|
else ("unlimited" if user_data.get("unlimited") else "finite")
|
||||||
)
|
)
|
||||||
|
imported_user_rate_limit = self._normalize_imported_user_rate_limit(user_data)
|
||||||
|
|
||||||
if existing_user:
|
if existing_user:
|
||||||
user_id = existing_user.id
|
user_id = existing_user.id
|
||||||
@@ -2413,6 +2455,7 @@ class AdminImportUsersAdapter(AdminApiAdapter):
|
|||||||
existing_user.allowed_providers = user_data.get("allowed_providers")
|
existing_user.allowed_providers = user_data.get("allowed_providers")
|
||||||
existing_user.allowed_api_formats = user_data.get("allowed_api_formats")
|
existing_user.allowed_api_formats = user_data.get("allowed_api_formats")
|
||||||
existing_user.allowed_models = user_data.get("allowed_models")
|
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(
|
existing_user.model_capability_settings = user_data.get(
|
||||||
"model_capability_settings"
|
"model_capability_settings"
|
||||||
)
|
)
|
||||||
@@ -2449,6 +2492,7 @@ class AdminImportUsersAdapter(AdminApiAdapter):
|
|||||||
allowed_providers=user_data.get("allowed_providers"),
|
allowed_providers=user_data.get("allowed_providers"),
|
||||||
allowed_api_formats=user_data.get("allowed_api_formats"),
|
allowed_api_formats=user_data.get("allowed_api_formats"),
|
||||||
allowed_models=user_data.get("allowed_models"),
|
allowed_models=user_data.get("allowed_models"),
|
||||||
|
rate_limit=imported_user_rate_limit,
|
||||||
model_capability_settings=user_data.get("model_capability_settings"),
|
model_capability_settings=user_data.get("model_capability_settings"),
|
||||||
is_active=user_data.get("is_active", True),
|
is_active=user_data.get("is_active", True),
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -18,7 +18,7 @@ from src.core.exceptions import InvalidRequestException, NotFoundException, tran
|
|||||||
from src.core.logger import logger
|
from src.core.logger import logger
|
||||||
from src.database import get_db, get_db_context
|
from src.database import get_db, get_db_context
|
||||||
from src.models.admin_requests import UpdateUserRequest
|
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.models.database import ApiKey, User, UserRole, Wallet
|
||||||
from src.services.cache.user_cache import UserCacheService
|
from src.services.cache.user_cache import UserCacheService
|
||||||
from src.services.system.config import SystemConfigService
|
from src.services.system.config import SystemConfigService
|
||||||
@@ -57,6 +57,7 @@ def _serialize_user(
|
|||||||
"allowed_providers": user.allowed_providers,
|
"allowed_providers": user.allowed_providers,
|
||||||
"allowed_api_formats": user.allowed_api_formats,
|
"allowed_api_formats": user.allowed_api_formats,
|
||||||
"allowed_models": user.allowed_models,
|
"allowed_models": user.allowed_models,
|
||||||
|
"rate_limit": user.rate_limit,
|
||||||
"unlimited": WalletService.is_unlimited_wallet(resolved_wallet),
|
"unlimited": WalletService.is_unlimited_wallet(resolved_wallet),
|
||||||
"is_active": user.is_active,
|
"is_active": user.is_active,
|
||||||
"created_at": user.created_at.isoformat(),
|
"created_at": user.created_at.isoformat(),
|
||||||
@@ -89,6 +90,7 @@ def _create_user_sync(
|
|||||||
allowed_providers=request.allowed_providers,
|
allowed_providers=request.allowed_providers,
|
||||||
allowed_api_formats=request.allowed_api_formats,
|
allowed_api_formats=request.allowed_api_formats,
|
||||||
allowed_models=request.allowed_models,
|
allowed_models=request.allowed_models,
|
||||||
|
rate_limit=request.rate_limit,
|
||||||
)
|
)
|
||||||
return _serialize_user(db, user), {
|
return _serialize_user(db, user), {
|
||||||
"action": "create_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(
|
def _toggle_user_key_lock_sync(
|
||||||
user_id: str,
|
user_id: str,
|
||||||
key_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)
|
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")
|
@router.patch("/{user_id}/api-keys/{key_id}/lock")
|
||||||
async def toggle_user_api_key_lock(
|
async def toggle_user_api_key_lock(
|
||||||
user_id: str,
|
user_id: str,
|
||||||
@@ -700,6 +777,33 @@ class AdminDeleteUserKeyAdapter(AdminApiAdapter):
|
|||||||
return response
|
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):
|
class AdminToggleUserKeyLockAdapter(AdminApiAdapter):
|
||||||
"""切换用户普通 API Key 的锁定状态"""
|
"""切换用户普通 API Key 的锁定状态"""
|
||||||
|
|
||||||
|
|||||||
@@ -19,7 +19,9 @@ from src.core.logger import logger
|
|||||||
from src.database.database import create_session
|
from src.database.database import create_session
|
||||||
from src.models.database import ApiKey, AuditEventType, User
|
from src.models.database import ApiKey, AuditEventType, User
|
||||||
from src.services.auth.service import AuthService
|
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.audit import AuditService
|
||||||
|
from src.services.system.config import SystemConfigService
|
||||||
from src.services.usage.service import UsageService
|
from src.services.usage.service import UsageService
|
||||||
from src.services.wallet import WalletService
|
from src.services.wallet import WalletService
|
||||||
from src.utils.perf import PerfRecorder
|
from src.utils.perf import PerfRecorder
|
||||||
@@ -123,6 +125,9 @@ class ApiRequestPipeline:
|
|||||||
auth_duration = PerfRecorder.stop(auth_start, "pipeline_auth", labels=perf_labels)
|
auth_duration = PerfRecorder.stop(auth_start, "pipeline_auth", labels=perf_labels)
|
||||||
_record_perf_metric("auth_ms", auth_duration)
|
_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
|
raw_body = None
|
||||||
should_eager_read_body = http_request.method in {"POST", "PUT", "PATCH"} and getattr(
|
should_eager_read_body = http_request.method in {"POST", "PUT", "PATCH"} and getattr(
|
||||||
adapter, "eager_request_body", True
|
adapter, "eager_request_body", True
|
||||||
@@ -267,6 +272,55 @@ class ApiRequestPipeline:
|
|||||||
# Internal helpers
|
# 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(
|
async def _authenticate_client(
|
||||||
self, request: Request, db: Session, adapter: ApiAdapter, **_kw: object
|
self, request: Request, db: Session, adapter: ApiAdapter, **_kw: object
|
||||||
) -> tuple[User, ApiKey]:
|
) -> tuple[User, ApiKey]:
|
||||||
|
|||||||
@@ -16,7 +16,8 @@ from src.api.base.pipeline import get_pipeline
|
|||||||
from src.core.logger import logger
|
from src.core.logger import logger
|
||||||
from src.database import get_db
|
from src.database import get_db
|
||||||
from src.models.database import ApiKey, AuditLog
|
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"])
|
router = APIRouter(prefix="/api/monitoring", tags=["Monitoring"])
|
||||||
pipeline = get_pipeline()
|
pipeline = get_pipeline()
|
||||||
@@ -146,10 +147,6 @@ class UserRateLimitStatusAdapter(AuthenticatedApiAdapter):
|
|||||||
if not user:
|
if not user:
|
||||||
raise HTTPException(status_code=401, detail="未登录")
|
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 = (
|
api_keys = (
|
||||||
db.query(ApiKey)
|
db.query(ApiKey)
|
||||||
.filter(ApiKey.user_id == user.id, ApiKey.is_active.is_(True))
|
.filter(ApiKey.user_id == user.id, ApiKey.is_active.is_(True))
|
||||||
@@ -157,31 +154,84 @@ class UserRateLimitStatusAdapter(AuthenticatedApiAdapter):
|
|||||||
.all()
|
.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 = []
|
rate_limit_info = []
|
||||||
for key in api_keys:
|
for key in api_keys:
|
||||||
try:
|
if limiter is not None:
|
||||||
headers = rate_limiter.get_rate_limit_headers(key)
|
if key.is_standalone:
|
||||||
except Exception as exc:
|
user_limit = key.rate_limit if key.rate_limit is not None else system_default
|
||||||
logger.warning(f"无法获取Key {key.id} 的限流信息: {exc}")
|
user_scope_key = limiter.get_standalone_rpm_key(key.id)
|
||||||
headers = {}
|
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(
|
rate_limit_info.append(
|
||||||
{
|
{
|
||||||
"api_key_name": key.name or f"Key-{key.id}",
|
"api_key_name": key.name or f"Key-{key.id}",
|
||||||
"limit": headers.get("X-RateLimit-Limit"),
|
"limit": None,
|
||||||
"remaining": headers.get("X-RateLimit-Remaining"),
|
"remaining": None,
|
||||||
"reset_time": headers.get("X-RateLimit-Reset"),
|
"scope": None,
|
||||||
"window": headers.get("X-RateLimit-Window"),
|
"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}
|
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
|
|
||||||
|
|||||||
@@ -33,6 +33,7 @@ from src.models.api import (
|
|||||||
PublicGlobalModelListResponse,
|
PublicGlobalModelListResponse,
|
||||||
PublicGlobalModelResponse,
|
PublicGlobalModelResponse,
|
||||||
UpdateApiKeyProvidersRequest,
|
UpdateApiKeyProvidersRequest,
|
||||||
|
UpdateMyApiKeyRequest,
|
||||||
UpdatePreferencesRequest,
|
UpdatePreferencesRequest,
|
||||||
UpdateProfileRequest,
|
UpdateProfileRequest,
|
||||||
)
|
)
|
||||||
@@ -146,6 +147,7 @@ def _create_my_api_key_sync(user_id: str, request: CreateMyApiKeyRequest) -> dic
|
|||||||
db=db,
|
db=db,
|
||||||
user_id=user_id,
|
user_id=user_id,
|
||||||
name=request.name,
|
name=request.name,
|
||||||
|
rate_limit=request.rate_limit,
|
||||||
)
|
)
|
||||||
except ValueError as exc:
|
except ValueError as exc:
|
||||||
raise InvalidRequestException(str(exc)) from 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,
|
"name": api_key.name,
|
||||||
"key": plain_key,
|
"key": plain_key,
|
||||||
"key_display": api_key.get_display_key(),
|
"key_display": api_key.get_display_key(),
|
||||||
|
"rate_limit": api_key.rate_limit,
|
||||||
"message": "API密钥创建成功",
|
"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(
|
def _update_api_key_providers_sync(
|
||||||
user_id: str,
|
user_id: str,
|
||||||
api_key_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)
|
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}")
|
@router.patch("/api-keys/{key_id}")
|
||||||
async def toggle_my_api_key(key_id: str, request: Request, db: Session = Depends(get_db)) -> Any:
|
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(),
|
"created_at": key.created_at.isoformat(),
|
||||||
"total_requests": real_stats["total_requests"],
|
"total_requests": real_stats["total_requests"],
|
||||||
"total_cost_usd": real_stats["total_cost_usd"],
|
"total_cost_usd": real_stats["total_cost_usd"],
|
||||||
|
"rate_limit": key.rate_limit,
|
||||||
"allowed_providers": key.allowed_providers,
|
"allowed_providers": key.allowed_providers,
|
||||||
"force_capabilities": key.force_capabilities,
|
"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
|
@dataclass
|
||||||
class DeleteMyApiKeyAdapter(AuthenticatedApiAdapter):
|
class DeleteMyApiKeyAdapter(AuthenticatedApiAdapter):
|
||||||
"""删除 API 密钥的适配器"""
|
"""删除 API 密钥的适配器"""
|
||||||
|
|||||||
@@ -109,8 +109,6 @@ class Config:
|
|||||||
# 支付回调安全配置(公开回调入口必须携带该共享密钥)
|
# 支付回调安全配置(公开回调入口必须携带该共享密钥)
|
||||||
self.payment_callback_secret = os.getenv("PAYMENT_CALLBACK_SECRET", "").strip()
|
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"))
|
self.public_api_rate_limit = int(os.getenv("PUBLIC_API_RATE_LIMIT", "60"))
|
||||||
|
|
||||||
# 异常处理配置
|
# 异常处理配置
|
||||||
|
|||||||
11
src/main.py
11
src/main.py
@@ -45,6 +45,7 @@ if TYPE_CHECKING:
|
|||||||
from src.services.model.fetch_scheduler import ModelFetchScheduler
|
from src.services.model.fetch_scheduler import ModelFetchScheduler
|
||||||
from src.services.provider_keys.pool_quota_probe_scheduler import PoolQuotaProbeScheduler
|
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.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.maintenance_scheduler import MaintenanceScheduler
|
||||||
from src.services.system.scheduler import TaskScheduler
|
from src.services.system.scheduler import TaskScheduler
|
||||||
from src.services.task.polling.task_poller import TaskPollerService
|
from src.services.task.polling.task_poller import TaskPollerService
|
||||||
@@ -99,6 +100,7 @@ class LifecycleState:
|
|||||||
|
|
||||||
redis_client: Redis | None = None
|
redis_client: Redis | None = None
|
||||||
concurrency_manager: ConcurrencyManager | None = None
|
concurrency_manager: ConcurrencyManager | None = None
|
||||||
|
user_rpm_limiter: UserRpmLimiter | None = None
|
||||||
plugin_manager: PluginManager | None = None
|
plugin_manager: PluginManager | None = None
|
||||||
available_modules: list[ModuleDefinition] = field(default_factory=list)
|
available_modules: list[ModuleDefinition] = field(default_factory=list)
|
||||||
task_coordinator: StartupTaskCoordinator | None = None
|
task_coordinator: StartupTaskCoordinator | None = None
|
||||||
@@ -181,6 +183,11 @@ async def _initialize_core_infrastructure(state: LifecycleState) -> None:
|
|||||||
|
|
||||||
state.concurrency_manager = await get_concurrency_manager()
|
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("初始化批量提交器...")
|
logger.info("初始化批量提交器...")
|
||||||
from src.core.batch_committer import init_batch_committer
|
from src.core.batch_committer import init_batch_committer
|
||||||
@@ -541,6 +548,10 @@ async def _run_shutdown(state: LifecycleState) -> None:
|
|||||||
if state.concurrency_manager:
|
if state.concurrency_manager:
|
||||||
await state.concurrency_manager.close()
|
await state.concurrency_manager.close()
|
||||||
|
|
||||||
|
logger.info("关闭用户/API Key RPM 限流器...")
|
||||||
|
if state.user_rpm_limiter:
|
||||||
|
await state.user_rpm_limiter.close()
|
||||||
|
|
||||||
# 关闭全局Redis客户端
|
# 关闭全局Redis客户端
|
||||||
logger.info("关闭全局Redis客户端...")
|
logger.info("关闭全局Redis客户端...")
|
||||||
from src.clients.redis_client import close_redis_client
|
from src.clients.redis_client import close_redis_client
|
||||||
|
|||||||
@@ -8,7 +8,6 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import hashlib
|
|
||||||
import time
|
import time
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
@@ -50,7 +49,6 @@ class PluginMiddleware:
|
|||||||
self._notification_cache_expires: float = 0.0
|
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
|
self.public_api_rate_limit = config.public_api_rate_limit
|
||||||
|
|
||||||
# 完全跳过限流的路径(静态资源、文档等)
|
# 完全跳过限流的路径(静态资源、文档等)
|
||||||
@@ -69,14 +67,6 @@ class PluginMiddleware:
|
|||||||
"/api/monitoring/", # 监控端点
|
"/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:
|
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||||||
"""ASGI 入口点"""
|
"""ASGI 入口点"""
|
||||||
if scope["type"] != "http":
|
if scope["type"] != "http":
|
||||||
@@ -291,13 +281,6 @@ class PluginMiddleware:
|
|||||||
|
|
||||||
return "unknown"
|
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(
|
async def _get_rate_limit_key_and_config(
|
||||||
self, request: Request
|
self, request: Request
|
||||||
) -> tuple[str | None, int | None]:
|
) -> tuple[str | None, int | None]:
|
||||||
@@ -305,7 +288,6 @@ class PluginMiddleware:
|
|||||||
获取速率限制的key和配置
|
获取速率限制的key和配置
|
||||||
|
|
||||||
策略说明:
|
策略说明:
|
||||||
- /v1/messages, /v1/chat/completions 等 LLM API: 按 API Key 限流
|
|
||||||
- /api/public/* 端点: 使用服务器级别 IP 限制
|
- /api/public/* 端点: 使用服务器级别 IP 限制
|
||||||
- /api/admin/* 端点: 跳过(在 skip_rate_limit_paths 中跳过)
|
- /api/admin/* 端点: 跳过(在 skip_rate_limit_paths 中跳过)
|
||||||
- /api/auth/* 端点: 跳过(由路由层的 IPRateLimiter 处理)
|
- /api/auth/* 端点: 跳过(由路由层的 IPRateLimiter 处理)
|
||||||
@@ -315,30 +297,6 @@ class PluginMiddleware:
|
|||||||
"""
|
"""
|
||||||
path = request.url.path
|
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
|
# /api/public/* 端点: 使用服务器级别 IP 地址作为限制 key
|
||||||
if path.startswith("/api/public/"):
|
if path.startswith("/api/public/"):
|
||||||
client_ip = self._get_client_ip(request)
|
client_ip = self._get_client_ip(request)
|
||||||
|
|||||||
@@ -692,6 +692,11 @@ class UpdateUserRequest(BaseModel):
|
|||||||
allowed_providers: list[str] | None = Field(None, description="允许使用的提供商 ID 列表")
|
allowed_providers: list[str] | None = Field(None, description="允许使用的提供商 ID 列表")
|
||||||
allowed_api_formats: list[str] | None = Field(None, description="允许使用的 API 格式列表")
|
allowed_api_formats: list[str] | None = Field(None, description="允许使用的 API 格式列表")
|
||||||
allowed_models: list[str] | None = Field(None, description="允许使用的模型名称列表")
|
allowed_models: list[str] | None = Field(None, description="允许使用的模型名称列表")
|
||||||
|
rate_limit: int | None = Field(
|
||||||
|
None,
|
||||||
|
ge=0,
|
||||||
|
description="每分钟请求限制;null 表示继承系统默认,0 表示不限制",
|
||||||
|
)
|
||||||
|
|
||||||
@field_validator("username")
|
@field_validator("username")
|
||||||
@classmethod
|
@classmethod
|
||||||
|
|||||||
@@ -252,6 +252,11 @@ class CreateUserRequest(BaseModel):
|
|||||||
allowed_models: list[str] | None = Field(
|
allowed_models: list[str] | None = Field(
|
||||||
default=None, description="允许使用的模型名称列表,null表示无限制"
|
default=None, description="允许使用的模型名称列表,null表示无限制"
|
||||||
)
|
)
|
||||||
|
rate_limit: int | None = Field(
|
||||||
|
default=None,
|
||||||
|
ge=0,
|
||||||
|
description="每分钟请求限制;null 表示继承系统默认,0 表示不限制",
|
||||||
|
)
|
||||||
|
|
||||||
@field_validator("initial_gift_usd", mode="before")
|
@field_validator("initial_gift_usd", mode="before")
|
||||||
@classmethod
|
@classmethod
|
||||||
@@ -335,6 +340,11 @@ class UpdateUserRequest(BaseModel):
|
|||||||
allowed_providers: list[str] | None = None # 允许使用的提供商 ID 列表
|
allowed_providers: list[str] | None = None # 允许使用的提供商 ID 列表
|
||||||
allowed_api_formats: list[str] | None = None # 允许使用的 API 格式列表
|
allowed_api_formats: list[str] | None = None # 允许使用的 API 格式列表
|
||||||
allowed_models: list[str] | None = None # 允许使用的模型名称列表
|
allowed_models: list[str] | None = None # 允许使用的模型名称列表
|
||||||
|
rate_limit: int | None = Field(
|
||||||
|
default=None,
|
||||||
|
ge=0,
|
||||||
|
description="每分钟请求限制;null 表示继承系统默认,0 表示不限制",
|
||||||
|
)
|
||||||
is_active: bool | None = None
|
is_active: bool | None = None
|
||||||
|
|
||||||
@field_validator("allowed_api_formats")
|
@field_validator("allowed_api_formats")
|
||||||
@@ -351,7 +361,11 @@ class CreateApiKeyRequest(BaseModel):
|
|||||||
allowed_providers: list[str] | None = None # 允许使用的提供商 ID 列表
|
allowed_providers: list[str] | None = None # 允许使用的提供商 ID 列表
|
||||||
allowed_api_formats: list[str] | None = None # 允许使用的 API 格式列表
|
allowed_api_formats: list[str] | None = None # 允许使用的 API 格式列表
|
||||||
allowed_models: list[str] | None = None # 允许使用的模型名称列表
|
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 = 永不过期,数字 = 多少天后过期
|
expire_days: int | None = None # None = 永不过期,数字 = 多少天后过期
|
||||||
expires_at: str | None = None # ISO 日期字符串,如 "2025-12-31",优先于 expire_days
|
expires_at: str | None = None # ISO 日期字符串,如 "2025-12-31",优先于 expire_days
|
||||||
initial_balance_usd: float | None = Field(
|
initial_balance_usd: float | None = Field(
|
||||||
@@ -382,6 +396,7 @@ class UserResponse(BaseModel):
|
|||||||
allowed_providers: list[str] | None = None # 允许使用的提供商 ID 列表
|
allowed_providers: list[str] | None = None # 允许使用的提供商 ID 列表
|
||||||
allowed_api_formats: list[str] | None = None # 允许使用的 API 格式列表
|
allowed_api_formats: list[str] | None = None # 允许使用的 API 格式列表
|
||||||
allowed_models: list[str] | None = None # 允许使用的模型名称列表
|
allowed_models: list[str] | None = None # 允许使用的模型名称列表
|
||||||
|
rate_limit: int | None = None
|
||||||
unlimited: bool = False
|
unlimited: bool = False
|
||||||
is_active: bool
|
is_active: bool
|
||||||
created_at: datetime
|
created_at: datetime
|
||||||
@@ -402,7 +417,7 @@ class ApiKeyResponse(BaseModel):
|
|||||||
total_cost_usd: float
|
total_cost_usd: float
|
||||||
allowed_providers: list[str] | None
|
allowed_providers: list[str] | None
|
||||||
allowed_models: list[str] | None
|
allowed_models: list[str] | None
|
||||||
rate_limit: int
|
rate_limit: int | None
|
||||||
is_active: bool
|
is_active: bool
|
||||||
expires_at: datetime | None = None
|
expires_at: datetime | None = None
|
||||||
is_standalone: bool = False
|
is_standalone: bool = False
|
||||||
@@ -772,6 +787,18 @@ class CreateMyApiKeyRequest(BaseModel):
|
|||||||
"""创建我的API密钥请求"""
|
"""创建我的API密钥请求"""
|
||||||
|
|
||||||
name: str
|
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):
|
class ProviderConfig(BaseModel):
|
||||||
|
|||||||
@@ -110,6 +110,9 @@ class User(Base):
|
|||||||
allowed_providers = Column(JSON, nullable=True) # 允许使用的提供商 ID 列表
|
allowed_providers = Column(JSON, nullable=True) # 允许使用的提供商 ID 列表
|
||||||
allowed_api_formats = Column(JSON, nullable=True) # 允许使用的 API 格式列表
|
allowed_api_formats = Column(JSON, nullable=True) # 允许使用的 API 格式列表
|
||||||
allowed_models = Column(JSON, nullable=True) # 允许使用的模型名称列表
|
allowed_models = Column(JSON, nullable=True) # 允许使用的模型名称列表
|
||||||
|
rate_limit = Column(
|
||||||
|
Integer, nullable=True, default=None
|
||||||
|
) # 每分钟请求限制,NULL=继承系统默认,0=不限制,N=N RPM
|
||||||
|
|
||||||
# Key 能力配置
|
# Key 能力配置
|
||||||
model_capability_settings = Column(JSON, nullable=True) # 用户针对特定模型的能力配置
|
model_capability_settings = Column(JSON, nullable=True) # 用户针对特定模型的能力配置
|
||||||
@@ -209,7 +212,9 @@ class ApiKey(Base):
|
|||||||
allowed_providers = Column(JSON, nullable=True) # 允许使用的提供商 ID 列表
|
allowed_providers = Column(JSON, nullable=True) # 允许使用的提供商 ID 列表
|
||||||
allowed_api_formats = Column(JSON, nullable=True) # 允许使用的 API 格式列表
|
allowed_api_formats = Column(JSON, nullable=True) # 允许使用的 API 格式列表
|
||||||
allowed_models = Column(JSON, nullable=True) # 允许使用的模型名称列表
|
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) # 并发请求限制
|
concurrent_limit = Column(Integer, default=5, nullable=True) # 并发请求限制
|
||||||
|
|
||||||
# Key 能力配置
|
# Key 能力配置
|
||||||
|
|||||||
@@ -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.concurrency_manager import ConcurrencyManager
|
||||||
from src.services.rate_limit.detector import RateLimitDetector
|
from src.services.rate_limit.detector import RateLimitDetector
|
||||||
from src.services.rate_limit.ip_limiter import IPRateLimiter
|
from src.services.rate_limit.ip_limiter import IPRateLimiter
|
||||||
|
from src.services.rate_limit.user_rpm_limiter import UserRpmLimiter, get_user_rpm_limiter
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"AdaptiveConcurrencyManager", # 向后兼容
|
"AdaptiveConcurrencyManager", # 向后兼容
|
||||||
@@ -19,5 +20,7 @@ __all__ = [
|
|||||||
"ConcurrencyManager",
|
"ConcurrencyManager",
|
||||||
"IPRateLimiter",
|
"IPRateLimiter",
|
||||||
"RateLimitDetector",
|
"RateLimitDetector",
|
||||||
|
"UserRpmLimiter",
|
||||||
"get_adaptive_rpm_manager",
|
"get_adaptive_rpm_manager",
|
||||||
|
"get_user_rpm_limiter",
|
||||||
]
|
]
|
||||||
|
|||||||
354
src/services/rate_limit/user_rpm_limiter.py
Normal file
354
src/services/rate_limit/user_rpm_limiter.py
Normal 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
|
||||||
@@ -61,6 +61,10 @@ class ApiKeyService:
|
|||||||
if final_expires_at is None and expire_days:
|
if final_expires_at is None and expire_days:
|
||||||
final_expires_at = datetime.now(timezone.utc) + timedelta(days=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(
|
api_key = ApiKey(
|
||||||
user_id=user_id,
|
user_id=user_id,
|
||||||
key_hash=key_hash,
|
key_hash=key_hash,
|
||||||
@@ -69,7 +73,7 @@ class ApiKeyService:
|
|||||||
allowed_providers=allowed_providers,
|
allowed_providers=allowed_providers,
|
||||||
allowed_api_formats=allowed_api_formats,
|
allowed_api_formats=allowed_api_formats,
|
||||||
allowed_models=allowed_models,
|
allowed_models=allowed_models,
|
||||||
rate_limit=rate_limit,
|
rate_limit=normalized_rate_limit,
|
||||||
concurrent_limit=concurrent_limit,
|
concurrent_limit=concurrent_limit,
|
||||||
expires_at=final_expires_at,
|
expires_at=final_expires_at,
|
||||||
is_standalone=is_standalone,
|
is_standalone=is_standalone,
|
||||||
@@ -144,7 +148,8 @@ class ApiKeyService:
|
|||||||
# 允许显式设置为空数组/None 的字段(NULL=不限制,[]=全部禁用)
|
# 允许显式设置为空数组/None 的字段(NULL=不限制,[]=全部禁用)
|
||||||
nullable_list_fields = {"allowed_providers", "allowed_api_formats", "allowed_models"}
|
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"}
|
nullable_fields = {"expires_at", "rate_limit"}
|
||||||
|
|
||||||
for field, value in kwargs.items():
|
for field, value in kwargs.items():
|
||||||
@@ -154,7 +159,9 @@ class ApiKeyService:
|
|||||||
if field in nullable_list_fields:
|
if field in nullable_list_fields:
|
||||||
setattr(api_key, field, value)
|
setattr(api_key, field, value)
|
||||||
elif field in nullable_fields:
|
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)
|
setattr(api_key, field, value)
|
||||||
elif value is not None:
|
elif value is not None:
|
||||||
setattr(api_key, field, value)
|
setattr(api_key, field, value)
|
||||||
@@ -180,39 +187,6 @@ class ApiKeyService:
|
|||||||
logger.info(f"删除API密钥: ID {key_id}")
|
logger.info(f"删除API密钥: ID {key_id}")
|
||||||
return True
|
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
|
@staticmethod
|
||||||
def cleanup_expired_keys(db: Session, auto_delete: bool = False) -> int:
|
def cleanup_expired_keys(db: Session, auto_delete: bool = False) -> int:
|
||||||
"""清理过期的API密钥
|
"""清理过期的API密钥
|
||||||
|
|||||||
@@ -38,6 +38,7 @@ class UserService:
|
|||||||
allowed_providers: list[str] | None = None,
|
allowed_providers: list[str] | None = None,
|
||||||
allowed_api_formats: list[str] | None = None,
|
allowed_api_formats: list[str] | None = None,
|
||||||
allowed_models: list[str] | None = None,
|
allowed_models: list[str] | None = None,
|
||||||
|
rate_limit: int | None = None,
|
||||||
) -> User:
|
) -> User:
|
||||||
"""创建新用户。"""
|
"""创建新用户。"""
|
||||||
|
|
||||||
@@ -74,6 +75,7 @@ class UserService:
|
|||||||
allowed_providers=allowed_providers,
|
allowed_providers=allowed_providers,
|
||||||
allowed_api_formats=allowed_api_formats,
|
allowed_api_formats=allowed_api_formats,
|
||||||
allowed_models=allowed_models,
|
allowed_models=allowed_models,
|
||||||
|
rate_limit=rate_limit,
|
||||||
)
|
)
|
||||||
user.set_password(password)
|
user.set_password(password)
|
||||||
|
|
||||||
@@ -218,6 +220,7 @@ class UserService:
|
|||||||
"allowed_providers",
|
"allowed_providers",
|
||||||
"allowed_api_formats",
|
"allowed_api_formats",
|
||||||
"allowed_models",
|
"allowed_models",
|
||||||
|
"rate_limit",
|
||||||
]
|
]
|
||||||
|
|
||||||
# 允许设置为 None 的字段(表示无限制)
|
# 允许设置为 None 的字段(表示无限制)
|
||||||
@@ -225,6 +228,7 @@ class UserService:
|
|||||||
"allowed_providers",
|
"allowed_providers",
|
||||||
"allowed_api_formats",
|
"allowed_api_formats",
|
||||||
"allowed_models",
|
"allowed_models",
|
||||||
|
"rate_limit",
|
||||||
]
|
]
|
||||||
|
|
||||||
for field, value in kwargs.items():
|
for field, value in kwargs.items():
|
||||||
|
|||||||
@@ -1,7 +1,9 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from contextlib import contextmanager
|
||||||
from datetime import datetime, timezone
|
from datetime import datetime, timezone
|
||||||
from types import SimpleNamespace
|
from types import SimpleNamespace
|
||||||
|
from typing import Generator
|
||||||
from unittest.mock import MagicMock
|
from unittest.mock import MagicMock
|
||||||
|
|
||||||
import pytest
|
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 (
|
from src.api.admin.users.routes import (
|
||||||
AdminGetUserKeyFullKeyAdapter,
|
AdminGetUserKeyFullKeyAdapter,
|
||||||
AdminToggleUserKeyLockAdapter,
|
AdminToggleUserKeyLockAdapter,
|
||||||
|
AdminUpdateUserKeyAdapter,
|
||||||
)
|
)
|
||||||
from src.api.admin.users.routes import router as admin_users_router
|
from src.api.admin.users.routes import router as admin_users_router
|
||||||
from src.core.exceptions import InvalidRequestException, NotFoundException
|
from src.core.exceptions import InvalidRequestException, NotFoundException
|
||||||
@@ -25,6 +28,15 @@ from src.database import get_db
|
|||||||
from src.models.api import CreateApiKeyRequest
|
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:
|
def _build_context(db: MagicMock) -> SimpleNamespace:
|
||||||
return SimpleNamespace(
|
return SimpleNamespace(
|
||||||
db=db,
|
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
|
*, adapter: object, http_request: object, db: MagicMock, mode: object
|
||||||
) -> object:
|
) -> object:
|
||||||
_ = http_request, mode
|
_ = http_request, mode
|
||||||
|
try:
|
||||||
|
payload = await http_request.json()
|
||||||
|
except Exception:
|
||||||
|
payload = {}
|
||||||
context = SimpleNamespace(
|
context = SimpleNamespace(
|
||||||
db=db,
|
db=db,
|
||||||
request=SimpleNamespace(state=SimpleNamespace()),
|
request=SimpleNamespace(state=SimpleNamespace()),
|
||||||
user=SimpleNamespace(id="admin-1"),
|
user=SimpleNamespace(id="admin-1"),
|
||||||
ensure_json_body=lambda: {},
|
ensure_json_body=lambda: payload,
|
||||||
add_audit_metadata=lambda **_: None,
|
add_audit_metadata=lambda **_: None,
|
||||||
)
|
)
|
||||||
return await adapter.handle(context)
|
return await adapter.handle(context)
|
||||||
@@ -82,10 +98,11 @@ def _build_admin_api_keys_app(db: MagicMock, monkeypatch: pytest.MonkeyPatch) ->
|
|||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@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()
|
db = MagicMock()
|
||||||
api_key = SimpleNamespace(id="key-1", user_id="user-1", is_standalone=False, is_locked=False)
|
api_key = SimpleNamespace(id="key-1", user_id="user-1", is_standalone=False, is_locked=False)
|
||||||
_mock_query_first(db, api_key)
|
_mock_query_first(db, api_key)
|
||||||
|
_patch_get_db_context(monkeypatch, db)
|
||||||
|
|
||||||
adapter = AdminToggleUserKeyLockAdapter(user_id="user-1", key_id="key-1")
|
adapter = AdminToggleUserKeyLockAdapter(user_id="user-1", key_id="key-1")
|
||||||
result = await adapter.handle(_build_context(db))
|
result = await adapter.handle(_build_context(db))
|
||||||
@@ -98,9 +115,12 @@ async def test_toggle_user_key_lock_adapter_success() -> None:
|
|||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@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()
|
db = MagicMock()
|
||||||
_mock_query_first(db, None)
|
_mock_query_first(db, None)
|
||||||
|
_patch_get_db_context(monkeypatch, db)
|
||||||
|
|
||||||
adapter = AdminToggleUserKeyLockAdapter(user_id="user-1", key_id="key-standalone")
|
adapter = AdminToggleUserKeyLockAdapter(user_id="user-1", key_id="key-standalone")
|
||||||
with pytest.raises(NotFoundException):
|
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
|
@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()
|
db = MagicMock()
|
||||||
normal_key = SimpleNamespace(
|
normal_key = SimpleNamespace(
|
||||||
id="key-user",
|
id="key-user",
|
||||||
@@ -182,6 +204,7 @@ async def test_standalone_toggle_adapters_reject_normal_user_key() -> None:
|
|||||||
updated_at=datetime.now(timezone.utc),
|
updated_at=datetime.now(timezone.utc),
|
||||||
)
|
)
|
||||||
_mock_query_first(db, normal_key)
|
_mock_query_first(db, normal_key)
|
||||||
|
_patch_get_db_context(monkeypatch, db)
|
||||||
context = _build_context(db)
|
context = _build_context(db)
|
||||||
|
|
||||||
with pytest.raises(InvalidRequestException):
|
with pytest.raises(InvalidRequestException):
|
||||||
@@ -195,6 +218,7 @@ def test_user_key_lock_route_path_smoke(monkeypatch: pytest.MonkeyPatch) -> None
|
|||||||
db = MagicMock()
|
db = MagicMock()
|
||||||
api_key = SimpleNamespace(id="key-5", user_id="user-2", is_standalone=False, is_locked=False)
|
api_key = SimpleNamespace(id="key-5", user_id="user-2", is_standalone=False, is_locked=False)
|
||||||
_mock_query_first(db, api_key)
|
_mock_query_first(db, api_key)
|
||||||
|
_patch_get_db_context(monkeypatch, db)
|
||||||
client = _build_admin_users_app(db, monkeypatch)
|
client = _build_admin_users_app(db, monkeypatch)
|
||||||
|
|
||||||
response = client.patch("/api/admin/users/user-2/api-keys/key-5/lock")
|
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"}
|
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:
|
def test_standalone_lock_route_removed(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||||
client = _build_admin_api_keys_app(MagicMock(), monkeypatch)
|
client = _build_admin_api_keys_app(MagicMock(), monkeypatch)
|
||||||
response = client.patch("/api/admin/api-keys/key-1/lock")
|
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:
|
def test_standalone_list_route_does_not_expose_is_locked(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||||
db = MagicMock()
|
db = MagicMock()
|
||||||
|
_patch_get_db_context(monkeypatch, db)
|
||||||
api_key = SimpleNamespace(
|
api_key = SimpleNamespace(
|
||||||
id="sa-key-1",
|
id="sa-key-1",
|
||||||
user_id="admin-1",
|
user_id="admin-1",
|
||||||
@@ -306,6 +397,7 @@ async def test_create_standalone_key_adapter_preserves_empty_restriction_lists(
|
|||||||
monkeypatch: pytest.MonkeyPatch,
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
) -> None:
|
) -> None:
|
||||||
db = MagicMock()
|
db = MagicMock()
|
||||||
|
_patch_get_db_context(monkeypatch, db)
|
||||||
captured: dict[str, object] = {}
|
captured: dict[str, object] = {}
|
||||||
created_key = SimpleNamespace(
|
created_key = SimpleNamespace(
|
||||||
id="sa-key-3",
|
id="sa-key-3",
|
||||||
@@ -365,6 +457,7 @@ async def test_update_standalone_key_adapter_preserves_empty_restriction_lists(
|
|||||||
monkeypatch: pytest.MonkeyPatch,
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
) -> None:
|
) -> None:
|
||||||
db = MagicMock()
|
db = MagicMock()
|
||||||
|
_patch_get_db_context(monkeypatch, db)
|
||||||
existing_key = SimpleNamespace(id="sa-key-4", is_standalone=True)
|
existing_key = SimpleNamespace(id="sa-key-4", is_standalone=True)
|
||||||
_mock_query_first(db, existing_key)
|
_mock_query_first(db, existing_key)
|
||||||
|
|
||||||
|
|||||||
@@ -49,6 +49,7 @@ def test_list_users_uses_wallet_batch_lookup(monkeypatch: pytest.MonkeyPatch) ->
|
|||||||
allowed_providers=None,
|
allowed_providers=None,
|
||||||
allowed_api_formats=None,
|
allowed_api_formats=None,
|
||||||
allowed_models=None,
|
allowed_models=None,
|
||||||
|
rate_limit=None,
|
||||||
is_active=True,
|
is_active=True,
|
||||||
created_at=now,
|
created_at=now,
|
||||||
updated_at=now,
|
updated_at=now,
|
||||||
@@ -62,6 +63,7 @@ def test_list_users_uses_wallet_batch_lookup(monkeypatch: pytest.MonkeyPatch) ->
|
|||||||
allowed_providers=None,
|
allowed_providers=None,
|
||||||
allowed_api_formats=None,
|
allowed_api_formats=None,
|
||||||
allowed_models=None,
|
allowed_models=None,
|
||||||
|
rate_limit=None,
|
||||||
is_active=True,
|
is_active=True,
|
||||||
created_at=now,
|
created_at=now,
|
||||||
updated_at=None,
|
updated_at=None,
|
||||||
@@ -101,24 +103,14 @@ async def test_create_user_adapter_preserves_empty_restriction_lists(
|
|||||||
db = MagicMock()
|
db = MagicMock()
|
||||||
captured: dict[str, Any] = {}
|
captured: dict[str, Any] = {}
|
||||||
|
|
||||||
def _create_user(**kwargs: Any) -> SimpleNamespace:
|
def _create_user_sync(request: Any, role: Any) -> tuple[dict[str, Any], dict[str, Any]]:
|
||||||
captured.update(kwargs)
|
captured["allowed_providers"] = request.allowed_providers
|
||||||
return SimpleNamespace(
|
captured["allowed_api_formats"] = request.allowed_api_formats
|
||||||
id="user-3",
|
captured["allowed_models"] = request.allowed_models
|
||||||
email="u3@example.com",
|
captured["role"] = role
|
||||||
username="user3",
|
return {"id": "user-3"}, {}
|
||||||
role=SimpleNamespace(value="user"),
|
|
||||||
is_active=True,
|
|
||||||
allowed_providers=[],
|
|
||||||
allowed_api_formats=[],
|
|
||||||
allowed_models=[],
|
|
||||||
)
|
|
||||||
|
|
||||||
monkeypatch.setattr("src.api.admin.users.routes.UserService.create_user", _create_user)
|
monkeypatch.setattr("src.api.admin.users.routes._create_user_sync", _create_user_sync)
|
||||||
monkeypatch.setattr(
|
|
||||||
"src.api.admin.users.routes._serialize_user",
|
|
||||||
lambda _db, user: {"id": user.id},
|
|
||||||
)
|
|
||||||
|
|
||||||
context = SimpleNamespace(
|
context = SimpleNamespace(
|
||||||
db=db,
|
db=db,
|
||||||
|
|||||||
120
tests/api/test_monitoring_user_rate_limit.py
Normal file
120
tests/api/test_monitoring_user_rate_limit.py
Normal file
@@ -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()
|
||||||
@@ -18,6 +18,7 @@ from src.api.base.adapter import ApiMode
|
|||||||
from src.api.base.pipeline import ApiRequestPipeline
|
from src.api.base.pipeline import ApiRequestPipeline
|
||||||
from src.core.enums import UserRole
|
from src.core.enums import UserRole
|
||||||
from src.core.modules.hooks import AUTH_TOKEN_PREFIX_AUTHENTICATORS
|
from src.core.modules.hooks import AUTH_TOKEN_PREFIX_AUTHENTICATORS
|
||||||
|
from src.services.rate_limit.user_rpm_limiter import RpmCheckResult
|
||||||
|
|
||||||
|
|
||||||
class TestPipelineBalanceCalculation:
|
class TestPipelineBalanceCalculation:
|
||||||
@@ -499,6 +500,166 @@ class TestPipelineAuthentication:
|
|||||||
assert "锁定" in str(exc_info.value.detail)
|
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:
|
class TestPipelineTokenPrefixAuth:
|
||||||
"""Tests token-prefix auth isolation."""
|
"""Tests token-prefix auth isolation."""
|
||||||
|
|
||||||
|
|||||||
110
tests/api/test_user_me_api_key_routes.py
Normal file
110
tests/api/test_user_me_api_key_routes.py
Normal 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
tests/services/test_user_rpm_limiter.py
Normal file
148
tests/services/test_user_rpm_limiter.py
Normal 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_hash == legacy_hash
|
||||||
assert key_encrypted == legacy_encrypted
|
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
|
||||||
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user