i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-09 04:30:20 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(pool,scheduling): 号池调度维度、配额冷却机制与管理后台重构

Browse Source 
- 新增 scheduling_dimensions 模块，为每个 Key 计算多维调度状态(手动/冷却/熔断/成本/健康)
- 新增 quota_cooldown 模块，统一判定 Key 的有效冷却原因
- Pool 管理后台 API 扩展 Key 详情字段(调度状态/维度/配额/OAuth 信息)
- 前端 Pool 管理页面重写，支持调度状态展示、批量清理封禁 Key
- Handler 基类增加请求调度元数据采集，stream telemetry 增强
- 请求时间线组件增强，支持 attempted 候选展示
- Kiro OAuth 凭证导入解析改进
- 新增 usage 表 provider_key 索引迁移
- 补充调度维度、配额冷却、候选枚举等单元测试

Closes #197

Co-authored-by: AAEE86 <33052466+AAEE86@users.noreply.github.com>
This commit is contained in:
fawney19
2026-03-03 09:22:20 +08:00
co-authored by AAEE86
parent f787b1b02a
commit 11997c024e
40 changed files with 4419 additions and 823 deletions
Download Patch File Download Diff File Expand all files Collapse all files
alembic/versions/20260303_1000_a1b2c3d4e5f6_add_idx_usage_provider_key.py
+45
View File
@@ -0,0 +1,45 @@
"""add_idx_usage_provider_key
Add composite index on usage(provider_id, provider_api_key_id) to support
the pool management page's per-key usage stats aggregation query.
Revision ID: a1b2c3d4e5f6
Revises: dd0278c0a28c
Create Date: 2026-03-03 10:00:00.000000+00:00
"""
from __future__ import annotations
import sqlalchemy as sa
from alembic import op
revision = "a1b2c3d4e5f6"
down_revision = "dd0278c0a28c"
branch_labels = None
depends_on = None
INDEX_NAME = "idx_usage_provider_key"
TABLE = "usage"
COLUMNS = ["provider_id", "provider_api_key_id"]
def upgrade() -> None:
bind = op.get_bind()
result = bind.execute(
sa.text("SELECT 1 FROM pg_indexes WHERE indexname = :name"),
{"name": INDEX_NAME},
).fetchone()
if result:
return
op.create_index(INDEX_NAME, TABLE, COLUMNS)
def downgrade() -> None:
bind = op.get_bind()
result = bind.execute(
sa.text("SELECT 1 FROM pg_indexes WHERE indexname = :name"),
{"name": INDEX_NAME},
).fetchone()
if not result:
return
op.drop_index(INDEX_NAME, table_name=TABLE)
frontend/src/api/endpoints/pool.ts
+73 -20
View File
@@ -1,4 +1,5 @@
import client from '../client'
import type { AllowedModels, ProxyConfig } from './types/provider'
export interface PoolKeyStatus {
key_id: string
@@ -78,15 +79,80 @@ export interface PoolKeyDetail {
key_name: string
is_active: boolean
auth_type: string
oauth_expires_at?: number | null
oauth_invalid_at?: number | null
oauth_invalid_reason?: string | null
oauth_plan_type?: string | null
quota_updated_at?: number | null
health_score?: number
circuit_breaker_open?: boolean
api_formats?: string[]
rate_multipliers?: Record<string, number> | null
internal_priority?: number
rpm_limit?: number | null
cache_ttl_minutes?: number
max_probe_interval_minutes?: number
note?: string | null
allowed_models?: AllowedModels
capabilities?: Record<string, boolean> | null
auto_fetch_models?: boolean
locked_models?: string[] | null
model_include_patterns?: string[] | null
model_exclude_patterns?: string[] | null
proxy?: ProxyConfig | null
account_quota: string | null
cooldown_reason: string | null
cooldown_ttl_seconds: number | null
cost_window_usage: number
cost_limit: number | null
request_count: number
total_tokens: number
total_cost_usd: number
sticky_sessions: number
lru_score: number | null
created_at: string | null
last_used_at: string | null
scheduling_status?: 'available' | 'degraded' | 'blocked'
scheduling_reason?:
| 'available'
| 'manual_disabled'
| 'cooldown'
| 'circuit_open'
| 'cost_exhausted'
| 'cost_soft'
| 'cost'
| 'health_low'
| 'health_degraded'
| 'health'
| string
scheduling_label?: string
scheduling_reasons?: PoolSchedulingReason[]
scheduling_score?: number
candidate_eligible?: boolean
scheduling_blocked_count?: number
scheduling_degraded_count?: number
scheduling_dimensions?: PoolSchedulingDimension[]
}
export interface PoolSchedulingReason {
code: string
label: string
blocking: boolean
source: 'manual' | 'pool' | 'health' | 'policy' | string
ttl_seconds?: number | null
detail?: string | null
}
export interface PoolSchedulingDimension {
code: string
label: string
status: 'ok' | 'degraded' | 'blocked' | string
blocking: boolean
source: 'manual' | 'pool' | 'health' | 'policy' | string
weight: number
score: number
ttl_seconds?: number | null
detail?: string | null
}
export interface PoolKeysPageResponse {
@@ -103,18 +169,6 @@ export interface PoolKeysQuery {
status?: 'all' | 'active' | 'cooldown' | 'inactive'
}
export interface PoolKeyImportItem {
name: string
api_key: string
auth_type?: string
}
export interface BatchImportResponse {
imported: number
skipped: number
errors: { index: number; reason: string }[]
}
export interface PoolBatchAction {
key_ids: string[]
action: 'enable' | 'disable' | 'delete' | 'clear_cooldown' | 'reset_cost'
@@ -133,14 +187,6 @@ export async function listPoolKeys(
return response.data
}
export async function batchImportPoolKeys(
providerId: string,
keys: PoolKeyImportItem[],
): Promise<BatchImportResponse> {
const response = await client.post(`/api/admin/pool/${providerId}/keys/batch-import`, { keys })
return response.data
}
export async function batchActionPoolKeys(
providerId: string,
body: PoolBatchAction,
@@ -148,3 +194,10 @@ export async function batchActionPoolKeys(
const response = await client.post(`/api/admin/pool/${providerId}/keys/batch-action`, body)
return response.data
}
export async function cleanupBannedPoolKeys(
providerId: string,
): Promise<{ affected: number; message: string }> {
const response = await client.post(`/api/admin/pool/${providerId}/keys/cleanup-banned`)
return response.data
}
frontend/src/components/common/JsonImportInput.vue
+281
View File
@@ -0,0 +1,281 @@
<template>
<div class="space-y-4">
<input
ref="fileInputRef"
type="file"
:accept="accept"
multiple
class="hidden"
@change="handleFileSelect"
>
<div
v-if="!showManualInput"
class="rounded-xl border-2 border-dashed transition-colors cursor-pointer"
:class="isDragging
? 'border-primary bg-primary/5'
: 'border-border hover:border-muted-foreground/40'"
@click="fileInputRef?.click()"
@dragover.prevent="isDragging = true"
@dragleave.prevent="isDragging = false"
@drop.prevent="handleFileDrop"
>
<div class="flex flex-col items-center justify-center py-10 gap-2">
<div class="w-9 h-9 rounded-full bg-muted/60 flex items-center justify-center">
<Upload class="w-4 h-4 text-muted-foreground" />
</div>
<div class="text-center">
<p class="text-xs font-medium">
{{ dropTitle }}
</p>
<p class="text-[11px] text-muted-foreground mt-0.5">
{{ dropHint }}
</p>
</div>
</div>
</div>
<div
v-else
class="space-y-1.5"
>
<Label v-if="manualLabel">
{{ manualLabel }}
</Label>
<Textarea
:model-value="modelValue"
:disabled="disabled"
:placeholder="manualPlaceholder"
:class="textareaClass"
spellcheck="false"
@update:model-value="emit('update:modelValue', $event)"
/>
<p
v-if="manualDescription"
class="text-xs text-muted-foreground"
>
{{ manualDescription }}
</p>
</div>
<div class="flex items-center justify-center pt-1">
<button
v-if="!showManualInput"
type="button"
class="text-sm text-muted-foreground hover:text-foreground transition-colors"
@click="showManualInput = true"
>
{{ pasteToggleText }}
</button>
<button
v-else
type="button"
class="text-sm text-muted-foreground hover:text-foreground transition-colors"
@click="switchToFileMode"
>
{{ fileToggleText }}
</button>
</div>
</div>
</template>
<script setup lang="ts">
import { computed, ref, watch } from 'vue'
import { Upload } from 'lucide-vue-next'
import { Label, Textarea } from '@/components/ui'
interface ImportInputErrorPayload {
message: string
title?: string
}
const props = withDefaults(defineProps<{
modelValue: string
disabled?: boolean
resetKey?: string | number
accept?: string
dropTitle?: string
dropHint?: string
manualLabel?: string
manualPlaceholder?: string
manualDescription?: string
pasteToggleText?: string
fileToggleText?: string
textareaClass?: string
}>(), {
disabled: false,
resetKey: '',
accept: '.json,.txt',
dropTitle: '拖入导入文件或点击选择',
dropHint: '支持 .json / .txt，可多选',
manualLabel: '',
manualPlaceholder: '',
manualDescription: '',
pasteToggleText: '或手动粘贴 JSON',
fileToggleText: '或选择 JSON 文件导入',
textareaClass: 'min-h-[220px] text-xs font-mono break-all !rounded-xl',
})
const emit = defineEmits<{
'update:modelValue': [value: string]
error: [payload: ImportInputErrorPayload]
}>()
const showManualInput = ref(false)
const isDragging = ref(false)
const fileInputRef = ref<HTMLInputElement | null>(null)
const acceptParts = computed(() => {
return props.accept
.split(',')
.map(part => part.trim().toLowerCase())
.filter(Boolean)
})
const acceptedExtensions = computed(() => acceptParts.value.filter(part => part.startsWith('.')))
const acceptedMimeTypes = computed(() => acceptParts.value.filter(part => part.includes('/')))
function resetUiState() {
showManualInput.value = false
isDragging.value = false
if (fileInputRef.value) {
fileInputRef.value.value = ''
}
}
function emitError(message: string, title?: string) {
emit('error', { message, title })
}
function isValidFileType(file: File): boolean {
const name = file.name.toLowerCase()
const type = (file.type || '').toLowerCase()
const extensionAllowed = acceptedExtensions.value.length > 0
&& acceptedExtensions.value.some(ext => name.endsWith(ext))
const mimeAllowed = acceptedMimeTypes.value.length > 0
&& acceptedMimeTypes.value.some((mimeType) => {
if (mimeType.endsWith('/*')) {
const prefix = mimeType.slice(0, -1)
return type.startsWith(prefix)
}
return type === mimeType
})
if (acceptedExtensions.value.length === 0 && acceptedMimeTypes.value.length === 0) {
return true
}
return extensionAllowed || mimeAllowed
}
function readFileAsText(file: File): Promise<string> {
return new Promise((resolve, reject) => {
const reader = new FileReader()
reader.onload = (event) => {
const content = event.target?.result
if (typeof content === 'string') {
resolve(content)
return
}
reject(new Error('读取失败'))
}
reader.onerror = () => reject(new Error('读取失败'))
reader.readAsText(file)
})
}
function mergeFileContents(contents: string[]): string {
const items: unknown[] = []
for (const raw of contents) {
const trimmed = raw.trim()
if (!trimmed) continue
try {
const parsed = JSON.parse(trimmed)
if (Array.isArray(parsed)) {
items.push(...parsed)
} else {
items.push(parsed)
}
continue
} catch {
// Fallback to line mode.
}
const lines = trimmed
.split('\n')
.map(line => line.trim())
.filter(line => line && !line.startsWith('#'))
items.push(...lines)
}
if (items.length === 1) {
return typeof items[0] === 'string' ? items[0] : JSON.stringify(items[0], null, 2)
}
return JSON.stringify(items, null, 2)
}
async function readFiles(files: File[]) {
const validFiles = files.filter(isValidFileType)
if (validFiles.length === 0) {
emitError('仅支持 .json 或 .txt 文件', '格式错误')
return
}
if (validFiles.length < files.length) {
emitError(`已忽略 ${files.length - validFiles.length} 个不支持的文件`, '提示')
}
try {
const contents = await Promise.all(validFiles.map(readFileAsText))
const merged = validFiles.length === 1 ? contents[0] : mergeFileContents(contents)
emit('update:modelValue', merged)
showManualInput.value = true
} catch {
emitError('文件读取失败', '错误')
} finally {
if (fileInputRef.value) {
fileInputRef.value.value = ''
}
}
}
function handleFileSelect(event: Event) {
const input = event.target as HTMLInputElement
const files = input.files
if (!files || files.length === 0) return
void readFiles(Array.from(files))
}
function handleFileDrop(event: DragEvent) {
isDragging.value = false
const files = event.dataTransfer?.files
if (!files || files.length === 0) return
void readFiles(Array.from(files))
}
function switchToFileMode() {
showManualInput.value = false
emit('update:modelValue', '')
if (fileInputRef.value) {
fileInputRef.value.value = ''
}
}
watch(
() => props.resetKey,
() => {
resetUiState()
},
)
watch(
() => props.modelValue,
(value) => {
if (value.trim() && !showManualInput.value) {
showManualInput.value = true
}
},
)
</script>
frontend/src/features/pool/components/PoolImportDialog.vue
-156
View File
@@ -1,156 +0,0 @@
<template>
<Dialog
:model-value="modelValue"
title="批量导入账号"
description="以 JSON 格式批量导入 API Key 到号池"
size="lg"
@update:model-value="emit('update:modelValue', $event)"
>
<div class="space-y-4">
<div class="space-y-1.5">
<Label>JSON 数据</Label>
<textarea
v-model="jsonText"
class="w-full h-48 p-3 text-sm font-mono border rounded-lg bg-background resize-none focus:outline-none focus:ring-2 focus:ring-ring"
placeholder="[{&quot;name&quot;: &quot;key-01&quot;, &quot;api_key&quot;: &quot;sk-xxx&quot;}, ...]"
/>
<p class="text-xs text-muted-foreground">
格式: [{"name": "名称", "api_key": "密钥", "auth_type": "api_key"}]，auth_type 可选
</p>
</div>
<div
v-if="parseError"
class="text-sm text-destructive"
>
{{ parseError }}
</div>
<div
v-if="parsedCount > 0"
class="text-sm text-muted-foreground"
>
已解析 {{ parsedCount }} 个账号
</div>
<div
v-if="importResult"
class="space-y-1 text-sm"
>
<p class="text-green-600">
成功导入: {{ importResult.imported }}
</p>
<p
v-if="importResult.errors.length > 0"
class="text-destructive"
>
失败: {{ importResult.errors.length }}
</p>
<div
v-for="err in importResult.errors.slice(0, 5)"
:key="err.index"
class="text-xs text-destructive"
>
#{{ err.index }}: {{ err.reason }}
</div>
</div>
</div>
<template #footer>
<Button
variant="outline"
:disabled="loading"
@click="emit('update:modelValue', false)"
>
{{ importResult ? '关闭' : '取消' }}
</Button>
<Button
v-if="!importResult"
:disabled="loading || parsedCount === 0"
@click="handleImport"
>
{{ loading ? '导入中...' : `导入 ${parsedCount} 个账号` }}
</Button>
</template>
</Dialog>
</template>
<script setup lang="ts">
import { ref, computed, watch } from 'vue'
import { Dialog, Button, Label } from '@/components/ui'
import { useToast } from '@/composables/useToast'
import { parseApiError } from '@/utils/errorParser'
import { batchImportPoolKeys } from '@/api/endpoints/pool'
import type { BatchImportResponse, PoolKeyImportItem } from '@/api/endpoints/pool'
const props = defineProps<{
modelValue: boolean
providerId: string
}>()
const emit = defineEmits<{
'update:modelValue': [value: boolean]
imported: []
}>()
const { error: showError } = useToast()
const jsonText = ref('')
const loading = ref(false)
const parseError = ref('')
const importResult = ref<BatchImportResponse | null>(null)
const parsedKeys = computed<PoolKeyImportItem[]>(() => {
if (!jsonText.value.trim()) return []
try {
const data = JSON.parse(jsonText.value)
if (!Array.isArray(data)) {
return []
}
return data.map((item: Record<string, unknown>) => ({
name: String(item.name || ''),
api_key: String(item.api_key || ''),
auth_type: String(item.auth_type || 'api_key'),
}))
} catch {
return []
}
})
watch(jsonText, (val) => {
if (!val.trim()) {
parseError.value = ''
return
}
try {
const data = JSON.parse(val)
parseError.value = Array.isArray(data) ? '' : 'JSON 必须是数组格式'
} catch {
parseError.value = 'JSON 格式无效'
}
})
const parsedCount = computed(() => parsedKeys.value.length)
watch(() => props.modelValue, (v) => {
if (v) {
jsonText.value = ''
importResult.value = null
parseError.value = ''
}
})
async function handleImport() {
if (!parsedKeys.value.length) return
loading.value = true
try {
importResult.value = await batchImportPoolKeys(props.providerId, parsedKeys.value)
if (importResult.value.imported > 0) {
emit('imported')
}
} catch (err) {
showError(parseApiError(err))
} finally {
loading.value = false
}
}
</script>
frontend/src/features/providers/components/OAuthAccountDialog.vue
+15 -144
View File
@@ -404,70 +404,19 @@
class="flex flex-col gap-3 justify-center transition-opacity duration-150"
:class="mode === 'import' ? 'opacity-100' : 'opacity-0 pointer-events-none'"
>
<input
ref="fileInputRef"
type="file"
accept=".json,.txt"
multiple
class="hidden"
@change="handleFileSelect"
>
<!-- 拖拽模式 -->
<div
v-if="!showManualInput"
class="rounded-xl border-2 border-dashed transition-colors cursor-pointer"
:class="isDragging
? 'border-primary bg-primary/5'
: 'border-border hover:border-muted-foreground/40'"
@click="fileInputRef?.click()"
@dragover.prevent="isDragging = true"
@dragleave.prevent="isDragging = false"
@drop.prevent="handleFileDrop"
>
<div class="flex flex-col items-center justify-center py-12 gap-2">
<div class="w-9 h-9 rounded-full bg-muted/60 flex items-center justify-center">
<Upload class="w-4 h-4 text-muted-foreground" />
</div>
<div class="text-center">
<p class="text-xs font-medium">
拖入授权文件或点击选择
</p>
<p class="text-[11px] text-muted-foreground mt-0.5">
支持 .json / .txt，可多选
</p>
</div>
</div>
</div>
<!-- 粘贴模式 -->
<Textarea
v-else
<JsonImportInput
v-model="importText"
:disabled="importing"
placeholder="粘贴 Refresh Token 或 JSON 内容"
class="min-h-[200px] text-xs font-mono break-all !rounded-xl"
spellcheck="false"
:reset-key="importInputResetKey"
drop-title="拖入授权文件或点击选择"
drop-hint="支持 .json / .txt，可多选"
manual-placeholder="粘贴 Refresh Token 或 JSON 内容"
paste-toggle-text="或手动粘贴 Refresh Token"
file-toggle-text="或选择 JSON 文件导入"
textarea-class="min-h-[200px] text-xs font-mono break-all !rounded-xl"
@error="handleImportInputError"
/>
<!-- 底部切换链接 -->
<div class="flex items-center justify-center pt-1">
<button
v-if="!showManualInput"
class="text-sm text-muted-foreground hover:text-foreground transition-colors"
@click="showManualInput = true"
>
或手动粘贴 Refresh Token
</button>
<button
v-else
class="text-sm text-muted-foreground hover:text-foreground transition-colors"
@click="showManualInput = false; importText = ''"
>
或选择 JSON 文件导入
</button>
</div>
<div
v-if="importTask"
class="rounded-xl border border-border bg-muted/20 p-3 space-y-2"
@@ -554,7 +503,7 @@ import {
ComboboxTrigger,
ComboboxViewport,
} from 'radix-vue'
import { UserPlus, Copy, ExternalLink, Upload, Globe, AlertCircle, ShieldCheck, ChevronsUpDown, Check } from 'lucide-vue-next'
import { UserPlus, Copy, ExternalLink, Globe, AlertCircle, ShieldCheck, ChevronsUpDown, Check } from 'lucide-vue-next'
import { useToast } from '@/composables/useToast'
import { useClipboard } from '@/composables/useClipboard'
import { useTotp } from '@/composables/useTotp'
@@ -575,6 +524,7 @@ import type {
} from '@/api/endpoints/provider_oauth'
import ProxyNodeSelect from './ProxyNodeSelect.vue'
import { useProxyNodesStore } from '@/stores/proxy-nodes'
import JsonImportInput from '@/components/common/JsonImportInput.vue'
const props = defineProps<{
open: boolean
@@ -712,9 +662,7 @@ let countdownTimer: ReturnType<typeof setInterval> | null = null
// 导入状态
const importText = ref('')
const importing = ref(false)
const isDragging = ref(false)
const showManualInput = ref(false)
const fileInputRef = ref<HTMLInputElement | null>(null)
const importInputResetKey = ref(0)
const importTask = ref<OAuthBatchImportTaskStatusResponse | null>(null)
let importPollTimer: ReturnType<typeof setTimeout> | null = null
const importPolling = ref(false)
@@ -847,14 +795,10 @@ function resetForm() {
importText.value = ''
importing.value = false
importTask.value = null
isDragging.value = false
showManualInput.value = false
importInputResetKey.value += 1
proxyPopoverOpen.value = false
selectedProxyNodeId.value = ''
mode.value = 'oauth'
if (fileInputRef.value) {
fileInputRef.value.value = ''
}
}
function switchMode(newMode: DialogMode) {
@@ -978,81 +922,8 @@ function parseImportText(text: string): { refresh_token: string; name?: string }
return { refresh_token: trimmed }
}
function readFileAsText(file: File): Promise<string> {
return new Promise((resolve, reject) => {
const reader = new FileReader()
reader.onload = (e) => {
const content = e.target?.result
if (typeof content === 'string') resolve(content)
else reject(new Error('读取失败'))
}
reader.onerror = () => reject(new Error('读取失败'))
reader.readAsText(file)
})
}
function isValidFileType(file: File): boolean {
return file.name.endsWith('.json') || file.name.endsWith('.txt')
|| file.type === 'application/json' || file.type === 'text/plain'
}
/** 合并多个文件内容为统一的凭据文本（JSON 数组） */
function mergeFileContents(contents: string[]): string {
const items: unknown[] = []
for (const raw of contents) {
const trimmed = raw.trim()
if (!trimmed) continue
try {
const parsed = JSON.parse(trimmed)
if (Array.isArray(parsed)) {
items.push(...parsed)
} else {
items.push(parsed)
}
continue
} catch {
// not JSON
}
// 按行拆分（纯 Token）
const lines = trimmed.split('\n').filter(l => l.trim() && !l.trim().startsWith('#'))
items.push(...lines.map(l => l.trim()))
}
if (items.length === 1) {
// 单条：保持原始格式
return typeof items[0] === 'string' ? items[0] : JSON.stringify(items[0], null, 2)
}
return JSON.stringify(items, null, 2)
}
async function readFiles(files: File[]) {
const validFiles = files.filter(isValidFileType)
if (validFiles.length === 0) {
showError('仅支持 .json 或 .txt 文件', '格式错误')
return
}
if (validFiles.length < files.length) {
showError(`已忽略 ${files.length - validFiles.length} 个不支持的文件`, '提示')
}
try {
const contents = await Promise.all(validFiles.map(readFileAsText))
const merged = validFiles.length === 1 ? contents[0] : mergeFileContents(contents)
importText.value = merged
showManualInput.value = true
} catch {
showError('文件读取失败', '错误')
}
}
function handleFileSelect(event: Event) {
const input = event.target as HTMLInputElement
const files = input.files
if (files && files.length > 0) readFiles(Array.from(files))
}
function handleFileDrop(event: DragEvent) {
isDragging.value = false
const files = event.dataTransfer?.files
if (files && files.length > 0) readFiles(Array.from(files))
function handleImportInputError(payload: { message: string; title?: string }) {
showError(payload.message, payload.title)
}
async function handleImport() {
frontend/src/features/usage/components/HorizontalRequestTimeline.vue
+277 -53
View File
@@ -67,7 +67,7 @@
<div
class="node-dot"
:class="[
getStatusColorClass(group.primaryStatus),
getStatusColorClass(getDisplayStatus(group.primary)),
{ 'is-first-selected': isGroupSelected(group) && selectedAttemptIndex === 0 }
]"
@click.stop="selectFirstAttempt(group)"
@@ -83,7 +83,7 @@
:key="attempt.id"
class="sub-dot"
:class="[
getStatusColorClass(attempt.status),
getStatusColorClass(getDisplayStatus(attempt)),
{ active: selectedAttemptIndex === idx + 1 }
]"
:title="attempt.key_name || `Key ${idx + 2}`"
@@ -115,9 +115,9 @@
<div class="panel-title">
<span
class="title-dot"
:class="getStatusColorClass(currentAttempt.status)"
:class="getStatusColorClass(getDisplayStatus(currentAttempt))"
/>
<span class="title-text">{{ selectedGroup.providerName }}</span>
<span class="title-text">{{ currentGroupTitle }}</span>
<a
v-if="currentAttempt.provider_website"
:href="currentAttempt.provider_website"
@@ -130,7 +130,7 @@
</a>
<span
class="status-tag"
:class="getStatusColorClass(currentAttempt.status)"
:class="getStatusColorClass(getDisplayStatus(currentAttempt))"
>
{{ currentAttempt.status_code || getStatusLabel(currentAttempt.status) }}
</span>
@@ -188,12 +188,12 @@
<span class="info-value mono">{{ formatLatency(currentAttempt.extra_data.first_byte_time_ms) }}</span>
</div>
<div
v-if="currentAttempt.extra_data?.provider_api_format"
v-if="currentAttemptFormatDisplay"
class="info-item"
>
<span class="info-label">格式</span>
<span class="info-value">
<code class="format-code">{{ formatApiFormat(currentAttempt.extra_data.provider_api_format) }}</code>
<code class="format-code">{{ currentAttemptFormatDisplay }}</code>
</span>
</div>
<div
@@ -433,6 +433,7 @@ interface NodeGroup {
endIndex: number
hasConversion: boolean // 组内是否有格式转换候选
providerApiFormat: string | null // 提供商 API 格式（如 openai:cli）
isPoolGroup?: boolean
}
// 用量数据类型
@@ -464,8 +465,12 @@ const props = defineProps<{
requestId: string
/** 外部传入的状态码，用于覆盖 trace.final_status 的判断 */
overrideStatusCode?: number
/** 请求侧 API 格式（客户端入口格式） */
requestApiFormat?: string | null
/** 用量和费用数据 */
usageData?: UsageData | null
/** 请求元数据（用于号池调度组装） */
requestMetadata?: Record<string, unknown> | null
}>()
// 用量数据（从 props 获取）
@@ -594,21 +599,54 @@ const proxyTimingBreakdown = (proxy: Record<string, unknown>): string => {
return parts.join(' / ')
}
const TIMELINE_STATUS: CandidateRecord['status'][] = [
'success',
'failed',
'skipped',
'cancelled',
'pending',
'streaming',
'available',
'unused',
'stream_interrupted',
]
const STATUS_PRIORITY: Record<string, number> = {
available: 0,
unused: 0,
skipped: 1,
failed: 2,
cancelled: 2,
stream_interrupted: 2,
pending: 3,
streaming: 3,
success: 4,
}
const toInt = (value: unknown, defaultValue = 0): number => {
const num = Number(value)
return Number.isFinite(num) ? Math.trunc(num) : defaultValue
}
const makeAttemptKey = (candidateIndex: number, retryIndex: number): string => {
return `${candidateIndex}:${retryIndex}`
}
const normalizeTimelineStatus = (value: unknown): CandidateRecord['status'] => {
if (typeof value !== 'string') return 'failed'
const normalized = value.trim().toLowerCase()
if ((TIMELINE_STATUS as string[]).includes(normalized)) {
return normalized as CandidateRecord['status']
}
// 兜底：内部调度轨迹里非标准状态统一按失败展示
return 'failed'
}
// 候选时间线（按实际执行顺序排序）
const timeline = computed<CandidateRecord[]>(() => {
const rawTimeline = computed<CandidateRecord[]>(() => {
if (!trace.value) return []
return [...trace.value.candidates]
.filter(c => [
'success',
'failed',
'skipped',
'cancelled',
'pending',
'streaming',
'available',
'unused',
'stream_interrupted'
].includes(c.status))
.filter(c => TIMELINE_STATUS.includes(c.status))
.sort((a, b) => {
const startedA = a.started_at ? new Date(a.started_at).getTime() : Infinity
const startedB = b.started_at ? new Date(b.started_at).getTime() : Infinity
@@ -620,18 +658,109 @@ const timeline = computed<CandidateRecord[]>(() => {
})
})
// 将相同 Provider 的所有请求合并为组（同提供商的 Key 放在子节点）
const groupedTimeline = computed<NodeGroup[]>(() => {
if (!timeline.value || timeline.value.length === 0) return []
const schedulingAudit = computed<Record<string, unknown> | null>(() => {
const metadata = props.requestMetadata
if (!metadata || typeof metadata !== 'object' || Array.isArray(metadata)) return null
const raw = metadata.scheduling_audit
if (!raw || typeof raw !== 'object' || Array.isArray(raw)) return null
return raw as Record<string, unknown>
})
const poolAttemptCandidates = computed<CandidateRecord[]>(() => {
const audit = schedulingAudit.value
if (!audit) return []
const attempts = audit.attempts
if (!Array.isArray(attempts) || attempts.length === 0) return []
const traceMap = new Map<string, CandidateRecord>()
for (const candidate of rawTimeline.value) {
traceMap.set(makeAttemptKey(candidate.candidate_index, candidate.retry_index), candidate)
}
return attempts
.map((item, index) => {
if (!item || typeof item !== 'object' || Array.isArray(item)) return null
const raw = item as Record<string, unknown>
const candidateIndex = toInt(raw.candidate_index, index)
const retryIndex = toInt(raw.retry_index, 0)
const key = makeAttemptKey(candidateIndex, retryIndex)
const fromTrace = traceMap.get(key)
const merged: CandidateRecord = fromTrace
? { ...fromTrace }
: {
id: `pool-${props.requestId}-${candidateIndex}-${retryIndex}-${index}`,
request_id: props.requestId,
candidate_index: candidateIndex,
retry_index: retryIndex,
provider_id: undefined,
provider_name: undefined,
endpoint_id: undefined,
key_id: undefined,
key_name: undefined,
status: 'failed',
is_cached: false,
created_at: new Date(0).toISOString(),
}
merged.status = normalizeTimelineStatus(raw.status ?? merged.status)
if (typeof raw.provider_id === 'string') merged.provider_id = raw.provider_id
if (typeof raw.provider_name === 'string') merged.provider_name = raw.provider_name
if (typeof raw.endpoint_id === 'string') merged.endpoint_id = raw.endpoint_id
if (typeof raw.key_id === 'string') merged.key_id = raw.key_id
if (typeof raw.key_name === 'string') merged.key_name = raw.key_name
if (typeof raw.status_code === 'number') merged.status_code = raw.status_code
if (typeof raw.error_type === 'string') merged.error_type = raw.error_type
return merged
})
.filter((item): item is CandidateRecord => item !== null)
})
const poolAttemptKeySet = computed<Set<string>>(() => {
return new Set(
poolAttemptCandidates.value.map((item) => makeAttemptKey(item.candidate_index, item.retry_index)),
)
})
const timeline = computed<CandidateRecord[]>(() => {
if (poolAttemptCandidates.value.length === 0) return rawTimeline.value
return rawTimeline.value.filter(
(candidate) => !poolAttemptKeySet.value.has(makeAttemptKey(candidate.candidate_index, candidate.retry_index)),
)
})
const AUTH_TYPE_PROVIDER_LABEL_MAP: Record<string, string> = {
codex: 'Codex',
kiro: 'Kiro',
antigravity: 'Antigravity',
claude_code: 'Claude Code',
gemini_cli: 'Gemini CLI',
}
const normalizeProviderName = (value: string): string => {
const text = value.trim()
if (!text) return '未知'
// 管理后台中常见“xx反代”命名，展示时保留提供商品牌名即可
return text.replace(/反代$/u, '').trim() || text
}
const getProviderDisplayName = (attempt: CandidateRecord | null | undefined): string => {
if (!attempt) return '未知'
const authType = String(attempt.key_auth_type || '').trim().toLowerCase()
if (authType && AUTH_TYPE_PROVIDER_LABEL_MAP[authType]) {
return AUTH_TYPE_PROVIDER_LABEL_MAP[authType]
}
const providerName = String(attempt.provider_name || '').trim()
return providerName ? normalizeProviderName(providerName) : '未知'
}
const buildProviderGroups = (items: CandidateRecord[]): NodeGroup[] => {
const groups: NodeGroup[] = []
let currentGroup: NodeGroup | null = null
timeline.value.forEach((candidate, index) => {
// 使用 provider_name 作为分组 key（同一个提供商的所有 Key 合并）
items.forEach((candidate, index) => {
const providerKey = candidate.provider_name || '未知'
// 如果属于同一个 Provider，合并到当前组
if (currentGroup && currentGroup.id === providerKey) {
currentGroup.allAttempts.push(candidate)
currentGroup.retryCount++
@@ -640,37 +769,67 @@ const groupedTimeline = computed<NodeGroup[]>(() => {
if (candidate.extra_data?.needs_conversion) {
currentGroup.hasConversion = true
}
// 按优先级提升组状态：success > streaming/pending > failed/cancelled/stream_interrupted > skipped > available/unused
const statusPriority: Record<string, number> = {
available: 0, unused: 0, skipped: 1,
failed: 2, cancelled: 2, stream_interrupted: 2,
pending: 3, streaming: 3, success: 4,
}
const currentPriority = statusPriority[currentGroup.primaryStatus] ?? 0
const newPriority = statusPriority[candidate.status] ?? 0
const currentPriority = STATUS_PRIORITY[currentGroup.primaryStatus] ?? 0
const newPriority = STATUS_PRIORITY[candidate.status] ?? 0
if (newPriority > currentPriority) {
currentGroup.primaryStatus = candidate.status
}
} else {
// 新建一个组
currentGroup = {
id: providerKey,
providerName: candidate.provider_name || '未知',
primary: candidate,
primaryStatus: candidate.status,
allAttempts: [candidate],
retryCount: 0,
totalLatency: candidate.latency_ms || 0,
startIndex: index,
endIndex: index,
hasConversion: candidate.extra_data?.needs_conversion === true,
providerApiFormat: candidate.extra_data?.provider_api_format || null,
}
groups.push(currentGroup)
return
}
currentGroup = {
id: providerKey,
providerName: getProviderDisplayName(candidate),
primary: candidate,
primaryStatus: candidate.status,
allAttempts: [candidate],
retryCount: 0,
totalLatency: candidate.latency_ms || 0,
startIndex: index,
endIndex: index,
hasConversion: candidate.extra_data?.needs_conversion === true,
providerApiFormat: candidate.extra_data?.provider_api_format || null,
isPoolGroup: false,
}
groups.push(currentGroup)
})
return groups
}
// 将相同 Provider 的所有请求合并为组（同提供商的 Key 放在子节点）
const groupedTimeline = computed<NodeGroup[]>(() => {
const providerGroups = buildProviderGroups(timeline.value)
const poolAttempts = poolAttemptCandidates.value
if (poolAttempts.length === 0) {
return providerGroups
}
const poolPrimaryStatus = poolAttempts.reduce((best, current) => {
const bestPriority = STATUS_PRIORITY[best] ?? 0
const currentPriority = STATUS_PRIORITY[current.status] ?? 0
return currentPriority > bestPriority ? current.status : best
}, poolAttempts[0].status)
const successAttempt = poolAttempts.find((item) => item.status === 'success')
const poolPrimary = successAttempt || poolAttempts[poolAttempts.length - 1] || poolAttempts[0]
const poolGroup: NodeGroup = {
id: '__pool_group__',
providerName: getProviderDisplayName(poolPrimary),
primary: poolPrimary,
primaryStatus: poolPrimaryStatus,
allAttempts: poolAttempts,
retryCount: Math.max(0, poolAttempts.length - 1),
totalLatency: poolAttempts.reduce((sum, item) => sum + (item.latency_ms || 0), 0),
startIndex: 0,
endIndex: poolAttempts.length - 1,
hasConversion: poolAttempts.some((item) => item.extra_data?.needs_conversion === true),
providerApiFormat: null,
isPoolGroup: true,
}
return [poolGroup, ...providerGroups]
})
// 格式转换分界点索引（首个 hasConversion=true 的 group index）
@@ -687,16 +846,16 @@ const conversionBoundaryIndex = computed(() => {
// 优先使用 latency_ms，因为它与 Usage.response_time_ms 使用相同的时间基准
// 避免 finished_at - started_at 带来的额外延迟（数据库操作时间）
const totalTraceLatency = computed(() => {
if (!timeline.value || timeline.value.length === 0) return 0
if (!rawTimeline.value || rawTimeline.value.length === 0) return 0
// 查找成功的候选，使用其 latency_ms
const successCandidate = timeline.value.find(c => c.status === 'success')
const successCandidate = rawTimeline.value.find(c => c.status === 'success')
if (successCandidate?.latency_ms != null) {
return successCandidate.latency_ms
}
// 如果没有成功的候选，查找失败但有 latency_ms 的候选
const failedWithLatency = timeline.value.find(c => c.status === 'failed' && c.latency_ms != null)
const failedWithLatency = rawTimeline.value.find(c => c.status === 'failed' && c.latency_ms != null)
if (failedWithLatency?.latency_ms != null) {
return failedWithLatency.latency_ms
}
@@ -705,7 +864,7 @@ const totalTraceLatency = computed(() => {
let earliestStart: number | null = null
let latestEnd: number | null = null
for (const candidate of timeline.value) {
for (const candidate of rawTimeline.value) {
if (candidate.started_at) {
const startTime = new Date(candidate.started_at).getTime()
if (earliestStart === null || startTime < earliestStart) {
@@ -738,6 +897,49 @@ const currentAttempt = computed(() => {
return selectedGroup.value.allAttempts[selectedAttemptIndex.value] || selectedGroup.value.primary
})
const currentGroupTitle = computed(() => {
if (!selectedGroup.value || !currentAttempt.value) return ''
if (selectedGroup.value.isPoolGroup) {
return getProviderDisplayName(currentAttempt.value)
}
return selectedGroup.value.providerName
})
const normalizeFormatSignature = (value: string): string => {
return value.trim().toLowerCase()
}
const currentAttemptFormatDisplay = computed(() => {
const attempt = currentAttempt.value
if (!attempt) return ''
const extra = (
attempt.extra_data && typeof attempt.extra_data === 'object' && !Array.isArray(attempt.extra_data)
? attempt.extra_data
: {}
) as Record<string, unknown>
const providerRaw = typeof extra.provider_api_format === 'string' ? extra.provider_api_format : ''
const clientRawFromExtra = typeof extra.client_api_format === 'string' ? extra.client_api_format : ''
const requestRaw = clientRawFromExtra || (typeof props.requestApiFormat === 'string' ? props.requestApiFormat : '')
if (!providerRaw && !requestRaw) return ''
const providerText = providerRaw ? formatApiFormat(providerRaw) : ''
const requestText = requestRaw ? formatApiFormat(requestRaw) : ''
const convertedByFlag = extra.needs_conversion === true
const convertedByDiff = Boolean(
providerRaw &&
requestRaw &&
normalizeFormatSignature(providerRaw) !== normalizeFormatSignature(requestRaw),
)
if ((convertedByFlag || convertedByDiff) && requestText && providerText) {
return `${requestText} -> ${providerText}`
}
return providerText || requestText
})
// 计算当前尝试启用的能力标签（请求需要的能力）
const activeCapabilities = computed(() => {
if (!currentAttempt.value?.required_capabilities) return []
@@ -1018,6 +1220,28 @@ const getStatusColorClass = (status: string) => {
}
return classes[status] || 'status-available'
}
// 展示状态：进行中态优先（包括 started 但未 finished 的中间态），再按 HTTP 状态码兜底
const getDisplayStatus = (attempt: CandidateRecord | null | undefined): string => {
if (!attempt) return 'available'
const hasFinished = Boolean(attempt.finished_at)
const isExplicitPending = (attempt.status === 'pending' || attempt.status === 'streaming') && !hasFinished
const isImplicitPending = Boolean(
attempt.started_at &&
!hasFinished &&
!['failed', 'cancelled', 'skipped', 'stream_interrupted'].includes(attempt.status),
)
if (isExplicitPending || isImplicitPending) {
return 'pending'
}
const code = attempt.status_code
if (typeof code === 'number') {
if (code >= 200 && code < 300) return 'success'
if (code >= 400) return 'failed'
}
return attempt.status
}
</script>
<style scoped>
frontend/src/features/usage/components/RequestDetailDrawer.vue
+6 -50
View File
@@ -378,57 +378,14 @@
</div>
</Card>
<!-- 号池调度摘要 -->
<Card v-if="poolSummary">
<div class="p-3 sm:p-4">
<div class="text-xs text-muted-foreground mb-2 font-medium">
号池调度
</div>
<div class="flex items-center gap-3 flex-wrap text-sm">
<span class="font-mono">
{{ poolSummary.total_keys }} 候选
</span>
<span class="text-muted-foreground">|</span>
<span class="font-mono">
{{ poolSummary.attempted }} 尝试
</span>
<template v-if="poolSummary.skipped_cooldown > 0">
<span class="text-muted-foreground">|</span>
<span class="font-mono text-amber-600 dark:text-amber-400">
{{ poolSummary.skipped_cooldown }} 冷却跳过
</span>
</template>
<template v-if="poolSummary.skipped_cost > 0">
<span class="text-muted-foreground">|</span>
<span class="font-mono text-orange-600 dark:text-orange-400">
{{ poolSummary.skipped_cost }} 成本跳过
</span>
</template>
<template v-if="poolSummary.sticky_session">
<span class="text-muted-foreground">|</span>
<Badge
variant="outline"
class="text-[10px] px-1.5 py-0 h-4"
>
粘性会话
</Badge>
</template>
<template v-if="poolSummary.success_reason">
<span class="text-muted-foreground">|</span>
<span class="text-xs text-muted-foreground">
{{ poolSummary.success_reason }}
</span>
</template>
</div>
</div>
</Card>
<!-- 请求链路追踪卡片 -->
<div v-if="detail.request_id || detail.id">
<HorizontalRequestTimeline
ref="timelineRef"
:request-id="detail.request_id || detail.id"
:override-status-code="detail.status_code"
:request-api-format="detail.api_format || null"
:request-metadata="traceRequestMetadata"
/>
</div>
@@ -781,11 +738,10 @@ const isDark = computed(() => {
return document.documentElement.classList.contains('dark')
})
// 号池调度摘要
const poolSummary = computed(() => {
const ps = detail.value?.metadata?.pool_summary as Record<string, unknown> | undefined
if (!ps || !ps.enabled) return null
return ps
const traceRequestMetadata = computed<Record<string, unknown> | null>(() => {
const meta = detail.value?.metadata
if (!meta || typeof meta !== 'object' || Array.isArray(meta)) return null
return meta as Record<string, unknown>
})
// 检测是否有提供商请求头
frontend/src/views/admin/PoolManagement.vue
+1313 -299
View File
File diff suppressed because it is too large Load Diff
src/api/admin/monitoring/trace.py
+9 -3
View File
@@ -163,13 +163,19 @@ class AdminGetRequestTraceAdapter(AdminApiAdapter):
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
db = context.db
# 只查询 candidates
candidates = RequestCandidateService.get_candidates_by_request_id(db, self.request_id)
# 查询所有候选后，默认只展示已发生调度结果的子集：
# - 过滤 available/unused（预创建但未实际参与本次调度）
# - 若过滤后为空（例如请求尚未开始），回退到全量，避免前端空白
all_candidates = RequestCandidateService.get_candidates_by_request_id(db, self.request_id)
# 如果没有数据，返回 404
if not candidates:
if not all_candidates:
raise HTTPException(status_code=404, detail="Request not found")
candidates = [
c for c in all_candidates if c.status not in ("available", "unused")
] or all_candidates
# 计算总延迟（只统计已完成的候选：success, failed, cancelled）
# 使用显式的 is not None 检查，避免过滤掉 0ms 的快速响应
total_latency = sum(
src/api/admin/pool/routes.py
+541 -21
View File
@@ -9,12 +9,14 @@ Provides endpoints for managing account pools at scale:
from __future__ import annotations
import asyncio
import json
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any
from fastapi import APIRouter, Depends, Query, Request
from sqlalchemy import func
from sqlalchemy.orm import Session
from src.api.base.admin_adapter import AdminApiAdapter
@@ -24,9 +26,14 @@ from src.core.crypto import crypto_service
from src.core.exceptions import NotFoundException
from src.core.logger import logger
from src.database import get_db
from src.models.database import Provider, ProviderAPIKey
from src.models.database import Provider, ProviderAPIKey, Usage
from src.services.provider.pool import redis_ops as pool_redis
from src.services.provider.pool.config import parse_pool_config
from src.services.provider.pool.scheduling_dimensions import (
PoolSchedulingSnapshot,
evaluate_pool_scheduling_dimensions,
summarize_pool_scheduling_dimensions,
)
from .schemas import (
BatchActionRequest,
@@ -38,6 +45,8 @@ from .schemas import (
PoolKeysPageResponse,
PoolOverviewItem,
PoolOverviewResponse,
PoolSchedulingDimension,
PoolSchedulingReason,
)
router = APIRouter(prefix="/api/admin/pool", tags=["pool-management"])
@@ -108,6 +117,33 @@ async def batch_import_keys(
ALLOWED_ACTIONS = {"enable", "disable", "delete", "clear_cooldown", "reset_cost"}
_COOLDOWN_REASON_LABELS: dict[str, str] = {
"rate_limited_429": "429 限流",
"forbidden_403": "403 禁止",
"overloaded_529": "529 过载",
"auth_failed_401": "401 认证失败",
"payment_required_402": "402 欠费",
"server_error_500": "500 错误",
}
_ACCOUNT_BLOCK_REASON_KEYWORDS: tuple[str, ...] = (
"account_block",
"account blocked",
"account has been disabled",
"account disabled",
"organization has been disabled",
"organization_disabled",
"validation_required",
"verify your account",
"forbidden",
"suspended",
"封禁",
"封号",
"被封",
"访问被禁止",
"账号异常",
)
def _to_float(value: Any) -> float | None:
if isinstance(value, bool):
@@ -125,6 +161,67 @@ def _to_float(value: Any) -> float | None:
return None
def _is_truthy_flag(value: Any) -> bool:
if isinstance(value, bool):
return value
if isinstance(value, (int, float)):
return value != 0
if isinstance(value, str):
normalized = value.strip().lower()
return normalized in {"1", "true", "yes", "y"}
return False
def _is_known_banned_reason(reason: str | None) -> bool:
if not reason:
return False
text = str(reason).strip()
if not text:
return False
lowered = text.lower()
# 结构化账号级别封禁标记（如 [ACCOUNT_BLOCK] ...）
try:
from src.services.provider.oauth_token import is_account_level_block
if is_account_level_block(text):
return True
except Exception:
pass
return any(keyword in lowered for keyword in _ACCOUNT_BLOCK_REASON_KEYWORDS)
def _is_known_banned_key(key: ProviderAPIKey, provider_type: str) -> bool:
upstream_metadata = getattr(key, "upstream_metadata", None)
normalized_provider = provider_type.strip().lower()
provider_bucket: dict[str, Any] | None = None
if isinstance(upstream_metadata, dict):
maybe_bucket = upstream_metadata.get(normalized_provider)
if isinstance(maybe_bucket, dict):
provider_bucket = maybe_bucket
if normalized_provider == "kiro" and provider_bucket:
if _is_truthy_flag(provider_bucket.get("is_banned")):
return True
if normalized_provider == "antigravity" and provider_bucket:
if _is_truthy_flag(provider_bucket.get("is_forbidden")):
return True
for source in (provider_bucket, upstream_metadata):
if not isinstance(source, dict):
continue
if _is_truthy_flag(source.get("is_banned")):
return True
if _is_truthy_flag(source.get("is_forbidden")):
return True
if _is_truthy_flag(source.get("account_disabled")):
return True
return _is_known_banned_reason(getattr(key, "oauth_invalid_reason", None))
def _format_percent(value: float) -> str:
clamped = max(0.0, min(value, 100.0))
return f"{clamped:.1f}%"
@@ -177,11 +274,7 @@ def _build_kiro_account_quota(upstream_metadata: dict[str, Any]) -> str | None:
remaining = 100.0 - usage_percentage
current_usage = _to_float(kiro.get("current_usage"))
usage_limit = _to_float(kiro.get("usage_limit"))
if (
current_usage is not None
and usage_limit is not None
and usage_limit > 0
):
if current_usage is not None and usage_limit is not None and usage_limit > 0:
return (
f"剩余 {_format_percent(remaining)} "
f"({_format_quota_value(current_usage)}/{_format_quota_value(usage_limit)})"
@@ -247,6 +340,215 @@ def _build_account_quota(provider_type: str, upstream_metadata: Any) -> str | No
return None
def _extract_quota_updated_at(provider_type: str, upstream_metadata: Any) -> int | None:
if not isinstance(upstream_metadata, dict):
return None
normalized_type = provider_type.strip().lower()
if normalized_type == "codex":
source = upstream_metadata.get("codex")
elif normalized_type == "antigravity":
source = upstream_metadata.get("antigravity")
elif normalized_type == "kiro":
source = upstream_metadata.get("kiro")
else:
return None
if not isinstance(source, dict):
return None
updated_at = _to_float(source.get("updated_at"))
if updated_at is None or updated_at <= 0:
return None
# 部分上游可能返回毫秒时间戳，统一转换为秒
if updated_at > 1_000_000_000_000:
updated_at /= 1000
return int(updated_at)
def _normalize_oauth_plan_type(plan_type: Any, provider_type: str) -> str | None:
if not isinstance(plan_type, str):
return None
text = plan_type.strip()
if not text:
return None
ptype = provider_type.strip().lower()
if ptype and text.lower().startswith(ptype):
trimmed = text[len(ptype) :].strip(" :-_")
if trimmed:
text = trimmed
return text or None
def _derive_oauth_plan_type(key: ProviderAPIKey, provider_type: str) -> str | None:
# Prefer persisted normalized field
persisted = _normalize_oauth_plan_type(getattr(key, "oauth_plan_type", None), provider_type)
if persisted:
return persisted
if str(getattr(key, "auth_type", "") or "").strip().lower() != "oauth":
return None
# Fallback 1: encrypted auth_config (common for Codex/Antigravity)
auth_config_raw = getattr(key, "auth_config", None)
if auth_config_raw:
try:
decrypted = crypto_service.decrypt(auth_config_raw)
auth_config = json.loads(decrypted)
if isinstance(auth_config, dict):
for plan_key in ("plan_type", "tier", "plan", "subscription_plan"):
normalized = _normalize_oauth_plan_type(
auth_config.get(plan_key), provider_type
)
if normalized:
return normalized
except Exception:
pass
# Fallback 2: upstream_metadata
upstream_metadata = getattr(key, "upstream_metadata", None)
if not isinstance(upstream_metadata, dict):
return None
provider_bucket = upstream_metadata.get(provider_type.strip().lower())
candidates: list[dict[str, Any]] = []
if isinstance(provider_bucket, dict):
candidates.append(provider_bucket)
candidates.append(upstream_metadata)
for source in candidates:
for plan_key in ("plan_type", "tier", "subscription_title", "subscription_plan"):
normalized = _normalize_oauth_plan_type(source.get(plan_key), provider_type)
if normalized:
return normalized
return None
def _compute_health_aggregate(
health_by_format: Any, circuit_breaker_by_format: Any
) -> tuple[float, bool]:
"""从按格式健康数据聚合出列表展示字段。"""
health_map = health_by_format if isinstance(health_by_format, dict) else {}
circuit_map = circuit_breaker_by_format if isinstance(circuit_breaker_by_format, dict) else {}
if health_map:
scores = [
float(item.get("health_score") or 1.0)
for item in health_map.values()
if isinstance(item, dict)
]
health_score = min(scores) if scores else 1.0
else:
health_score = 1.0
any_circuit_open = any(
bool(item.get("open", False)) for item in circuit_map.values() if isinstance(item, dict)
)
return health_score, any_circuit_open
def _format_cooldown_detail(raw: str | None) -> str | None:
if not raw:
return None
return _COOLDOWN_REASON_LABELS.get(raw, raw)
def _build_pool_scheduling_state(
*,
is_active: bool,
cooldown_reason: str | None,
cooldown_ttl_seconds: int | None,
circuit_breaker_open: bool,
cost_window_usage: int,
cost_limit: int | None,
cost_soft_threshold_percent: int,
health_score: float,
) -> tuple[
str,
str,
str,
list[PoolSchedulingReason],
float,
bool,
int,
int,
list[PoolSchedulingDimension],
]:
"""Build unified scheduling state for frontend display."""
snapshot = PoolSchedulingSnapshot(
is_active=is_active,
cooldown_reason=cooldown_reason,
cooldown_ttl_seconds=cooldown_ttl_seconds,
circuit_breaker_open=circuit_breaker_open,
cost_window_usage=cost_window_usage,
cost_limit=cost_limit,
cost_soft_threshold_percent=cost_soft_threshold_percent,
health_score=health_score,
)
dimensions_raw = evaluate_pool_scheduling_dimensions(snapshot)
summary = summarize_pool_scheduling_dimensions(dimensions_raw)
scheduling_dimensions: list[PoolSchedulingDimension] = []
scheduling_reasons: list[PoolSchedulingReason] = []
for item in dimensions_raw:
detail = item.detail
if item.code == "cooldown":
detail = _format_cooldown_detail(detail)
model = PoolSchedulingDimension(
code=item.code,
label=item.label,
status=item.status,
blocking=bool(item.blocking or item.status == "blocked"),
source=item.source,
weight=item.weight,
score=item.score,
ttl_seconds=item.ttl_seconds,
detail=detail,
)
scheduling_dimensions.append(model)
if item.status != "ok":
scheduling_reasons.append(
PoolSchedulingReason(
code=item.code,
label=item.label,
blocking=bool(item.blocking or item.status == "blocked"),
source=item.source,
ttl_seconds=item.ttl_seconds,
detail=detail,
)
)
return (
summary.status,
summary.reason,
summary.label,
scheduling_reasons,
summary.score,
summary.candidate_eligible,
summary.blocked_count,
summary.degraded_count,
scheduling_dimensions,
)
def _mask_proxy_password(proxy_config: Any) -> dict[str, Any] | None:
if not isinstance(proxy_config, dict):
return None
masked = dict(proxy_config)
password = masked.get("password")
if isinstance(password, str) and password:
masked["password"] = "******"
return masked
@router.post("/{provider_id}/keys/batch-action", response_model=BatchActionResponse)
async def batch_action_keys(
provider_id: str,
@@ -259,6 +561,17 @@ async def batch_action_keys(
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
@router.post("/{provider_id}/keys/cleanup-banned", response_model=BatchActionResponse)
async def cleanup_banned_keys(
provider_id: str,
request: Request,
db: Session = Depends(get_db),
) -> BatchActionResponse:
"""Delete known banned/suspended accounts for the provider."""
adapter = AdminCleanupBannedKeysAdapter(provider_id=provider_id)
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
# ---------------------------------------------------------------------------
# Adapters
# ---------------------------------------------------------------------------
@@ -380,29 +693,135 @@ class AdminListPoolKeysAdapter(AdminApiAdapter):
if pcfg
else asyncio.sleep(0, result={})
)
cooldowns, cooldown_ttls, lru_scores, cost_totals = await asyncio.gather(
cooldowns, cooldown_ttls, lru_scores, cost_totals, sticky_counts = await asyncio.gather(
pool_redis.batch_get_cooldowns(pid, key_ids),
pool_redis.batch_get_cooldown_ttls(pid, key_ids),
_lru_coro,
_cost_coro,
pool_redis.batch_get_key_sticky_counts(pid, key_ids),
)
else:
cooldowns, cooldown_ttls, lru_scores, cost_totals = {}, {}, {}, {}
# Sticky session count per key is expensive (SCAN+MGET per key).
# Only compute when the page is small enough to avoid timeout.
sticky_counts: dict[str, int] = {}
if key_ids and len(key_ids) <= 30:
counts = await asyncio.gather(
*(pool_redis.get_key_sticky_count(pid, kid) for kid in key_ids)
cooldowns, cooldown_ttls, lru_scores, cost_totals, sticky_counts = (
{},
{},
{},
{},
{},
)
sticky_counts = dict(zip(key_ids, counts))
usage_stats_by_key: dict[str, dict[str, Any]] = {}
if key_ids:
usage_rows = (
db.query(
Usage.provider_api_key_id.label("key_id"),
func.count(Usage.id).label("request_count"),
func.coalesce(func.sum(Usage.total_tokens), 0).label("total_tokens"),
func.coalesce(func.sum(Usage.total_cost_usd), 0.0).label("total_cost_usd"),
func.max(Usage.created_at).label("last_used_at"),
)
.filter(
Usage.provider_id == pid,
Usage.provider_api_key_id.in_(key_ids),
Usage.status.notin_(["pending", "streaming"]),
)
.group_by(Usage.provider_api_key_id)
.all()
)
usage_stats_by_key = {
str(row.key_id): {
"request_count": int(row.request_count or 0),
"total_tokens": int(row.total_tokens or 0),
"total_cost_usd": float(row.total_cost_usd or 0.0),
"last_used_at": getattr(row, "last_used_at", None),
}
for row in usage_rows
if getattr(row, "key_id", None)
}
key_details: list[PoolKeyDetail] = []
for k in keys:
kid = str(k.id)
cd_reason = cooldowns.get(kid)
cd_ttl = cooldown_ttls.get(kid) if cd_reason else None
health_score, any_circuit_open = _compute_health_aggregate(
getattr(k, "health_by_format", None),
getattr(k, "circuit_breaker_by_format", None),
)
cost_usage = int(cost_totals.get(kid, 0) or 0)
cost_limit = pcfg.cost_limit_per_key_tokens if pcfg else None
(
scheduling_status,
scheduling_reason,
scheduling_label,
scheduling_reasons,
scheduling_score,
candidate_eligible,
scheduling_blocked_count,
scheduling_degraded_count,
scheduling_dimensions,
) = _build_pool_scheduling_state(
is_active=bool(k.is_active),
cooldown_reason=cd_reason,
cooldown_ttl_seconds=cd_ttl,
circuit_breaker_open=any_circuit_open,
cost_window_usage=cost_usage,
cost_limit=cost_limit,
cost_soft_threshold_percent=(pcfg.cost_soft_threshold_percent if pcfg else 80),
health_score=health_score,
)
raw_allowed_models = getattr(k, "allowed_models", None)
allowed_models = (
[str(item) for item in raw_allowed_models]
if isinstance(raw_allowed_models, list)
else None
)
raw_locked_models = getattr(k, "locked_models", None)
locked_models = (
[str(item) for item in raw_locked_models]
if isinstance(raw_locked_models, list)
else None
)
raw_include_patterns = getattr(k, "model_include_patterns", None)
include_patterns = (
[str(item) for item in raw_include_patterns]
if isinstance(raw_include_patterns, list)
else None
)
raw_exclude_patterns = getattr(k, "model_exclude_patterns", None)
exclude_patterns = (
[str(item) for item in raw_exclude_patterns]
if isinstance(raw_exclude_patterns, list)
else None
)
capabilities = (
{str(name): bool(enabled) for name, enabled in k.capabilities.items()}
if isinstance(getattr(k, "capabilities", None), dict)
else None
)
rate_multipliers: dict[str, float] | None = None
if isinstance(getattr(k, "rate_multipliers", None), dict):
converted: dict[str, float] = {}
for fmt, raw_val in k.rate_multipliers.items():
num_val = _to_float(raw_val)
if num_val is None:
continue
converted[str(fmt)] = num_val
rate_multipliers = converted or None
api_formats = (
[str(fmt) for fmt in getattr(k, "api_formats", []) if isinstance(fmt, str)]
if isinstance(getattr(k, "api_formats", None), list)
else []
)
key_usage_stats = usage_stats_by_key.get(kid, {})
key_request_count = int(
key_usage_stats.get("request_count") or getattr(k, "request_count", 0) or 0
)
key_total_tokens = int(key_usage_stats.get("total_tokens") or 0)
key_total_cost_usd = float(key_usage_stats.get("total_cost_usd") or 0.0)
key_last_used_at = getattr(k, "last_used_at", None) or key_usage_stats.get(
"last_used_at"
)
key_details.append(
PoolKeyDetail(
@@ -410,22 +829,66 @@ class AdminListPoolKeysAdapter(AdminApiAdapter):
key_name=k.name or "",
is_active=bool(k.is_active),
auth_type=str(getattr(k, "auth_type", "api_key") or "api_key"),
oauth_expires_at=(
int(k.oauth_expires_at.timestamp())
if getattr(k, "oauth_expires_at", None)
else None
),
oauth_invalid_at=(
int(k.oauth_invalid_at.timestamp())
if getattr(k, "oauth_invalid_at", None)
else None
),
oauth_invalid_reason=getattr(k, "oauth_invalid_reason", None),
oauth_plan_type=_derive_oauth_plan_type(k, provider_type),
quota_updated_at=_extract_quota_updated_at(
provider_type,
getattr(k, "upstream_metadata", None),
),
health_score=health_score,
circuit_breaker_open=any_circuit_open,
api_formats=api_formats,
rate_multipliers=rate_multipliers,
internal_priority=int(getattr(k, "internal_priority", 50) or 50),
rpm_limit=getattr(k, "rpm_limit", None),
cache_ttl_minutes=int(getattr(k, "cache_ttl_minutes", 5) or 5),
max_probe_interval_minutes=int(
getattr(k, "max_probe_interval_minutes", 32) or 32
),
note=getattr(k, "note", None),
allowed_models=allowed_models,
capabilities=capabilities,
auto_fetch_models=bool(getattr(k, "auto_fetch_models", False)),
locked_models=locked_models,
model_include_patterns=include_patterns,
model_exclude_patterns=exclude_patterns,
proxy=_mask_proxy_password(getattr(k, "proxy", None)),
account_quota=_build_account_quota(
provider_type,
getattr(k, "upstream_metadata", None),
),
cooldown_reason=cd_reason,
cooldown_ttl_seconds=cd_ttl,
cost_window_usage=cost_totals.get(kid, 0),
cost_limit=pcfg.cost_limit_per_key_tokens if pcfg else None,
cost_window_usage=cost_usage,
cost_limit=cost_limit,
request_count=key_request_count,
total_tokens=key_total_tokens,
total_cost_usd=key_total_cost_usd,
sticky_sessions=sticky_counts.get(kid, 0),
lru_score=lru_scores.get(kid),
created_at=(
k.created_at.isoformat() if getattr(k, "created_at", None) else None
),
last_used_at=(
k.last_used_at.isoformat() if getattr(k, "last_used_at", None) else None
),
last_used_at=(key_last_used_at.isoformat() if key_last_used_at else None),
scheduling_status=scheduling_status,
scheduling_reason=scheduling_reason,
scheduling_label=scheduling_label,
scheduling_reasons=scheduling_reasons,
scheduling_score=scheduling_score,
candidate_eligible=candidate_eligible,
scheduling_blocked_count=scheduling_blocked_count,
scheduling_degraded_count=scheduling_degraded_count,
scheduling_dimensions=scheduling_dimensions,
)
)
@@ -447,6 +910,9 @@ class AdminBatchImportKeysAdapter(AdminApiAdapter):
provider = db.query(Provider).filter(Provider.id == self.provider_id).first()
if not provider:
raise NotFoundException("Provider not found", "provider")
key_proxy: dict[str, Any] | None = None
if self.body.proxy_node_id and self.body.proxy_node_id.strip():
key_proxy = {"node_id": self.body.proxy_node_id.strip(), "enabled": True}
imported = 0
skipped = 0
@@ -466,6 +932,7 @@ class AdminBatchImportKeysAdapter(AdminApiAdapter):
name=item.name or f"imported-{idx}",
api_key=encrypted_key,
auth_type=item.auth_type or "api_key",
proxy=key_proxy,
is_active=True,
created_at=now,
updated_at=now,
@@ -591,3 +1058,56 @@ class AdminBatchActionKeysAdapter(AdminApiAdapter):
affected=affected,
message=f"{affected} keys {action_labels.get(self.body.action, self.body.action)}",
)
@dataclass
class AdminCleanupBannedKeysAdapter(AdminApiAdapter):
provider_id: str = ""
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
db = context.db
provider = db.query(Provider).filter(Provider.id == self.provider_id).first()
if not provider:
raise NotFoundException("Provider not found", "provider")
pid = str(provider.id)
provider_type = str(getattr(provider, "provider_type", "") or "").strip().lower()
keys = db.query(ProviderAPIKey).filter(ProviderAPIKey.provider_id == pid).all()
banned_keys = [key for key in keys if _is_known_banned_key(key, provider_type)]
if not banned_keys:
return BatchActionResponse(affected=0, message="未发现已知封号账号")
banned_key_ids = [str(key.id) for key in banned_keys]
for key in banned_keys:
db.delete(key)
try:
db.commit()
except Exception as exc:
db.rollback()
logger.error("cleanup banned keys commit failed: {}", exc)
return BatchActionResponse(affected=0, message=f"commit failed: {exc}")
# 清理 Redis 中可能残留的状态，避免删除后仍有旧状态占用资源。
cleanup_coros = []
for kid in banned_key_ids:
cleanup_coros.append(pool_redis.clear_cooldown(pid, kid))
cleanup_coros.append(pool_redis.clear_cost(pid, kid))
if cleanup_coros:
await asyncio.gather(*cleanup_coros, return_exceptions=True)
admin_name = context.user.username if context.user else "admin"
logger.warning(
"Pool cleanup banned by {}: provider={}, affected={}, key_ids={}",
admin_name,
self.provider_id[:8],
len(banned_key_ids),
[kid[:8] for kid in banned_key_ids],
)
return BatchActionResponse(
affected=len(banned_key_ids),
message=f"已清理 {len(banned_key_ids)} 个已知封号账号",
)
src/api/admin/pool/schemas.py
+66
View File
@@ -2,6 +2,8 @@
from __future__ import annotations
from typing import Any
from pydantic import BaseModel, ConfigDict, Field
# ---------------------------------------------------------------------------
@@ -32,6 +34,31 @@ class PoolOverviewResponse(BaseModel):
# ---------------------------------------------------------------------------
class PoolSchedulingReason(BaseModel):
"""Structured scheduling reason for a key."""
code: str
label: str
blocking: bool = False
source: str = "pool" # manual / pool / health / policy
ttl_seconds: int | None = None
detail: str | None = None
class PoolSchedulingDimension(BaseModel):
"""Detailed scheduling dimension status."""
code: str
label: str
status: str = "ok" # ok / degraded / blocked
blocking: bool = False
source: str = "pool"
weight: int = 1
score: float = 1.0 # normalized 0~1
ttl_seconds: int | None = None
detail: str | None = None
class PoolKeyDetail(BaseModel):
"""Detailed status of a single pool key."""
@@ -39,15 +66,50 @@ class PoolKeyDetail(BaseModel):
key_name: str
is_active: bool
auth_type: str = "api_key"
oauth_expires_at: int | None = None
oauth_invalid_at: int | None = None
oauth_invalid_reason: str | None = None
oauth_plan_type: str | None = None
quota_updated_at: int | None = None
# 健康度聚合字段（与 Provider Key 列表口径一致）
health_score: float = 1.0
circuit_breaker_open: bool = False
# 编辑/权限/代理所需字段
api_formats: list[str] = Field(default_factory=list)
rate_multipliers: dict[str, float] | None = None
internal_priority: int = 50
rpm_limit: int | None = None
cache_ttl_minutes: int = 5
max_probe_interval_minutes: int = 32
note: str | None = None
allowed_models: list[str] | None = None
capabilities: dict[str, bool] | None = None
auto_fetch_models: bool = False
locked_models: list[str] | None = None
model_include_patterns: list[str] | None = None
model_exclude_patterns: list[str] | None = None
proxy: dict[str, Any] | None = None
account_quota: str | None = None
cooldown_reason: str | None = None
cooldown_ttl_seconds: int | None = None
cost_window_usage: int = 0
cost_limit: int | None = None
request_count: int = 0
total_tokens: int = 0
total_cost_usd: float = 0.0
sticky_sessions: int = 0
lru_score: float | None = None
created_at: str | None = None
last_used_at: str | None = None
scheduling_status: str = "available" # available / degraded / blocked
scheduling_reason: str = "available"
scheduling_label: str = "可用"
scheduling_reasons: list[PoolSchedulingReason] = Field(default_factory=list)
scheduling_score: float = 100.0
candidate_eligible: bool = True
scheduling_blocked_count: int = 0
scheduling_degraded_count: int = 0
scheduling_dimensions: list[PoolSchedulingDimension] = Field(default_factory=list)
model_config = ConfigDict(from_attributes=True)
@@ -76,6 +138,10 @@ class PoolKeyImportItem(BaseModel):
class BatchImportRequest(BaseModel):
keys: list[PoolKeyImportItem] = Field(..., max_length=500)
proxy_node_id: str | None = Field(
default=None,
description="导入时绑定到账号的代理节点 ID（可选）",
)
class BatchImportError(BaseModel):
src/api/admin/provider_oauth.py
+61 -9
View File
@@ -1368,6 +1368,61 @@ def _parse_kiro_import_input(raw_input: str) -> list[dict[str, Any]]:
返回: 凭据字典列表
"""
def _normalize_item(item: Any) -> dict[str, Any] | None:
"""规范化单条 Kiro 导入项，兼容导出结构。"""
if isinstance(item, str) and item.strip():
return {"refreshToken": item.strip()}
if not isinstance(item, dict):
return None
nested = item.get("auth_config") or item.get("authConfig")
if isinstance(nested, dict):
# 优先使用 auth_config，兼容导出对象形态：
# {"name": "...", "auth_config": {...}, ...}
merged = dict(nested)
# 若顶层也包含关键字段，允许覆盖 nested（便于手工修正）
for key in (
"provider_type",
"providerType",
"auth_method",
"authMethod",
"auth_type",
"authType",
"refresh_token",
"refreshToken",
"expires_at",
"expiresAt",
"profile_arn",
"profileArn",
"region",
"auth_region",
"authRegion",
"api_region",
"apiRegion",
"client_id",
"clientId",
"client_secret",
"clientSecret",
"machine_id",
"machineId",
"kiro_version",
"kiroVersion",
"system_version",
"systemVersion",
"node_version",
"nodeVersion",
"email",
"access_token",
"accessToken",
):
value = item.get(key)
if value is not None and value != "":
merged[key] = value
return merged
return item
raw = raw_input.strip()
if not raw:
return []
@@ -1380,18 +1435,15 @@ def _parse_kiro_import_input(raw_input: str) -> list[dict[str, Any]]:
if isinstance(parsed, list):
result: list[dict[str, Any]] = []
for item in parsed:
if isinstance(item, dict):
result.append(item)
elif isinstance(item, str) and item.strip():
result.append({"refreshToken": item.strip()})
normalized = _normalize_item(item)
if normalized:
result.append(normalized)
return result
if isinstance(parsed, dict):
# 兼容嵌套格式: {"auth_config": {...}} / {"authConfig": {...}}
nested = parsed.get("auth_config") or parsed.get("authConfig")
if isinstance(nested, dict):
return [nested]
return [parsed]
normalized = _normalize_item(parsed)
if normalized:
return [normalized]
except json.JSONDecodeError:
pass
src/api/handlers/base/base_handler.py
+278
View File
@@ -145,6 +145,284 @@ class BaseMessageHandler:
return None
return {"perf": self.perf_metrics}
@staticmethod
def _normalize_candidate_status(candidate: dict[str, Any]) -> str:
status = candidate.get("status")
if isinstance(status, str) and status.strip():
return status.strip().lower()
attempt_status = candidate.get("attempt_status")
if isinstance(attempt_status, str) and attempt_status.strip():
return attempt_status.strip().lower()
if candidate.get("skipped"):
return "skipped"
return ""
@staticmethod
def _to_int(value: Any, default: int = 0) -> int:
try:
return int(value)
except Exception:
return default
def _load_request_candidate_keys(self) -> list[Any]:
if not self.request_id:
return []
try:
from src.services.candidate.recorder import CandidateRecorder
return CandidateRecorder(self.db).get_candidate_keys(self.request_id)
except Exception:
return []
def _compact_candidate_key_snapshot(self, item: Any) -> dict[str, Any] | None:
raw: dict[str, Any] | None = None
if isinstance(item, dict):
raw = dict(item)
elif hasattr(item, "to_dict"):
try:
converted = item.to_dict()
if isinstance(converted, dict):
raw = dict(converted)
except Exception:
raw = None
if raw is None:
return None
status = self._normalize_candidate_status(raw)
candidate_index = raw.get("candidate_index", raw.get("index", 0))
retry_index = raw.get("retry_index", 0)
snapshot: dict[str, Any] = {
"candidate_index": self._to_int(candidate_index, 0),
"retry_index": self._to_int(retry_index, 0),
}
passthrough_fields = (
"provider_id",
"provider_name",
"endpoint_id",
"key_id",
"key_name",
"auth_type",
"priority",
"is_cached",
"skip_reason",
"error_type",
"status_code",
"latency_ms",
)
for field in passthrough_fields:
value = raw.get(field)
if value is not None and value != "":
snapshot[field] = value
if status:
snapshot["status"] = status
if raw.get("skipped"):
snapshot["skipped"] = True
snapshot.setdefault("status", "skipped")
if "selected" in raw:
snapshot["selected"] = bool(raw.get("selected"))
error_message = raw.get("error_message")
if isinstance(error_message, str) and error_message:
snapshot["error_message"] = error_message[:240]
return snapshot
def _collect_candidate_snapshots(
self,
*,
candidate_keys: list[Any] | None = None,
fallback_from_request: bool = False,
) -> list[dict[str, Any]]:
source = candidate_keys
if (not source) and fallback_from_request:
source = self._load_request_candidate_keys()
snapshots: list[dict[str, Any]] = []
for item in source or []:
snapshot = self._compact_candidate_key_snapshot(item)
if snapshot:
snapshots.append(snapshot)
snapshots.sort(
key=lambda it: (
self._to_int(it.get("candidate_index"), 0),
self._to_int(it.get("retry_index"), 0),
)
)
return snapshots[:64]
def _build_scheduling_audit(
self,
snapshots: list[dict[str, Any]],
*,
selected_key_id: str | None = None,
) -> dict[str, Any] | None:
if not snapshots:
return None
# "unused" means the candidate was pre-created for audit but never actually attempted.
executed_status_exclude = {"", "available", "pending", "skipped", "unused"}
executed_count = 0
attempts: list[dict[str, Any]] = []
account_map: dict[str, dict[str, Any]] = {}
candidate_indices: set[int] = set()
key_ids: set[str] = set()
for snapshot in snapshots:
status = str(snapshot.get("status", "") or "").lower()
if status in executed_status_exclude:
continue
executed_count += 1
candidate_index = self._to_int(snapshot.get("candidate_index"), 0)
retry_index = self._to_int(snapshot.get("retry_index"), 0)
key_id = snapshot.get("key_id")
key_name = snapshot.get("key_name")
provider_id = snapshot.get("provider_id")
provider_name = snapshot.get("provider_name")
candidate_indices.add(candidate_index)
if isinstance(key_id, str) and key_id:
key_ids.add(key_id)
if len(attempts) < 24:
attempts.append(
{
"candidate_index": candidate_index,
"retry_index": retry_index,
"provider_id": provider_id,
"provider_name": provider_name,
"key_id": key_id,
"key_name": key_name,
"status": status,
"status_code": snapshot.get("status_code"),
"error_type": snapshot.get("error_type"),
}
)
if not isinstance(key_id, str) or not key_id:
continue
account = account_map.get(key_id)
if account is None:
account = {
"key_id": key_id,
"key_name": key_name,
"provider_id": provider_id,
"provider_name": provider_name,
"attempts": 0,
"successes": 0,
"last_status": status,
}
account_map[key_id] = account
account["attempts"] = self._to_int(account.get("attempts"), 0) + 1
if status in {"success", "streaming"}:
account["successes"] = self._to_int(account.get("successes"), 0) + 1
account["last_status"] = status
if executed_count == 0:
return {
"mode": "internal",
"attempted_count": 0,
"account_count": 0,
"retry_occurred": False,
"failover_occurred": False,
"accounts": [],
"attempts": [],
}
selected_key_id_norm = str(selected_key_id) if selected_key_id else None
accounts = list(account_map.values())[:12]
selected_account: dict[str, Any] | None = None
if selected_key_id_norm and selected_key_id_norm in account_map:
selected_account = dict(account_map[selected_key_id_norm])
else:
for account in account_map.values():
if self._to_int(account.get("successes"), 0) > 0:
selected_account = dict(account)
selected_key_id_norm = str(account.get("key_id", ""))
break
if selected_account is not None:
for account in accounts:
if account.get("key_id") == selected_account.get("key_id"):
account["selected"] = True
failover_occurred = executed_count > 1 and (len(candidate_indices) > 1 or len(key_ids) > 1)
return {
"mode": "internal",
"attempted_count": executed_count,
"account_count": len(account_map),
"retry_occurred": executed_count > 1,
"failover_occurred": bool(failover_occurred),
"selected_key_id": selected_key_id_norm,
"selected_account": selected_account,
"accounts": accounts,
"attempts": attempts,
}
def _build_scheduling_metadata(
self,
*,
candidate_keys: list[Any] | None = None,
selected_key_id: str | None = None,
pool_summary: dict[str, Any] | None = None,
fallback_from_request: bool = False,
) -> dict[str, Any]:
snapshots = self._collect_candidate_snapshots(
candidate_keys=candidate_keys,
fallback_from_request=fallback_from_request,
)
metadata: dict[str, Any] = {}
if pool_summary:
metadata["pool_summary"] = pool_summary
if snapshots:
metadata["candidate_keys"] = snapshots
scheduling_audit = self._build_scheduling_audit(
snapshots,
selected_key_id=selected_key_id,
)
if scheduling_audit:
metadata["scheduling_audit"] = scheduling_audit
return metadata
def _merge_scheduling_metadata(
self,
request_metadata: dict[str, Any] | None,
*,
exec_result: Any | None = None,
selected_key_id: str | None = None,
candidate_keys: list[Any] | None = None,
pool_summary: dict[str, Any] | None = None,
fallback_from_request: bool = True,
) -> dict[str, Any] | None:
merged = dict(request_metadata or {})
resolved_candidate_keys = (
candidate_keys
if candidate_keys is not None
else getattr(exec_result, "candidate_keys", None)
)
resolved_key_id = selected_key_id or getattr(exec_result, "key_id", None)
resolved_pool_summary = (
pool_summary if pool_summary is not None else getattr(exec_result, "pool_summary", None)
)
merged.update(
self._build_scheduling_metadata(
candidate_keys=resolved_candidate_keys,
selected_key_id=resolved_key_id,
pool_summary=resolved_pool_summary,
fallback_from_request=fallback_from_request,
)
)
return merged or None
def _resolve_capability_requirements(
self,
model_name: str,
src/api/handlers/base/chat_handler_base.py
+15
View File
@@ -589,6 +589,21 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
ctx.key_id = key_id
if getattr(exec_result, "pool_summary", None):
ctx.pool_summary = exec_result.pool_summary
scheduling_metadata = (
self._merge_scheduling_metadata(
{},
exec_result=exec_result,
selected_key_id=key_id,
fallback_from_request=False,
)
or {}
)
candidate_keys = scheduling_metadata.get("candidate_keys")
if isinstance(candidate_keys, list):
ctx.candidate_keys = candidate_keys
scheduling_audit = scheduling_metadata.get("scheduling_audit")
if isinstance(scheduling_audit, dict):
ctx.scheduling_audit = scheduling_audit
# 同步整流状态（如果请求体被整流过）
ctx.rectified = request_body_ref.get("_rectified", False)
src/api/handlers/base/chat_sync_executor.py
+31 -4
View File
@@ -216,8 +216,11 @@ class ChatSyncExecutor:
request_metadata = handler._build_request_metadata() or {}
if ctx.sync_proxy_info:
request_metadata["proxy"] = ctx.sync_proxy_info
if getattr(exec_result, "pool_summary", None):
request_metadata["pool_summary"] = exec_result.pool_summary
request_metadata = handler._merge_scheduling_metadata(
request_metadata,
exec_result=exec_result,
selected_key_id=ctx.key_id,
)
total_cost = await handler.telemetry.record_success( # noqa: F841
provider=ctx.provider_name,
model=model,
@@ -251,7 +254,7 @@ class ChatSyncExecutor:
provider_api_key_id=ctx.key_id,
# 模型映射信息
target_model=ctx.mapped_model_result,
request_metadata=request_metadata or None,
request_metadata=request_metadata,
)
logger.debug(f"{handler.FORMAT_ID} 非流式响应完成")
@@ -273,6 +276,12 @@ class ChatSyncExecutor:
request_metadata = handler._build_request_metadata() or {}
if ctx.sync_proxy_info:
request_metadata["proxy"] = ctx.sync_proxy_info
request_metadata = handler._merge_scheduling_metadata(
request_metadata,
selected_key_id=ctx.key_id,
pool_summary=ctx.pool_summary,
fallback_from_request=True,
)
await handler.telemetry.record_failure(
provider=ctx.provider_name or "unknown",
model=model,
@@ -286,7 +295,7 @@ class ChatSyncExecutor:
provider_id=ctx.provider_id,
provider_endpoint_id=ctx.endpoint_id,
provider_api_key_id=ctx.key_id,
request_metadata=request_metadata or None,
request_metadata=request_metadata,
)
client_format = (ctx.client_api_format_for_error or "").upper()
provider_format = (ctx.provider_api_format_for_error or client_format).upper()
@@ -308,6 +317,12 @@ class ChatSyncExecutor:
request_metadata = handler._build_request_metadata() or {}
if ctx.sync_proxy_info:
request_metadata["proxy"] = ctx.sync_proxy_info
request_metadata = handler._merge_scheduling_metadata(
request_metadata,
selected_key_id=ctx.key_id,
pool_summary=ctx.pool_summary,
fallback_from_request=True,
)
client_format = (ctx.client_api_format_for_error or "").upper()
provider_format = (ctx.provider_api_format_for_error or client_format).upper()
payload = _build_error_json_payload(
@@ -372,6 +387,12 @@ class ChatSyncExecutor:
request_metadata = handler._build_request_metadata() or {}
if ctx.sync_proxy_info:
request_metadata["proxy"] = ctx.sync_proxy_info
request_metadata = handler._merge_scheduling_metadata(
request_metadata,
selected_key_id=ctx.key_id,
pool_summary=ctx.pool_summary,
fallback_from_request=True,
)
await handler.telemetry.record_failure(
provider=ctx.provider_name or "unknown",
model=model,
@@ -750,6 +771,12 @@ class ChatSyncExecutor:
stream_fail_metadata: dict[str, Any] | None = None
if ctx.proxy_info:
stream_fail_metadata = {"proxy": ctx.proxy_info}
stream_fail_metadata = handler._merge_scheduling_metadata(
stream_fail_metadata,
selected_key_id=ctx.key_id,
pool_summary=ctx.pool_summary,
fallback_from_request=True,
)
await handler.telemetry.record_failure(
provider=ctx.provider_name or "unknown",
src/api/handlers/base/cli_monitor_mixin.py
+28 -4
View File
@@ -299,7 +299,13 @@ class CliMonitorMixin:
if ctx.is_client_disconnected():
# 客户端取消：记录为 cancelled（不算系统失败）
request_metadata = {"perf": ctx.perf_metrics} if ctx.perf_metrics else None
request_metadata = self._merge_scheduling_metadata(
{"perf": ctx.perf_metrics} if ctx.perf_metrics else None,
selected_key_id=ctx.key_id,
candidate_keys=ctx.candidate_keys,
pool_summary=ctx.pool_summary,
fallback_from_request=True,
)
await bg_telemetry.record_cancelled(
provider=ctx.provider_name or "unknown",
model=ctx.model,
@@ -334,7 +340,13 @@ class CliMonitorMixin:
)
else:
# 服务端/上游异常：记录为失败
request_metadata = {"perf": ctx.perf_metrics} if ctx.perf_metrics else None
request_metadata = self._merge_scheduling_metadata(
{"perf": ctx.perf_metrics} if ctx.perf_metrics else None,
selected_key_id=ctx.key_id,
candidate_keys=ctx.candidate_keys,
pool_summary=ctx.pool_summary,
fallback_from_request=True,
)
await bg_telemetry.record_failure(
provider=ctx.provider_name or "unknown",
model=ctx.model,
@@ -412,7 +424,13 @@ class CliMonitorMixin:
f"provider={ctx.provider_name}, model={ctx.model}, "
f"in={ctx.input_tokens}, out={ctx.output_tokens}"
)
request_metadata = {"perf": ctx.perf_metrics} if ctx.perf_metrics else None
request_metadata = self._merge_scheduling_metadata(
{"perf": ctx.perf_metrics} if ctx.perf_metrics else None,
selected_key_id=ctx.key_id,
candidate_keys=ctx.candidate_keys,
pool_summary=ctx.pool_summary,
fallback_from_request=True,
)
total_cost = await bg_telemetry.record_success(
provider=ctx.provider_name,
model=ctx.model,
@@ -570,7 +588,13 @@ class CliMonitorMixin:
# 失败时返回给客户端的是 JSON 错误响应
client_response_headers = {"content-type": "application/json"}
request_metadata = {"perf": ctx.perf_metrics} if ctx.perf_metrics else None
request_metadata = self._merge_scheduling_metadata(
{"perf": ctx.perf_metrics} if ctx.perf_metrics else None,
selected_key_id=ctx.key_id,
candidate_keys=ctx.candidate_keys,
pool_summary=ctx.pool_summary,
fallback_from_request=True,
)
await self.telemetry.record_failure(
provider=ctx.provider_name or "unknown",
model=ctx.model,
src/api/handlers/base/cli_protocol.py
+11
View File
@@ -87,6 +87,17 @@ class CliHandlerProtocol(Protocol):
http_request: Any | None = ...,
) -> dict[str, Any] | None: ...
def _merge_scheduling_metadata(
self,
request_metadata: dict[str, Any] | None,
*,
exec_result: Any | None = ...,
selected_key_id: str | None = ...,
candidate_keys: list[Any] | None = ...,
pool_summary: dict[str, Any] | None = ...,
fallback_from_request: bool = ...,
) -> dict[str, Any] | None: ...
def _resolve_capability_requirements(
self,
model_name: str,
src/api/handlers/base/cli_stream_mixin.py
+15
View File
@@ -198,6 +198,21 @@ class CliStreamMixin:
ctx.key_id = key_id
if getattr(exec_result, "pool_summary", None):
ctx.pool_summary = exec_result.pool_summary
scheduling_metadata = (
self._merge_scheduling_metadata(
{},
exec_result=exec_result,
selected_key_id=key_id,
fallback_from_request=False,
)
or {}
)
candidate_keys = scheduling_metadata.get("candidate_keys")
if isinstance(candidate_keys, list):
ctx.candidate_keys = candidate_keys
scheduling_audit = scheduling_metadata.get("scheduling_audit")
if isinstance(scheduling_audit, dict):
ctx.scheduling_audit = scheduling_audit
# 同步整流状态（如果请求体被整流过）
ctx.rectified = request_body_ref.get("_rectified", False)
src/api/handlers/base/cli_sync_mixin.py
+21 -5
View File
@@ -101,6 +101,7 @@ class CliSyncMixin:
provider_id = None # Provider ID（用于失败记录）
endpoint_id = None # Endpoint ID（用于失败记录）
key_id = None # Key ID（用于失败记录）
exec_result = None
mapped_model_result = None # 映射后的目标模型名（用于 Usage 记录）
response_metadata_result: dict[str, Any] = {} # Provider 响应元数据
needs_conversion = False # 是否需要格式转换（由 candidate 决定）
@@ -558,8 +559,11 @@ class CliSyncMixin:
request_metadata = self._build_request_metadata() or {}
if sync_proxy_info:
request_metadata["proxy"] = sync_proxy_info
if getattr(exec_result, "pool_summary", None):
request_metadata["pool_summary"] = exec_result.pool_summary
request_metadata = self._merge_scheduling_metadata(
request_metadata,
exec_result=exec_result,
selected_key_id=key_id,
)
total_cost = await self.telemetry.record_success(
provider=provider_name,
model=model,
@@ -592,7 +596,7 @@ class CliSyncMixin:
target_model=mapped_model_result,
# Provider 响应元数据（如 Gemini 的 modelVersion）
response_metadata=response_metadata_result if response_metadata_result else None,
request_metadata=request_metadata or None,
request_metadata=request_metadata,
)
logger.info("{} 非流式响应处理完成", self.FORMAT_ID)
@@ -607,6 +611,12 @@ class CliSyncMixin:
request_metadata = self._build_request_metadata() or {}
if sync_proxy_info:
request_metadata["proxy"] = sync_proxy_info
request_metadata = self._merge_scheduling_metadata(
request_metadata,
selected_key_id=key_id,
pool_summary=getattr(exec_result, "pool_summary", None),
fallback_from_request=True,
)
await self.telemetry.record_failure(
provider=provider_name or "unknown",
model=model,
@@ -620,7 +630,7 @@ class CliSyncMixin:
api_format=api_format,
api_family=self.api_family,
endpoint_kind=self.endpoint_kind,
request_metadata=request_metadata or None,
request_metadata=request_metadata,
)
raise
@@ -645,6 +655,12 @@ class CliSyncMixin:
request_metadata = self._build_request_metadata() or {}
if sync_proxy_info:
request_metadata["proxy"] = sync_proxy_info
request_metadata = self._merge_scheduling_metadata(
request_metadata,
selected_key_id=key_id,
pool_summary=getattr(exec_result, "pool_summary", None),
fallback_from_request=True,
)
await self.telemetry.record_failure(
provider=provider_name or "unknown",
model=model,
@@ -667,7 +683,7 @@ class CliSyncMixin:
has_format_conversion=is_format_converted(provider_api_format, str(api_format)),
# 模型映射信息
target_model=mapped_model_result,
request_metadata=request_metadata or None,
request_metadata=request_metadata,
)
raise
src/api/handlers/base/stream_context.py
+7
View File
@@ -139,6 +139,10 @@ class StreamContext:
# 号池调度摘要（来自 ExecutionResult.pool_summary）
pool_summary: dict[str, Any] | None = None
# 候选轨迹（来自 ExecutionResult.candidate_keys，写入 usage metadata）
candidate_keys: list[dict[str, Any]] = field(default_factory=list)
# 内部调度审计摘要（重试/故障转移/账号使用轨迹）
scheduling_audit: dict[str, Any] | None = None
# 流式格式转换状态（跨 chunk 追踪）
stream_conversion_state: StreamState | None = None
@@ -175,6 +179,9 @@ class StreamContext:
self.final_usage = None
self.final_response = None
self.proxy_info = None
self.pool_summary = None
self.candidate_keys = []
self.scheduling_audit = None
self.stream_conversion_state = None
self.stream_conversion_event_count = 0
self.needs_conversion = False
src/api/handlers/base/stream_telemetry.py
+16
View File
@@ -210,6 +210,10 @@ class StreamTelemetryRecorder:
metadata["proxy"] = ctx.proxy_info
if ctx.pool_summary:
metadata["pool_summary"] = ctx.pool_summary
if ctx.candidate_keys:
metadata["candidate_keys"] = ctx.candidate_keys
if ctx.scheduling_audit:
metadata["scheduling_audit"] = ctx.scheduling_audit
await writer.record_success(
provider=ctx.provider_name or "unknown",
@@ -268,6 +272,12 @@ class StreamTelemetryRecorder:
metadata["perf"] = ctx.perf_metrics
if ctx.proxy_info:
metadata["proxy"] = ctx.proxy_info
if ctx.pool_summary:
metadata["pool_summary"] = ctx.pool_summary
if ctx.candidate_keys:
metadata["candidate_keys"] = ctx.candidate_keys
if ctx.scheduling_audit:
metadata["scheduling_audit"] = ctx.scheduling_audit
await writer.record_failure(
provider=ctx.provider_name or "unknown",
@@ -327,6 +337,12 @@ class StreamTelemetryRecorder:
metadata["perf"] = ctx.perf_metrics
if ctx.proxy_info:
metadata["proxy"] = ctx.proxy_info
if ctx.pool_summary:
metadata["pool_summary"] = ctx.pool_summary
if ctx.candidate_keys:
metadata["candidate_keys"] = ctx.candidate_keys
if ctx.scheduling_audit:
metadata["scheduling_audit"] = ctx.scheduling_audit
await writer.record_cancelled(
provider=ctx.provider_name or "unknown",
src/models/database.py
+1
View File
@@ -306,6 +306,7 @@ class Usage(Base):
Index("idx_usage_provider_model_created", "provider_name", "model", "created_at"),
Index("idx_usage_provider_created", "provider_name", "created_at"),
Index("idx_usage_model_created", "model", "created_at"),
Index("idx_usage_provider_key", "provider_id", "provider_api_key_id"),
)
id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()), index=True)
src/services/orchestration/candidate_resolver.py
+7
View File
@@ -214,6 +214,11 @@ class CandidateResolver:
provider = candidate.provider
endpoint = candidate.endpoint
key = candidate.key
pool_extra = (
getattr(candidate, "_pool_extra_data", None)
if isinstance(getattr(candidate, "_pool_extra_data", None), dict)
else {}
)
if candidate.is_skipped:
record_id = str(uuid.uuid4())
@@ -235,6 +240,7 @@ class CandidateResolver:
"needs_conversion": candidate.needs_conversion,
"provider_api_format": candidate.provider_api_format or None,
"mapping_matched_model": candidate.mapping_matched_model or None,
**pool_extra,
},
"required_capabilities": active_capabilities,
"created_at": datetime.now(timezone.utc),
@@ -269,6 +275,7 @@ class CandidateResolver:
"needs_conversion": candidate.needs_conversion,
"provider_api_format": candidate.provider_api_format or None,
"mapping_matched_model": candidate.mapping_matched_model or None,
**pool_extra,
},
"required_capabilities": active_capabilities,
"created_at": datetime.now(timezone.utc),
src/services/provider/adapters/kiro/models/credentials.py
+40 -8
View File
@@ -36,6 +36,28 @@ def _nonempty(s: str | None) -> str | None:
return None
def _normalize_auth_method(value: str | None) -> str:
method = (value or "").strip().lower()
if not method:
return "social"
# 历史/别名兼容：统一映射到 idc
if method in {
"idc",
"builder-id",
"builder_id",
"builderid",
"identity-center",
"identity_center",
"identitycenter",
"iam",
"device",
"device_authorization",
"device-auth",
}:
return "idc"
return method
def _parse_iso_to_epoch_seconds(value: object) -> int | None:
if not isinstance(value, str) or not value.strip():
return None
@@ -111,6 +133,11 @@ class KiroAuthConfig:
- 包含 clientId + clientSecret -> IdC
- 仅含 refreshToken -> Social
"""
explicit_method = _get_str(raw, "auth_method", "authMethod", "auth_type", "authType")
normalized_explicit = _normalize_auth_method(explicit_method)
if normalized_explicit != "social":
return normalized_explicit
client_id = raw.get("client_id") or raw.get("clientId")
client_secret = raw.get("client_secret") or raw.get("clientSecret")
@@ -136,7 +163,12 @@ class KiroAuthConfig:
return False, "refreshToken 不完整（含有 ...），请导出完整的 Token"
# IdC 类型需要 clientId 和 clientSecret
auth_method = KiroAuthConfig.infer_auth_method(raw)
explicit_method = _get_str(raw, "auth_method", "authMethod", "auth_type", "authType")
auth_method = (
_normalize_auth_method(explicit_method)
if explicit_method
else KiroAuthConfig.infer_auth_method(raw)
)
if auth_method == "idc":
client_id = raw.get("client_id") or raw.get("clientId")
client_secret = raw.get("client_secret") or raw.get("clientSecret")
@@ -155,8 +187,12 @@ class KiroAuthConfig:
provider_type = _get_str(raw, "provider_type", "providerType") or "kiro"
# 自动推断 auth_method（如果未显式指定）
explicit_method = _get_str(raw, "auth_method", "authMethod")
auth_method = explicit_method.lower() if explicit_method else cls.infer_auth_method(raw)
explicit_method = _get_str(raw, "auth_method", "authMethod", "auth_type", "authType")
auth_method = (
_normalize_auth_method(explicit_method)
if explicit_method
else cls.infer_auth_method(raw)
)
refresh_token = (_get_str(raw, "refresh_token", "refreshToken") or "").strip()
@@ -168,7 +204,7 @@ class KiroAuthConfig:
cfg = cls(
provider_type=provider_type,
auth_method=(auth_method or "social").lower(),
auth_method=_normalize_auth_method(auth_method),
refresh_token=refresh_token,
expires_at=int(expires_at),
profile_arn=_get_str(raw, "profile_arn", "profileArn"),
@@ -185,10 +221,6 @@ class KiroAuthConfig:
access_token=_get_str(raw, "access_token", "accessToken"),
)
# Normalize auth_method aliases.
if cfg.auth_method in {"builder-id", "builder_id", "iam"}:
cfg.auth_method = "idc"
return cfg
def to_dict(self) -> dict[str, Any]:
src/services/provider/pool/redis_ops.py
+44
View File
@@ -440,6 +440,50 @@ async def get_key_sticky_count(provider_id: str, key_id: str) -> int:
return 0
async def batch_get_key_sticky_counts(
provider_id: str,
key_ids: list[str],
) -> dict[str, int]:
"""Count sticky sessions for multiple keys in a single scan (admin only)."""
if not key_ids:
return {}
redis = await _get_redis()
if redis is None:
return {kid: 0 for kid in key_ids}
target_ids = set(key_ids)
counts: dict[str, int] = {kid: 0 for kid in key_ids}
try:
pattern = f"{PREFIX}:{provider_id}:sticky:*"
batch: list[bytes | str] = []
async for key in redis.scan_iter(match=pattern, count=200):
batch.append(key)
if len(batch) >= 200:
vals = await redis.mget(batch)
for val in vals:
if not val:
continue
bound_id = val.decode() if isinstance(val, bytes) else str(val)
if bound_id in target_ids:
counts[bound_id] = counts.get(bound_id, 0) + 1
batch.clear()
if batch:
vals = await redis.mget(batch)
for val in vals:
if not val:
continue
bound_id = val.decode() if isinstance(val, bytes) else str(val)
if bound_id in target_ids:
counts[bound_id] = counts.get(bound_id, 0) + 1
return counts
except Exception:
return {kid: 0 for kid in key_ids}
async def get_cooldown_ttl(provider_id: str, key_id: str) -> int | None:
"""Get remaining cooldown TTL in seconds. None = no cooldown."""
redis = await _get_redis()
src/services/provider/pool/scheduling_dimensions.py
+373
View File
@@ -0,0 +1,373 @@
"""Pool scheduling dimension registry and evaluation helpers.
This module keeps pool scheduling scoring isolated from API layer code.
Callers build a :class:`PoolSchedulingSnapshot` and evaluate it against
registered dimensions to obtain a normalized summary.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Protocol
PoolDimensionStatus = str # ok / degraded / blocked
@dataclass(frozen=True, slots=True)
class PoolSchedulingSnapshot:
"""Point-in-time scheduling inputs for one key."""
is_active: bool
cooldown_reason: str | None
cooldown_ttl_seconds: int | None
circuit_breaker_open: bool
cost_window_usage: int
cost_limit: int | None
cost_soft_threshold_percent: int = 80
health_score: float = 1.0
@dataclass(frozen=True, slots=True)
class PoolSchedulingDimensionResult:
"""Evaluation output for one scheduling dimension."""
code: str
label: str
status: PoolDimensionStatus = "ok"
blocking: bool = False
source: str = "pool"
weight: int = 1
score: float = 1.0
detail: str | None = None
ttl_seconds: int | None = None
@dataclass(frozen=True, slots=True)
class PoolSchedulingSummary:
"""Merged scheduling state across all dimensions."""
status: str # available / degraded / blocked
reason: str
label: str
score: float
candidate_eligible: bool
blocked_count: int
degraded_count: int
class PoolSchedulingDimension(Protocol):
"""Dimension evaluator protocol."""
code: str
label: str
source: str
weight: int
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
"""Evaluate one dimension from snapshot."""
@dataclass(frozen=True, slots=True)
class _ManualEnableDimension:
code: str = "manual_disabled"
label: str = "已禁用"
source: str = "manual"
weight: int = 8
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
if snapshot.is_active:
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="ok",
score=1.0,
)
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="blocked",
blocking=True,
score=0.0,
detail="账号被手动禁用",
)
@dataclass(frozen=True, slots=True)
class _CooldownDimension:
code: str = "cooldown"
label: str = "冷却中"
source: str = "pool"
weight: int = 7
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
if not snapshot.cooldown_reason:
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="ok",
score=1.0,
)
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="blocked",
blocking=True,
score=0.0,
detail=snapshot.cooldown_reason,
ttl_seconds=snapshot.cooldown_ttl_seconds,
)
@dataclass(frozen=True, slots=True)
class _CircuitBreakerDimension:
code: str = "circuit_open"
label: str = "熔断中"
source: str = "health"
weight: int = 6
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
if not snapshot.circuit_breaker_open:
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="ok",
score=1.0,
)
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="blocked",
blocking=True,
score=0.0,
)
@dataclass(frozen=True, slots=True)
class _CostDimension:
code: str = "cost"
label: str = "成本"
source: str = "pool"
weight: int = 5
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
limit = snapshot.cost_limit
usage = max(snapshot.cost_window_usage, 0)
if limit is None or limit <= 0:
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="ok",
score=1.0,
detail=f"{usage}/-",
)
ratio = usage / limit
detail = f"{usage}/{limit}"
if ratio >= 1.0:
return PoolSchedulingDimensionResult(
code="cost_exhausted",
label="成本超限",
source=self.source,
weight=self.weight,
status="blocked",
blocking=True,
score=0.0,
detail=detail,
)
soft_threshold = max(1, min(snapshot.cost_soft_threshold_percent, 100))
if ratio * 100 >= soft_threshold:
return PoolSchedulingDimensionResult(
code="cost_soft",
label="成本接近上限",
source=self.source,
weight=self.weight,
status="degraded",
score=0.45,
detail=detail,
)
if ratio >= 0.6:
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="degraded",
score=0.72,
detail=detail,
)
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="ok",
score=1.0,
detail=detail,
)
@dataclass(frozen=True, slots=True)
class _HealthDimension:
code: str = "health"
label: str = "健康度"
source: str = "health"
weight: int = 4
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
score = max(0.0, min(snapshot.health_score, 1.0))
detail = f"{score:.2f}"
if score < 0.5:
return PoolSchedulingDimensionResult(
code="health_low",
label="健康度过低",
source=self.source,
weight=self.weight,
status="degraded",
score=0.3,
detail=detail,
)
if score < 0.8:
return PoolSchedulingDimensionResult(
code="health_degraded",
label="健康度下降",
source=self.source,
weight=self.weight,
status="degraded",
score=0.65,
detail=detail,
)
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="ok",
score=1.0,
detail=detail,
)
_POOL_DIMENSION_REGISTRY: dict[str, PoolSchedulingDimension] = {}
_POOL_DIMENSION_ORDER: list[str] = []
def register_pool_scheduling_dimension(name: str, dimension: PoolSchedulingDimension) -> None:
"""Register a dimension evaluator by name."""
normalized = name.strip()
if not normalized:
return
if normalized not in _POOL_DIMENSION_ORDER:
_POOL_DIMENSION_ORDER.append(normalized)
_POOL_DIMENSION_REGISTRY[normalized] = dimension
def get_pool_scheduling_dimension(name: str) -> PoolSchedulingDimension | None:
"""Fetch a registered dimension evaluator."""
return _POOL_DIMENSION_REGISTRY.get(name.strip())
def list_pool_scheduling_dimensions() -> tuple[str, ...]:
"""List registered dimension names in evaluation order."""
return tuple(_POOL_DIMENSION_ORDER)
def evaluate_pool_scheduling_dimensions(
snapshot: PoolSchedulingSnapshot,
*,
dimension_names: tuple[str, ...] | None = None,
) -> list[PoolSchedulingDimensionResult]:
"""Evaluate snapshot across all registered dimensions."""
names = dimension_names or list_pool_scheduling_dimensions()
results: list[PoolSchedulingDimensionResult] = []
for name in names:
dimension = get_pool_scheduling_dimension(name)
if dimension is None:
continue
results.append(dimension.evaluate(snapshot))
return results
def summarize_pool_scheduling_dimensions(
dimensions: list[PoolSchedulingDimensionResult],
) -> PoolSchedulingSummary:
"""Summarize dimension outputs into a unified scheduling state."""
if not dimensions:
return PoolSchedulingSummary(
status="available",
reason="available",
label="可用",
score=100.0,
candidate_eligible=True,
blocked_count=0,
degraded_count=0,
)
blocked = [item for item in dimensions if item.status == "blocked" or item.blocking]
degraded = [item for item in dimensions if item.status == "degraded"]
total_weight = sum(max(item.weight, 1) for item in dimensions)
weighted_score = sum(
max(item.weight, 1) * max(min(item.score, 1.0), 0.0) for item in dimensions
) / max(total_weight, 1)
if blocked:
primary = blocked[0]
return PoolSchedulingSummary(
status="blocked",
reason=primary.code,
label=primary.label,
score=round(weighted_score * 100, 1),
candidate_eligible=False,
blocked_count=len(blocked),
degraded_count=len(degraded),
)
if degraded:
primary = degraded[0]
return PoolSchedulingSummary(
status="degraded",
reason=primary.code,
label=primary.label,
score=round(weighted_score * 100, 1),
candidate_eligible=True,
blocked_count=0,
degraded_count=len(degraded),
)
return PoolSchedulingSummary(
status="available",
reason="available",
label="可用",
score=round(weighted_score * 100, 1),
candidate_eligible=True,
blocked_count=0,
degraded_count=0,
)
def _register_default_dimensions() -> None:
register_pool_scheduling_dimension("manual", _ManualEnableDimension())
register_pool_scheduling_dimension("cooldown", _CooldownDimension())
register_pool_scheduling_dimension("circuit", _CircuitBreakerDimension())
register_pool_scheduling_dimension("cost", _CostDimension())
register_pool_scheduling_dimension("health", _HealthDimension())
_register_default_dimensions()
src/services/provider/pool/trace.py
+17 -3
View File
@@ -63,7 +63,12 @@ class PoolSchedulingTrace:
session_uuid: str | None = None
candidate_traces: dict[str, PoolCandidateTrace] = field(default_factory=dict)
def build_summary(self, success_key_id: str | None = None) -> dict[str, Any]:
def build_summary(
self,
success_key_id: str | None = None,
*,
attempted_key_ids: set[str] | None = None,
) -> dict[str, Any]:
"""Build compact dict for ``Usage.request_metadata["pool_summary"]``."""
skipped_cooldown = 0
skipped_cost = 0
@@ -74,8 +79,17 @@ class PoolSchedulingTrace:
skipped_cooldown += 1
elif t.skip_type == "cost_exhausted":
skipped_cost += 1
else:
attempted += 1
if attempted_key_ids is None:
# Backward-compatible behavior: count all schedulable keys.
attempted = sum(1 for t in self.candidate_traces.values() if not t.skipped)
else:
# Preferred behavior: count only keys that were actually executed.
attempted = sum(
1
for kid in attempted_key_ids
if kid in self.candidate_traces and not self.candidate_traces[kid].skipped
)
success_reason: str | None = None
if success_key_id and success_key_id in self.candidate_traces:
src/services/provider_keys/quota_cooldown.py
+35
View File
@@ -0,0 +1,35 @@
"""配额冷却判定工具。"""
from __future__ import annotations
from typing import Any
from src.core.logger import logger
from src.services.scheduling.quota_skipper import is_key_quota_exhausted
def resolve_effective_cooldown_reason(
*,
provider_type: str | None,
key: Any,
redis_reason: str | None,
) -> str | None:
"""返回 Key 的有效冷却原因。
规则：
- Redis 冷却存在时，优先返回 Redis 原因（429/403/quota_exhausted 等）。
- Redis 冷却不存在时，回退到 upstream_metadata 配额判断：
若账号级配额耗尽（Codex/Kiro），返回 ``quota_exhausted``。
"""
if redis_reason:
return redis_reason
try:
exhausted, _ = is_key_quota_exhausted(provider_type, key, model_name="")
except Exception:
logger.opt(exception=True).debug(
"quota_cooldown: is_key_quota_exhausted failed for key={}",
getattr(key, "id", "?"),
)
return None
return "quota_exhausted" if exhausted else None
src/services/scheduling/candidate_builder.py
+13 -43
View File
@@ -576,32 +576,21 @@ class CandidateBuilder:
if not active_keys:
continue
# --- Pool branch: select a single key internally ------
# Pool provider should still expose all key candidates here.
# Runtime pool scheduling/failover is handled later by TaskService._apply_pool_reorder.
use_random = all((key.cache_ttl_minutes or 0) == 0 for key in active_keys)
if pool_cfg is not None:
selected_key = await self._pool_select_key(
db, provider, pool_cfg, active_keys, request_body
)
if selected_key is None:
logger.debug(
"Pool[{}]: no schedulable key for endpoint {}",
str(provider.id)[:8],
endpoint_format_str,
)
continue
keys_to_check: list[ProviderAPIKey] = [selected_key]
else:
# --- Normal branch: check all keys ----
use_random = all((key.cache_ttl_minutes or 0) == 0 for key in active_keys)
if use_random and len(active_keys) > 1:
logger.debug(
" Provider {} 启用 Key 轮换模式 (endpoint_format={}, {} keys)",
provider.name,
endpoint_format_str,
len(active_keys),
)
keys_to_check = self._sorter.shuffle_keys_by_internal_priority(
active_keys, affinity_key, use_random
use_random = False
elif use_random and len(active_keys) > 1:
logger.debug(
" Provider {} 启用 Key 轮换模式 (endpoint_format={}, {} keys)",
provider.name,
endpoint_format_str,
len(active_keys),
)
keys_to_check = self._sorter.shuffle_keys_by_internal_priority(
active_keys, affinity_key, use_random
)
for key in keys_to_check:
# Key 级别检查（健康度/熔断按 provider_format bucket）
@@ -660,22 +649,3 @@ class CandidateBuilder:
candidates = candidates[:max_candidates]
return candidates
async def _pool_select_key(
self,
db: Session,
provider: Provider,
pool_cfg: "PoolConfig",
active_keys: list[ProviderAPIKey],
request_body: dict | None,
) -> ProviderAPIKey | None:
"""Select a single key via pool scheduling (sticky -> cooldown/cost -> LRU)."""
from src.services.provider.pool.hooks import get_pool_hook
from src.services.provider.pool.manager import PoolManager
provider_type = str(getattr(provider, "provider_type", "") or "")
hook = get_pool_hook(provider_type)
session_uuid = hook.extract_session_uuid(request_body) if hook and request_body else None
mgr = PoolManager(str(provider.id), pool_cfg)
release_db_connection_before_await(db)
return await mgr.select_key(session_uuid, active_keys)
src/services/task/service.py
+24 -1
View File
@@ -119,6 +119,7 @@ class TaskService:
allow_format_conversion=allow_format_conversion,
capability_requirements=capability_requirements,
max_candidates=max_candidates,
request_body=request_body,
)
candidate_keys = []
@@ -598,8 +599,22 @@ class TaskService:
# Build pool scheduling summary from traces collected during reorder.
if pool_traces and result.key_id:
try:
attempted_key_ids: set[str] = set()
for ck in result.candidate_keys or []:
status = str(getattr(ck, "status", "") or "").strip().lower()
if status in {"", "available", "pending", "skipped", "unused"}:
continue
kid = getattr(ck, "key_id", None)
if isinstance(kid, str) and kid:
attempted_key_ids.add(kid)
if not attempted_key_ids:
attempted_key_ids.add(str(result.key_id))
for pt in pool_traces:
summary = pt.build_summary(result.key_id)
summary = pt.build_summary(
result.key_id,
attempted_key_ids=attempted_key_ids,
)
if summary:
result.pool_summary = summary
break
@@ -1204,6 +1219,7 @@ class TaskService:
allow_format_conversion: bool = False,
capability_requirements: dict[str, bool] | None = None,
max_candidates: int | None = None,
request_body: dict[str, Any] | None = None,
) -> Any:
"""
Unified ASYNC submit entrypoint (Phase 3.2).
@@ -1300,6 +1316,7 @@ class TaskService:
request_id=request_id,
is_stream=False,
capability_requirements=capability_requirements,
request_body=request_body,
)
if not candidates:
@@ -1309,6 +1326,12 @@ class TaskService:
last_status_code=None,
)
# Account Pool: keep internal key failover order/skip behavior
# consistent with the SYNC path.
candidates, _pool_traces = await self._apply_pool_reorder(
candidates, request_body=request_body
)
if max_candidates is not None and max_candidates > 0:
candidates = candidates[:max_candidates]
src/services/usage/_recording_helpers.py
+1
View File
@@ -39,6 +39,7 @@ METADATA_KEEP_KEYS: frozenset[str] = frozenset(
"billing_updated_at",
"perf",
"pool_summary",
"scheduling_audit",
"_metadata_truncated",
}
)
tests/services/test_aware_scheduler_pool_candidate_enumeration.py
+57
View File
@@ -0,0 +1,57 @@
from unittest.mock import AsyncMock, MagicMock
import pytest
from src.services.scheduling.aware_scheduler import CacheAwareScheduler
def _mock_key(key_id: str, api_formats: list[str]) -> MagicMock:
key = MagicMock()
key.id = key_id
key.is_active = True
key.api_formats = api_formats
key.cache_ttl_minutes = 1
key.internal_priority = 1
return key
def _mock_endpoint(api_format: str) -> MagicMock:
endpoint = MagicMock()
endpoint.id = f"ep_{api_format.lower().replace(':', '_')}"
endpoint.is_active = True
endpoint.api_format = api_format
endpoint.api_family = api_format.split(":", 1)[0]
endpoint.endpoint_kind = api_format.split(":", 1)[1]
endpoint.format_acceptance_config = None
return endpoint
@pytest.mark.asyncio
async def test_pool_provider_enumerates_all_key_candidates() -> None:
scheduler = CacheAwareScheduler()
builder = scheduler._candidate_builder
builder._check_model_support = AsyncMock(return_value=(True, None, None, {"m"})) # type: ignore[method-assign]
builder._check_key_availability = MagicMock(return_value=(True, None, None)) # type: ignore[method-assign]
provider = MagicMock()
provider.id = "p_pool"
provider.name = "pool_provider"
provider.enable_format_conversion = False
provider.config = {"pool_advanced": {}}
provider.endpoints = [_mock_endpoint("openai:chat")]
provider.api_keys = [
_mock_key("k1", ["openai:chat"]),
_mock_key("k2", ["openai:chat"]),
]
candidates = await builder._build_candidates(
db=MagicMock(),
providers=[provider],
client_format="openai:chat",
model_name="dummy-model",
affinity_key="aff-1",
global_conversion_enabled=True,
)
assert len(candidates) == 2
assert {str(c.key.id) for c in candidates} == {"k1", "k2"}
tests/services/test_candidate_service_submit.py
+74
View File
@@ -310,3 +310,77 @@ async def test_submit_with_failover_filters_missing_billing_rule(
assert submit.await_count == 1
finally:
config.billing_require_rule = old
@pytest.mark.asyncio
async def test_submit_with_failover_applies_pool_reorder_before_submit(
monkeypatch: pytest.MonkeyPatch,
) -> None:
db = MagicMock()
svc = TaskService(db)
monkeypatch.setattr(
"src.services.system.config.SystemConfigService.get_config",
lambda *_args, **_kwargs: "provider",
)
monkeypatch.setattr(
"src.services.scheduling.aware_scheduler.get_cache_aware_scheduler",
AsyncMock(return_value=None),
)
pool_candidate_a = _make_candidate(
provider_id="pool-1",
provider_name="pool-provider",
endpoint_id="ep-1",
key_id="k-a",
key_name="key-a",
)
pool_candidate_b = _make_candidate(
provider_id="pool-1",
provider_name="pool-provider",
endpoint_id="ep-1",
key_id="k-b",
key_name="key-b",
)
fetch_candidates = AsyncMock(
return_value=(
[pool_candidate_a, pool_candidate_b],
"gm1",
)
)
monkeypatch.setattr(
"src.services.candidate.resolver.CandidateResolver.fetch_candidates",
fetch_candidates,
)
reordered = [pool_candidate_b, pool_candidate_a]
apply_pool_reorder = AsyncMock(return_value=(reordered, []))
monkeypatch.setattr(svc, "_apply_pool_reorder", apply_pool_reorder)
submit = AsyncMock(return_value=httpx.Response(200, json={"id": "task-pooled"}))
body = {"session_id": "sid-123"}
outcome = await svc.submit_with_failover(
api_format="openai:video",
model_name="sora",
affinity_key="a1",
user_api_key=MagicMock(),
request_id=None,
task_type="video",
submit_func=submit,
extract_external_task_id=lambda payload: payload.get("id"),
supported_auth_types={"api_key"},
allow_format_conversion=False,
max_candidates=10,
request_body=body,
)
assert outcome.external_task_id == "task-pooled"
assert outcome.candidate.key.id == "k-b"
assert submit.await_count == 1
fetch_candidates.assert_awaited_once()
assert fetch_candidates.await_args.kwargs.get("request_body") == body
apply_pool_reorder.assert_awaited_once_with(
[pool_candidate_a, pool_candidate_b],
request_body=body,
)
tests/services/test_pool_manager_trace.py
+37
View File
@@ -166,6 +166,43 @@ async def test_trace_build_summary_matches() -> None:
assert summary["success_key_id"] == "key-3"[:8]
@pytest.mark.asyncio
async def test_trace_build_summary_uses_attempted_key_ids_when_provided() -> None:
pool = PoolManager("prov-1", PoolConfig())
c1 = _make_candidate("key-1")
c2 = _make_candidate("key-2")
with (
patch(
"src.services.provider.pool.redis_ops.get_sticky_binding",
new_callable=AsyncMock,
return_value=None,
),
patch(
"src.services.provider.pool.redis_ops.batch_get_cooldowns",
new_callable=AsyncMock,
return_value={"key-1": (None, None), "key-2": (None, None)},
),
patch(
"src.services.provider.pool.redis_ops.get_lru_scores",
new_callable=AsyncMock,
return_value={},
),
):
result = await pool.reorder_candidates(None, [c1, c2])
trace = getattr(result[0], "_pool_scheduling_trace", None)
assert trace is not None
summary = trace.build_summary(
success_key_id="key-1",
attempted_key_ids={"key-1"},
)
assert summary["total_keys"] == 2
assert summary["attempted"] == 1
assert summary["success_key_id"] == "key-1"[:8]
@pytest.mark.asyncio
async def test_sticky_trace_info() -> None:
pool = PoolManager("prov-1", PoolConfig(sticky_session_ttl_seconds=3600))
tests/services/test_pool_scheduling_dimensions.py
+94
View File
@@ -0,0 +1,94 @@
"""Tests for pool scheduling dimension evaluation."""
from __future__ import annotations
from src.services.provider.pool.scheduling_dimensions import (
PoolSchedulingDimensionResult,
PoolSchedulingSnapshot,
evaluate_pool_scheduling_dimensions,
list_pool_scheduling_dimensions,
summarize_pool_scheduling_dimensions,
)
def _snapshot(**overrides: object) -> PoolSchedulingSnapshot:
base = {
"is_active": True,
"cooldown_reason": None,
"cooldown_ttl_seconds": None,
"circuit_breaker_open": False,
"cost_window_usage": 1200,
"cost_limit": 10000,
"cost_soft_threshold_percent": 80,
"health_score": 0.95,
}
base.update(overrides)
return PoolSchedulingSnapshot(**base) # type: ignore[arg-type]
def test_default_dimension_registry_contains_core_dimensions() -> None:
names = list_pool_scheduling_dimensions()
assert names == ("manual", "cooldown", "circuit", "cost", "health")
def test_summary_available_when_all_dimensions_ok() -> None:
dimensions = evaluate_pool_scheduling_dimensions(_snapshot())
summary = summarize_pool_scheduling_dimensions(dimensions)
assert summary.status == "available"
assert summary.reason == "available"
assert summary.candidate_eligible is True
assert summary.blocked_count == 0
assert summary.degraded_count == 0
assert summary.score == 100.0
def test_summary_blocked_when_manual_disabled() -> None:
dimensions = evaluate_pool_scheduling_dimensions(_snapshot(is_active=False))
summary = summarize_pool_scheduling_dimensions(dimensions)
assert summary.status == "blocked"
assert summary.reason == "manual_disabled"
assert summary.candidate_eligible is False
assert summary.blocked_count >= 1
assert summary.score < 100.0
def test_summary_degraded_when_cost_reaches_soft_threshold() -> None:
dimensions = evaluate_pool_scheduling_dimensions(
_snapshot(cost_window_usage=8200, cost_limit=10000, cost_soft_threshold_percent=80)
)
summary = summarize_pool_scheduling_dimensions(dimensions)
assert summary.status == "degraded"
assert summary.reason == "cost_soft"
assert summary.candidate_eligible is True
assert summary.blocked_count == 0
assert summary.degraded_count >= 1
def test_summary_blocked_on_cooldown_even_if_other_dimensions_ok() -> None:
dimensions = evaluate_pool_scheduling_dimensions(
_snapshot(cooldown_reason="rate_limited_429", cooldown_ttl_seconds=120)
)
summary = summarize_pool_scheduling_dimensions(dimensions)
assert summary.status == "blocked"
assert summary.reason == "cooldown"
assert summary.candidate_eligible is False
assert summary.blocked_count >= 1
def test_empty_summary_defaults_to_available() -> None:
summary = summarize_pool_scheduling_dimensions([])
assert summary.status == "available"
assert summary.reason == "available"
assert summary.score == 100.0
def test_dimension_result_keeps_degraded_health_details() -> None:
dimensions = evaluate_pool_scheduling_dimensions(_snapshot(health_score=0.65))
health = next((item for item in dimensions if item.code == "health_degraded"), None)
assert isinstance(health, PoolSchedulingDimensionResult)
assert health.status == "degraded"
assert health.detail == "0.65"
tests/services/test_provider_keys_quota_cooldown.py
+59
View File
@@ -0,0 +1,59 @@
from __future__ import annotations
from types import SimpleNamespace
from typing import Any, cast
from src.core.provider_types import ProviderType
from src.services.provider_keys.quota_cooldown import resolve_effective_cooldown_reason
def _key_with_metadata(upstream_metadata: dict[str, Any]) -> Any:
return cast(Any, SimpleNamespace(upstream_metadata=upstream_metadata))
def test_resolve_effective_cooldown_reason_prefers_redis_reason() -> None:
key = _key_with_metadata({"codex": {"primary_used_percent": 100.0}})
reason = resolve_effective_cooldown_reason(
provider_type=ProviderType.CODEX,
key=key,
redis_reason="rate_limited_429",
)
assert reason == "rate_limited_429"
def test_resolve_effective_cooldown_reason_fallbacks_to_codex_quota_exhausted() -> None:
key = _key_with_metadata({"codex": {"primary_used_percent": 100.0}})
reason = resolve_effective_cooldown_reason(
provider_type=ProviderType.CODEX,
key=key,
redis_reason=None,
)
assert reason == "quota_exhausted"
def test_resolve_effective_cooldown_reason_fallbacks_to_kiro_quota_exhausted() -> None:
key = _key_with_metadata({"kiro": {"remaining": 0}})
reason = resolve_effective_cooldown_reason(
provider_type=ProviderType.KIRO,
key=key,
redis_reason=None,
)
assert reason == "quota_exhausted"
def test_resolve_effective_cooldown_reason_returns_none_when_not_exhausted() -> None:
key = _key_with_metadata({"codex": {"primary_used_percent": 12.0}})
reason = resolve_effective_cooldown_reason(
provider_type=ProviderType.CODEX,
key=key,
redis_reason=None,
)
assert reason is None
tests/services/test_provider_oauth_kiro_import_parsing.py
+70
View File
@@ -0,0 +1,70 @@
from __future__ import annotations
import json
from src.api.admin.provider_oauth import _parse_kiro_import_input
from src.services.provider.adapters.kiro.models.credentials import KiroAuthConfig
def test_parse_kiro_import_input_array_unwraps_nested_auth_config() -> None:
raw = json.dumps(
[
{
"name": "acc-1",
"auth_config": {
"refresh_token": "rt-1",
"auth_method": "identity_center",
"client_id": "cid-1",
"client_secret": "csec-1",
},
}
]
)
parsed = _parse_kiro_import_input(raw)
assert len(parsed) == 1
assert parsed[0]["refresh_token"] == "rt-1"
assert parsed[0]["auth_method"] == "identity_center"
assert parsed[0]["client_id"] == "cid-1"
def test_parse_kiro_import_input_single_object_unwraps_auth_config() -> None:
raw = json.dumps(
{
"name": "acc-1",
"authConfig": {
"refreshToken": "rt-1",
"authType": "builder_id",
"clientId": "cid-1",
"clientSecret": "csec-1",
},
}
)
parsed = _parse_kiro_import_input(raw)
assert len(parsed) == 1
assert parsed[0]["refreshToken"] == "rt-1"
assert parsed[0]["authType"] == "builder_id"
def test_kiro_auth_config_from_dict_maps_device_alias_to_idc() -> None:
cfg = KiroAuthConfig.from_dict(
{
"refreshToken": "rt-1",
"auth_type": "identity_center",
"clientId": "cid-1",
"clientSecret": "csec-1",
}
)
assert cfg.auth_method == "idc"
def test_kiro_auth_config_validate_requires_idc_client_fields_when_explicit() -> None:
is_valid, message = KiroAuthConfig.validate_required_fields(
{
"refreshToken": "rt-1",
"auth_type": "builder_id",
}
)
assert is_valid is False
assert "clientId" in message
tests/unit/test_monitoring_trace_attempted_subset.py
+82
View File
@@ -0,0 +1,82 @@
from __future__ import annotations
import asyncio
from datetime import datetime, timezone
from types import SimpleNamespace
from unittest.mock import MagicMock
from src.api.admin.monitoring.trace import AdminGetRequestTraceAdapter
from src.services.request.candidate import RequestCandidateService
def _candidate(*, status: str, latency_ms: int | None = None) -> SimpleNamespace:
now = datetime.now(timezone.utc)
return SimpleNamespace(
id=f"cand-{status}",
request_id="req-1",
candidate_index=0,
retry_index=0,
provider_id=None,
endpoint_id=None,
key_id=None,
required_capabilities=None,
status=status,
skip_reason=None,
is_cached=False,
status_code=None,
error_type=None,
error_message=None,
latency_ms=latency_ms,
concurrent_requests=None,
extra_data=None,
created_at=now,
started_at=None,
finished_at=None,
)
def _context() -> SimpleNamespace:
return SimpleNamespace(
db=MagicMock(),
add_audit_metadata=lambda **_: None,
)
def test_trace_prefers_attempted_subset(monkeypatch: object) -> None:
candidates = [
_candidate(status="available"),
_candidate(status="unused"),
_candidate(status="failed", latency_ms=123),
]
monkeypatch.setattr(
RequestCandidateService,
"get_candidates_by_request_id",
lambda _db, _request_id: candidates,
)
adapter = AdminGetRequestTraceAdapter(request_id="req-1")
response = asyncio.run(adapter.handle(_context()))
assert response.total_candidates == 1
assert len(response.candidates) == 1
assert response.candidates[0].status == "failed"
assert response.total_latency_ms == 123
def test_trace_falls_back_to_all_when_no_attempted(monkeypatch: object) -> None:
candidates = [
_candidate(status="available"),
_candidate(status="unused"),
]
monkeypatch.setattr(
RequestCandidateService,
"get_candidates_by_request_id",
lambda _db, _request_id: candidates,
)
adapter = AdminGetRequestTraceAdapter(request_id="req-1")
response = asyncio.run(adapter.handle(_context()))
assert response.total_candidates == 2
assert len(response.candidates) == 2
assert {c.status for c in response.candidates} == {"available", "unused"}
tests/unit/test_pool_management_scheduling_state.py
+128
View File
@@ -0,0 +1,128 @@
"""Tests for pool management scheduling state assembly."""
from __future__ import annotations
from types import SimpleNamespace
from src.api.admin.pool.routes import (
_build_pool_scheduling_state,
_is_known_banned_key,
_is_known_banned_reason,
)
def test_pool_scheduling_state_manual_disabled_is_blocked() -> None:
(
status,
reason,
_label,
reasons,
score,
candidate_eligible,
blocked_count,
_degraded_count,
dimensions,
) = _build_pool_scheduling_state(
is_active=False,
cooldown_reason=None,
cooldown_ttl_seconds=None,
circuit_breaker_open=False,
cost_window_usage=0,
cost_limit=None,
cost_soft_threshold_percent=80,
health_score=1.0,
)
assert status == "blocked"
assert reason == "manual_disabled"
assert candidate_eligible is False
assert blocked_count >= 1
assert score < 100
assert any(item.code == "manual_disabled" for item in reasons)
assert any(item.code == "manual_disabled" for item in dimensions)
def test_pool_scheduling_state_cooldown_detail_is_mapped() -> None:
(
_status,
reason,
_label,
reasons,
_score,
_candidate_eligible,
_blocked_count,
_degraded_count,
dimensions,
) = _build_pool_scheduling_state(
is_active=True,
cooldown_reason="rate_limited_429",
cooldown_ttl_seconds=180,
circuit_breaker_open=False,
cost_window_usage=0,
cost_limit=None,
cost_soft_threshold_percent=80,
health_score=1.0,
)
assert reason == "cooldown"
cooldown_reason = next(item for item in reasons if item.code == "cooldown")
cooldown_dimension = next(item for item in dimensions if item.code == "cooldown")
assert cooldown_reason.detail == "429 限流"
assert cooldown_dimension.detail == "429 限流"
def test_pool_scheduling_state_cost_soft_is_degraded() -> None:
(
status,
reason,
_label,
_reasons,
_score,
candidate_eligible,
blocked_count,
degraded_count,
_dimensions,
) = _build_pool_scheduling_state(
is_active=True,
cooldown_reason=None,
cooldown_ttl_seconds=None,
circuit_breaker_open=False,
cost_window_usage=85,
cost_limit=100,
cost_soft_threshold_percent=80,
health_score=1.0,
)
assert status == "degraded"
assert reason == "cost_soft"
assert candidate_eligible is True
assert blocked_count == 0
assert degraded_count >= 1
def test_known_banned_reason_account_block_prefix() -> None:
assert _is_known_banned_reason("[ACCOUNT_BLOCK] Google 要求验证账号") is True
def test_known_banned_key_detects_kiro_banned_metadata() -> None:
key = SimpleNamespace(
upstream_metadata={"kiro": {"is_banned": True}},
oauth_invalid_reason=None,
)
assert _is_known_banned_key(key, "kiro") is True
def test_known_banned_key_detects_reason_keywords() -> None:
key = SimpleNamespace(
upstream_metadata={},
oauth_invalid_reason="AWS account temporarily suspended",
)
assert _is_known_banned_key(key, "antigravity") is True
def test_known_banned_key_does_not_treat_token_expired_as_banned() -> None:
key = SimpleNamespace(
upstream_metadata={"kiro": {"is_banned": False}},
oauth_invalid_reason="access token expired",
)
assert _is_known_banned_key(key, "kiro") is False
tests/unit/test_request_scheduling_metadata.py
+154
View File
@@ -0,0 +1,154 @@
from src.api.handlers.base.base_handler import BaseMessageHandler
from src.services.candidate.schema import CandidateKey
def _handler() -> BaseMessageHandler:
return BaseMessageHandler.__new__(BaseMessageHandler)
def test_scheduling_audit_detects_internal_failover() -> None:
handler = _handler()
metadata = handler._merge_scheduling_metadata(
{},
candidate_keys=[
{
"candidate_index": 0,
"retry_index": 0,
"provider_id": "p1",
"provider_name": "provider-a",
"key_id": "k1",
"key_name": "account-a",
"status": "failed",
"status_code": 429,
},
{
"candidate_index": 1,
"retry_index": 0,
"provider_id": "p1",
"provider_name": "provider-a",
"key_id": "k2",
"key_name": "account-b",
"status": "success",
"status_code": 200,
},
],
selected_key_id="k2",
fallback_from_request=False,
)
assert metadata is not None
audit = metadata.get("scheduling_audit")
assert isinstance(audit, dict)
assert audit.get("attempted_count") == 2
assert audit.get("account_count") == 2
assert audit.get("retry_occurred") is True
assert audit.get("failover_occurred") is True
assert audit.get("selected_key_id") == "k2"
accounts = audit.get("accounts")
assert isinstance(accounts, list)
assert any(a.get("key_id") == "k2" and a.get("selected") for a in accounts)
def test_scheduling_audit_distinguishes_retry_from_failover() -> None:
handler = _handler()
metadata = handler._merge_scheduling_metadata(
{},
candidate_keys=[
{
"candidate_index": 0,
"retry_index": 0,
"provider_id": "p1",
"provider_name": "provider-a",
"key_id": "k1",
"key_name": "account-a",
"status": "failed",
"status_code": 503,
},
{
"candidate_index": 0,
"retry_index": 1,
"provider_id": "p1",
"provider_name": "provider-a",
"key_id": "k1",
"key_name": "account-a",
"status": "success",
"status_code": 200,
},
],
selected_key_id="k1",
fallback_from_request=False,
)
assert metadata is not None
audit = metadata.get("scheduling_audit")
assert isinstance(audit, dict)
assert audit.get("attempted_count") == 2
assert audit.get("account_count") == 1
assert audit.get("retry_occurred") is True
assert audit.get("failover_occurred") is False
def test_scheduling_metadata_supports_candidate_key_dataclass() -> None:
handler = _handler()
metadata = handler._merge_scheduling_metadata(
{},
candidate_keys=[
CandidateKey(
candidate_index=0,
retry_index=0,
provider_id="p1",
provider_name="provider-a",
endpoint_id="e1",
key_id="k1",
key_name="account-a",
status="success",
status_code=200,
)
],
selected_key_id="k1",
fallback_from_request=False,
)
assert metadata is not None
snapshots = metadata.get("candidate_keys")
assert isinstance(snapshots, list)
assert snapshots[0]["status"] == "success"
assert snapshots[0]["key_id"] == "k1"
def test_scheduling_audit_excludes_unused_candidates() -> None:
handler = _handler()
metadata = handler._merge_scheduling_metadata(
{},
candidate_keys=[
{
"candidate_index": 0,
"retry_index": 0,
"provider_id": "p1",
"provider_name": "provider-a",
"key_id": "k1",
"key_name": "account-a",
"status": "success",
"status_code": 200,
},
{
"candidate_index": 1,
"retry_index": 0,
"provider_id": "p1",
"provider_name": "provider-a",
"key_id": "k2",
"key_name": "account-b",
"status": "unused",
},
],
selected_key_id="k1",
fallback_from_request=False,
)
assert metadata is not None
audit = metadata.get("scheduling_audit")
assert isinstance(audit, dict)
assert audit.get("attempted_count") == 1
assert audit.get("account_count") == 1
attempts = audit.get("attempts")
assert isinstance(attempts, list)
assert len(attempts) == 1
assert attempts[0].get("key_id") == "k1"