mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
feat(pool): 增加 OAuth 账号池管理功能
- 新增 pool manager / strategy / health_policy / cost_tracker / redis_ops 等核心模块 - 新增 pool admin API 路由与 schemas - 新增 OAuth 账号类型解析 (oauth_plan) - 前端增加 PoolManagement 页面、PoolConfigDialog、PoolImportDialog、PoolStatusCard 组件 - 补充 pool config / cost tracker / health policy / manager / strategy / trace 等测试
This commit is contained in:
149
frontend/src/api/endpoints/pool.ts
Normal file
149
frontend/src/api/endpoints/pool.ts
Normal file
@@ -0,0 +1,149 @@
|
||||
import client from '../client'
|
||||
|
||||
export interface PoolKeyStatus {
|
||||
key_id: string
|
||||
key_name: string
|
||||
is_active: boolean
|
||||
cooldown_reason: string | null
|
||||
cooldown_ttl_seconds: number | null
|
||||
cost_window_usage: number
|
||||
cost_limit: number | null
|
||||
sticky_sessions: number
|
||||
lru_score: number | null
|
||||
}
|
||||
|
||||
export interface PoolStatusResponse {
|
||||
provider_id: string
|
||||
provider_name: string
|
||||
pool_enabled: boolean
|
||||
total_keys: number
|
||||
total_sticky_sessions: number
|
||||
keys: PoolKeyStatus[]
|
||||
}
|
||||
|
||||
/**
|
||||
* 获取 Provider 的号池状态
|
||||
*/
|
||||
export async function getPoolStatus(providerId: string): Promise<PoolStatusResponse> {
|
||||
const response = await client.get(`/api/admin/providers/${providerId}/pool-status`)
|
||||
return response.data
|
||||
}
|
||||
|
||||
/**
|
||||
* 清除指定 Key 的号池冷却状态
|
||||
*/
|
||||
export async function clearPoolCooldown(
|
||||
providerId: string,
|
||||
keyId: string,
|
||||
): Promise<{ message: string }> {
|
||||
const response = await client.post(
|
||||
`/api/admin/providers/${providerId}/pool/clear-cooldown/${keyId}`,
|
||||
)
|
||||
return response.data
|
||||
}
|
||||
|
||||
/**
|
||||
* 重置指定 Key 的号池成本窗口
|
||||
*/
|
||||
export async function resetPoolCost(
|
||||
providerId: string,
|
||||
keyId: string,
|
||||
): Promise<{ message: string }> {
|
||||
const response = await client.post(
|
||||
`/api/admin/providers/${providerId}/pool/reset-cost/${keyId}`,
|
||||
)
|
||||
return response.data
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Pool management API (standalone page)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
export interface PoolOverviewItem {
|
||||
provider_id: string
|
||||
provider_name: string
|
||||
provider_type: string
|
||||
total_keys: number
|
||||
active_keys: number
|
||||
cooldown_count: number
|
||||
pool_enabled: boolean
|
||||
}
|
||||
|
||||
export interface PoolOverviewResponse {
|
||||
items: PoolOverviewItem[]
|
||||
}
|
||||
|
||||
export interface PoolKeyDetail {
|
||||
key_id: string
|
||||
key_name: string
|
||||
is_active: boolean
|
||||
auth_type: string
|
||||
cooldown_reason: string | null
|
||||
cooldown_ttl_seconds: number | null
|
||||
cost_window_usage: number
|
||||
cost_limit: number | null
|
||||
sticky_sessions: number
|
||||
lru_score: number | null
|
||||
created_at: string | null
|
||||
last_used_at: string | null
|
||||
}
|
||||
|
||||
export interface PoolKeysPageResponse {
|
||||
total: number
|
||||
page: number
|
||||
page_size: number
|
||||
keys: PoolKeyDetail[]
|
||||
}
|
||||
|
||||
export interface PoolKeysQuery {
|
||||
page?: number
|
||||
page_size?: number
|
||||
search?: string
|
||||
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'
|
||||
}
|
||||
|
||||
export async function getPoolOverview(): Promise<PoolOverviewResponse> {
|
||||
const response = await client.get('/api/admin/pool/overview')
|
||||
return response.data
|
||||
}
|
||||
|
||||
export async function listPoolKeys(
|
||||
providerId: string,
|
||||
params: PoolKeysQuery = {},
|
||||
): Promise<PoolKeysPageResponse> {
|
||||
const response = await client.get(`/api/admin/pool/${providerId}/keys`, { params })
|
||||
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,
|
||||
): Promise<{ affected: number; message: string }> {
|
||||
const response = await client.post(`/api/admin/pool/${providerId}/keys/batch-action`, body)
|
||||
return response.data
|
||||
}
|
||||
391
frontend/src/features/pool/components/PoolConfigDialog.vue
Normal file
391
frontend/src/features/pool/components/PoolConfigDialog.vue
Normal file
@@ -0,0 +1,391 @@
|
||||
<template>
|
||||
<Dialog
|
||||
:model-value="modelValue"
|
||||
title="号池配置"
|
||||
description="调整号池调度策略和健康检查参数"
|
||||
size="lg"
|
||||
@update:model-value="emit('update:modelValue', $event)"
|
||||
>
|
||||
<form
|
||||
class="space-y-5"
|
||||
@submit.prevent="handleSave"
|
||||
>
|
||||
<!-- 调度策略 -->
|
||||
<div class="space-y-3">
|
||||
<h3 class="text-sm font-medium border-b pb-2">
|
||||
调度策略
|
||||
</h3>
|
||||
|
||||
<div class="flex items-center justify-between p-3 border rounded-lg bg-muted/50">
|
||||
<div class="space-y-0.5">
|
||||
<span class="text-sm font-medium">LRU 调度</span>
|
||||
<p class="text-xs text-muted-foreground">
|
||||
优先选择最久未用的 Key
|
||||
</p>
|
||||
</div>
|
||||
<Switch
|
||||
:model-value="form.lru_enabled"
|
||||
@update:model-value="(v: boolean) => form.lru_enabled = v"
|
||||
/>
|
||||
</div>
|
||||
|
||||
<div class="grid grid-cols-2 gap-4">
|
||||
<div class="space-y-1.5">
|
||||
<Label>
|
||||
粘性会话 TTL
|
||||
<span class="text-xs text-muted-foreground">(秒)</span>
|
||||
</Label>
|
||||
<Input
|
||||
:model-value="form.sticky_session_ttl_seconds ?? ''"
|
||||
type="number"
|
||||
min="60"
|
||||
max="86400"
|
||||
placeholder="3600 (留空禁用)"
|
||||
@update:model-value="(v) => form.sticky_session_ttl_seconds = parseNum(v)"
|
||||
/>
|
||||
<p class="text-xs text-muted-foreground">
|
||||
同一对话始终路由到同一 Key
|
||||
</p>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- 冷却与健康 -->
|
||||
<div class="space-y-3">
|
||||
<h3 class="text-sm font-medium border-b pb-2">
|
||||
冷却与健康
|
||||
</h3>
|
||||
|
||||
<div class="flex items-center justify-between p-3 border rounded-lg bg-muted/50">
|
||||
<div class="space-y-0.5">
|
||||
<span class="text-sm font-medium">健康策略</span>
|
||||
<p class="text-xs text-muted-foreground">
|
||||
按上游错误码自动冷却/禁用 Key
|
||||
</p>
|
||||
</div>
|
||||
<Switch
|
||||
:model-value="form.health_policy_enabled"
|
||||
@update:model-value="(v: boolean) => form.health_policy_enabled = v"
|
||||
/>
|
||||
</div>
|
||||
|
||||
<div class="grid grid-cols-2 gap-4">
|
||||
<div class="space-y-1.5">
|
||||
<Label>
|
||||
429 冷却
|
||||
<span class="text-xs text-muted-foreground">(秒)</span>
|
||||
</Label>
|
||||
<Input
|
||||
:model-value="form.rate_limit_cooldown_seconds ?? ''"
|
||||
type="number"
|
||||
min="10"
|
||||
max="3600"
|
||||
placeholder="300"
|
||||
@update:model-value="(v) => form.rate_limit_cooldown_seconds = parseNum(v)"
|
||||
/>
|
||||
</div>
|
||||
<div class="space-y-1.5">
|
||||
<Label>
|
||||
529 冷却
|
||||
<span class="text-xs text-muted-foreground">(秒)</span>
|
||||
</Label>
|
||||
<Input
|
||||
:model-value="form.overload_cooldown_seconds ?? ''"
|
||||
type="number"
|
||||
min="5"
|
||||
max="600"
|
||||
placeholder="30"
|
||||
@update:model-value="(v) => form.overload_cooldown_seconds = parseNum(v)"
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- 成本控制 -->
|
||||
<div class="space-y-3">
|
||||
<h3 class="text-sm font-medium border-b pb-2">
|
||||
成本控制
|
||||
</h3>
|
||||
|
||||
<div class="grid grid-cols-2 gap-4">
|
||||
<div class="space-y-1.5">
|
||||
<Label>
|
||||
成本窗口
|
||||
<span class="text-xs text-muted-foreground">(秒)</span>
|
||||
</Label>
|
||||
<Input
|
||||
:model-value="form.cost_window_seconds ?? ''"
|
||||
type="number"
|
||||
min="3600"
|
||||
max="86400"
|
||||
placeholder="18000 (5 小时)"
|
||||
@update:model-value="(v) => form.cost_window_seconds = parseNum(v)"
|
||||
/>
|
||||
</div>
|
||||
<div class="space-y-1.5">
|
||||
<Label>
|
||||
Key 窗口限额
|
||||
<span class="text-xs text-muted-foreground">(tokens)</span>
|
||||
</Label>
|
||||
<Input
|
||||
:model-value="form.cost_limit_per_key_tokens ?? ''"
|
||||
type="number"
|
||||
min="0"
|
||||
placeholder="留空 = 不限"
|
||||
@update:model-value="(v) => form.cost_limit_per_key_tokens = parseNum(v)"
|
||||
/>
|
||||
</div>
|
||||
<div class="space-y-1.5">
|
||||
<Label>
|
||||
软阈值
|
||||
<span class="text-xs text-muted-foreground">(%)</span>
|
||||
</Label>
|
||||
<Input
|
||||
:model-value="form.cost_soft_threshold_percent ?? ''"
|
||||
type="number"
|
||||
min="0"
|
||||
max="100"
|
||||
placeholder="80"
|
||||
@update:model-value="(v) => form.cost_soft_threshold_percent = parseNum(v)"
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Claude Code 特有配置 -->
|
||||
<template v-if="providerType === 'claude_code'">
|
||||
<div class="space-y-3">
|
||||
<h3 class="text-sm font-medium border-b pb-2">
|
||||
Claude Code
|
||||
</h3>
|
||||
|
||||
<div class="space-y-3 p-3 border rounded-lg bg-muted/50">
|
||||
<div class="flex items-center justify-between">
|
||||
<div class="space-y-0.5">
|
||||
<span class="text-sm font-medium">会话数量控制</span>
|
||||
<p class="text-xs text-muted-foreground">
|
||||
限制同时活跃的会话数量
|
||||
</p>
|
||||
</div>
|
||||
<Switch
|
||||
:model-value="claudeForm.session_control_enabled"
|
||||
@update:model-value="(v: boolean) => claudeForm.session_control_enabled = v"
|
||||
/>
|
||||
</div>
|
||||
|
||||
<div
|
||||
v-if="claudeForm.session_control_enabled"
|
||||
class="grid grid-cols-2 gap-3"
|
||||
>
|
||||
<div class="space-y-1.5">
|
||||
<Label class="text-xs">最大会话数</Label>
|
||||
<Input
|
||||
:model-value="claudeForm.max_sessions ?? ''"
|
||||
type="number"
|
||||
min="1"
|
||||
max="1000"
|
||||
placeholder="例如 20"
|
||||
@update:model-value="(v) => claudeForm.max_sessions = parseNum(v)"
|
||||
/>
|
||||
</div>
|
||||
<div class="space-y-1.5">
|
||||
<Label class="text-xs">会话空闲超时 (分钟)</Label>
|
||||
<Input
|
||||
:model-value="claudeForm.session_idle_timeout_minutes ?? ''"
|
||||
type="number"
|
||||
min="1"
|
||||
max="1440"
|
||||
placeholder="默认 5"
|
||||
@update:model-value="(v) => claudeForm.session_idle_timeout_minutes = parseNum(v) ?? 5"
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="flex items-center justify-between p-3 border rounded-lg bg-muted/50">
|
||||
<div class="space-y-0.5">
|
||||
<span class="text-sm font-medium">TLS 指纹模拟</span>
|
||||
<p class="text-xs text-muted-foreground">
|
||||
模拟 Node.js / Claude Code 客户端的 TLS 指纹
|
||||
</p>
|
||||
</div>
|
||||
<Switch
|
||||
:model-value="claudeForm.enable_tls_fingerprint"
|
||||
@update:model-value="(v: boolean) => claudeForm.enable_tls_fingerprint = v"
|
||||
/>
|
||||
</div>
|
||||
|
||||
<div class="flex items-center justify-between p-3 border rounded-lg bg-muted/50">
|
||||
<div class="space-y-0.5">
|
||||
<span class="text-sm font-medium">会话 ID 伪装</span>
|
||||
<p class="text-xs text-muted-foreground">
|
||||
启用后在 15 分钟内固定 metadata.user_id 中的 session ID
|
||||
</p>
|
||||
</div>
|
||||
<Switch
|
||||
:model-value="claudeForm.session_id_masking_enabled"
|
||||
@update:model-value="(v: boolean) => claudeForm.session_id_masking_enabled = v"
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
</template>
|
||||
</form>
|
||||
|
||||
<template #footer>
|
||||
<Button
|
||||
variant="outline"
|
||||
:disabled="loading"
|
||||
@click="emit('update:modelValue', false)"
|
||||
>
|
||||
取消
|
||||
</Button>
|
||||
<Button
|
||||
:disabled="loading"
|
||||
@click="handleSave"
|
||||
>
|
||||
{{ loading ? '保存中...' : '保存' }}
|
||||
</Button>
|
||||
</template>
|
||||
</Dialog>
|
||||
</template>
|
||||
|
||||
<script setup lang="ts">
|
||||
import { ref, watch } from 'vue'
|
||||
import { Dialog, Button, Input, Label, Switch } from '@/components/ui'
|
||||
import { useToast } from '@/composables/useToast'
|
||||
import { parseApiError } from '@/utils/errorParser'
|
||||
import { updateProvider } from '@/api/endpoints'
|
||||
import type { PoolAdvancedConfig, ClaudeCodeAdvancedConfig } from '@/api/endpoints/types/provider'
|
||||
|
||||
const props = defineProps<{
|
||||
modelValue: boolean
|
||||
providerId: string
|
||||
providerType?: string
|
||||
currentConfig: PoolAdvancedConfig | null
|
||||
currentClaudeConfig?: ClaudeCodeAdvancedConfig | null
|
||||
}>()
|
||||
|
||||
const emit = defineEmits<{
|
||||
'update:modelValue': [value: boolean]
|
||||
saved: []
|
||||
}>()
|
||||
|
||||
const { success, error: showError } = useToast()
|
||||
const loading = ref(false)
|
||||
|
||||
const form = ref<PoolAdvancedConfig>({
|
||||
sticky_session_ttl_seconds: null,
|
||||
lru_enabled: true,
|
||||
cost_window_seconds: null,
|
||||
cost_limit_per_key_tokens: null,
|
||||
cost_soft_threshold_percent: null,
|
||||
rate_limit_cooldown_seconds: null,
|
||||
overload_cooldown_seconds: null,
|
||||
health_policy_enabled: true,
|
||||
})
|
||||
|
||||
interface ClaudeFormState {
|
||||
session_control_enabled: boolean
|
||||
max_sessions: number | undefined
|
||||
session_idle_timeout_minutes: number
|
||||
enable_tls_fingerprint: boolean
|
||||
session_id_masking_enabled: boolean
|
||||
}
|
||||
|
||||
const claudeForm = ref<ClaudeFormState>({
|
||||
session_control_enabled: true,
|
||||
max_sessions: undefined,
|
||||
session_idle_timeout_minutes: 5,
|
||||
enable_tls_fingerprint: true,
|
||||
session_id_masking_enabled: true,
|
||||
})
|
||||
|
||||
function parseNum(v: string | number): number | undefined {
|
||||
if (v === '' || v === null || v === undefined) return undefined
|
||||
const n = Number(v)
|
||||
return isNaN(n) ? undefined : n
|
||||
}
|
||||
|
||||
watch(() => props.modelValue, (v) => {
|
||||
if (v && props.currentConfig) {
|
||||
form.value = { ...props.currentConfig }
|
||||
} else if (v) {
|
||||
form.value = {
|
||||
sticky_session_ttl_seconds: null,
|
||||
lru_enabled: true,
|
||||
cost_window_seconds: null,
|
||||
cost_limit_per_key_tokens: null,
|
||||
cost_soft_threshold_percent: null,
|
||||
rate_limit_cooldown_seconds: null,
|
||||
overload_cooldown_seconds: null,
|
||||
health_policy_enabled: true,
|
||||
}
|
||||
}
|
||||
|
||||
// Claude Code 配置
|
||||
if (v && props.providerType === 'claude_code') {
|
||||
const cc = props.currentClaudeConfig
|
||||
if (cc) {
|
||||
// max_sessions 为 null 表示用户明确关闭了会话控制,其余情况默认开启
|
||||
const sessionOff = cc.max_sessions === null
|
||||
claudeForm.value = {
|
||||
session_control_enabled: !sessionOff,
|
||||
max_sessions: sessionOff ? undefined : (cc.max_sessions ?? undefined),
|
||||
session_idle_timeout_minutes: cc.session_idle_timeout_minutes ?? 5,
|
||||
enable_tls_fingerprint: cc.enable_tls_fingerprint ?? true,
|
||||
session_id_masking_enabled: cc.session_id_masking_enabled ?? true,
|
||||
}
|
||||
} else {
|
||||
// 默认值:全部开启
|
||||
claudeForm.value = {
|
||||
session_control_enabled: true,
|
||||
max_sessions: undefined,
|
||||
session_idle_timeout_minutes: 5,
|
||||
enable_tls_fingerprint: true,
|
||||
session_id_masking_enabled: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
async function handleSave() {
|
||||
loading.value = true
|
||||
try {
|
||||
const payload: Record<string, unknown> = {
|
||||
pool_advanced: {
|
||||
sticky_session_ttl_seconds: form.value.sticky_session_ttl_seconds ?? undefined,
|
||||
lru_enabled: form.value.lru_enabled,
|
||||
cost_window_seconds: form.value.cost_window_seconds ?? undefined,
|
||||
cost_limit_per_key_tokens: form.value.cost_limit_per_key_tokens ?? undefined,
|
||||
cost_soft_threshold_percent: form.value.cost_soft_threshold_percent ?? undefined,
|
||||
rate_limit_cooldown_seconds: form.value.rate_limit_cooldown_seconds ?? undefined,
|
||||
overload_cooldown_seconds: form.value.overload_cooldown_seconds ?? undefined,
|
||||
health_policy_enabled: form.value.health_policy_enabled,
|
||||
},
|
||||
}
|
||||
|
||||
// Claude Code 特有配置
|
||||
if (props.providerType === 'claude_code') {
|
||||
payload.claude_code_advanced = {
|
||||
max_sessions: claudeForm.value.session_control_enabled
|
||||
? (claudeForm.value.max_sessions ?? undefined)
|
||||
: null,
|
||||
session_idle_timeout_minutes: claudeForm.value.session_control_enabled
|
||||
? (claudeForm.value.session_idle_timeout_minutes ?? 5)
|
||||
: null,
|
||||
enable_tls_fingerprint: claudeForm.value.enable_tls_fingerprint,
|
||||
session_id_masking_enabled: claudeForm.value.session_id_masking_enabled,
|
||||
}
|
||||
}
|
||||
|
||||
await updateProvider(props.providerId, payload)
|
||||
success('号池配置已保存')
|
||||
emit('saved')
|
||||
emit('update:modelValue', false)
|
||||
} catch (err) {
|
||||
showError(parseApiError(err))
|
||||
} finally {
|
||||
loading.value = false
|
||||
}
|
||||
}
|
||||
</script>
|
||||
156
frontend/src/features/pool/components/PoolImportDialog.vue
Normal file
156
frontend/src/features/pool/components/PoolImportDialog.vue
Normal file
@@ -0,0 +1,156 @@
|
||||
<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="[{"name": "key-01", "api_key": "sk-xxx"}, ...]"
|
||||
/>
|
||||
<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>
|
||||
288
frontend/src/features/providers/components/PoolStatusCard.vue
Normal file
288
frontend/src/features/providers/components/PoolStatusCard.vue
Normal file
@@ -0,0 +1,288 @@
|
||||
<template>
|
||||
<Card class="overflow-hidden">
|
||||
<!-- Header -->
|
||||
<div class="p-4 border-b border-border/60">
|
||||
<div class="flex items-center justify-between">
|
||||
<div class="flex items-center gap-2">
|
||||
<h3 class="text-sm font-semibold">
|
||||
号池状态
|
||||
</h3>
|
||||
<Badge
|
||||
v-if="poolStatus"
|
||||
variant="secondary"
|
||||
class="text-xs"
|
||||
>
|
||||
{{ poolStatus.total_keys }} 个密钥
|
||||
</Badge>
|
||||
<Badge
|
||||
v-if="poolStatus && poolStatus.total_sticky_sessions > 0"
|
||||
variant="outline"
|
||||
class="text-xs"
|
||||
>
|
||||
{{ poolStatus.total_sticky_sessions }} 个粘性会话
|
||||
</Badge>
|
||||
</div>
|
||||
<RefreshButton
|
||||
:loading="refreshing"
|
||||
title="刷新号池状态"
|
||||
@click="refresh"
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Loading -->
|
||||
<div
|
||||
v-if="initialLoading"
|
||||
class="flex items-center justify-center py-8"
|
||||
>
|
||||
<div class="animate-spin rounded-full h-6 w-6 border-b-2 border-primary" />
|
||||
</div>
|
||||
|
||||
<!-- Pool not enabled -->
|
||||
<div
|
||||
v-else-if="poolStatus && !poolStatus.pool_enabled"
|
||||
class="p-6 text-center text-muted-foreground"
|
||||
>
|
||||
<p class="text-sm">
|
||||
号池未启用
|
||||
</p>
|
||||
<p class="text-xs mt-1">
|
||||
请在提供商编辑中配置号池参数
|
||||
</p>
|
||||
</div>
|
||||
|
||||
<!-- Key list -->
|
||||
<div
|
||||
v-else-if="poolStatus && poolStatus.keys.length > 0"
|
||||
class="divide-y divide-border/40"
|
||||
>
|
||||
<div
|
||||
v-for="key in poolStatus.keys"
|
||||
:key="key.key_id"
|
||||
class="px-4 py-3 hover:bg-muted/30 transition-colors"
|
||||
:class="{ 'opacity-40': !key.is_active }"
|
||||
>
|
||||
<!-- Row 1: name + cooldown + actions -->
|
||||
<div class="flex items-center justify-between gap-2">
|
||||
<div class="flex items-center gap-2 min-w-0">
|
||||
<span class="text-sm font-medium truncate">{{ key.key_name || '未命名' }}</span>
|
||||
<Badge
|
||||
v-if="key.cooldown_reason"
|
||||
variant="destructive"
|
||||
class="text-[10px] px-1.5 py-0 shrink-0"
|
||||
>
|
||||
{{ formatCooldownReason(key.cooldown_reason) }}
|
||||
</Badge>
|
||||
<span
|
||||
v-if="key.cooldown_ttl_seconds"
|
||||
class="text-[10px] text-destructive tabular-nums shrink-0"
|
||||
>
|
||||
{{ formatTTL(key.cooldown_ttl_seconds) }}
|
||||
</span>
|
||||
</div>
|
||||
<div class="flex items-center gap-0.5 shrink-0">
|
||||
<Button
|
||||
v-if="key.cooldown_reason"
|
||||
variant="ghost"
|
||||
size="icon"
|
||||
class="h-7 w-7 text-muted-foreground hover:text-green-600"
|
||||
title="清除冷却"
|
||||
:disabled="actionLoading === key.key_id"
|
||||
@click="handleClearCooldown(key.key_id)"
|
||||
>
|
||||
<RefreshCw
|
||||
class="w-3.5 h-3.5"
|
||||
:class="{ 'animate-spin': actionLoading === key.key_id }"
|
||||
/>
|
||||
</Button>
|
||||
<Button
|
||||
v-if="key.cost_limit != null && key.cost_window_usage > 0"
|
||||
variant="ghost"
|
||||
size="icon"
|
||||
class="h-7 w-7 text-muted-foreground hover:text-foreground"
|
||||
title="重置成本窗口"
|
||||
:disabled="actionLoading === key.key_id"
|
||||
@click="handleResetCost(key.key_id)"
|
||||
>
|
||||
<RotateCcw class="w-3.5 h-3.5" />
|
||||
</Button>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Row 2: cost + sticky + lru -->
|
||||
<div class="flex items-center gap-3 mt-1.5 text-[11px] text-muted-foreground">
|
||||
<!-- Cost with limit -->
|
||||
<div
|
||||
v-if="key.cost_limit != null"
|
||||
class="flex items-center gap-1.5 flex-1 min-w-0"
|
||||
>
|
||||
<span class="shrink-0">成本</span>
|
||||
<div class="flex-1 h-1.5 bg-border rounded-full overflow-hidden max-w-[120px]">
|
||||
<div
|
||||
class="h-full transition-all duration-300 rounded-full"
|
||||
:class="getCostBarColor(key.cost_window_usage, key.cost_limit)"
|
||||
:style="{
|
||||
width: `${Math.min((key.cost_window_usage / key.cost_limit) * 100, 100)}%`,
|
||||
}"
|
||||
/>
|
||||
</div>
|
||||
<span class="tabular-nums shrink-0">
|
||||
{{ formatTokens(key.cost_window_usage) }} / {{ formatTokens(key.cost_limit) }}
|
||||
</span>
|
||||
</div>
|
||||
<!-- Cost without limit -->
|
||||
<div
|
||||
v-else-if="key.cost_window_usage > 0"
|
||||
class="flex items-center gap-1"
|
||||
>
|
||||
<span>成本</span>
|
||||
<span class="tabular-nums">{{ formatTokens(key.cost_window_usage) }}</span>
|
||||
</div>
|
||||
|
||||
<span
|
||||
v-if="(key.cost_limit != null || key.cost_window_usage > 0) && key.sticky_sessions > 0"
|
||||
class="text-muted-foreground/40"
|
||||
>|</span>
|
||||
|
||||
<!-- Sticky sessions -->
|
||||
<span v-if="key.sticky_sessions > 0">
|
||||
{{ key.sticky_sessions }} 粘性会话
|
||||
</span>
|
||||
|
||||
<!-- LRU score -->
|
||||
<template v-if="key.lru_score != null">
|
||||
<span class="text-muted-foreground/40">|</span>
|
||||
<span>LRU {{ formatLruScore(key.lru_score) }}</span>
|
||||
</template>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Empty -->
|
||||
<div
|
||||
v-else
|
||||
class="p-6 text-center text-muted-foreground"
|
||||
>
|
||||
<p class="text-sm">
|
||||
暂无密钥数据
|
||||
</p>
|
||||
</div>
|
||||
</Card>
|
||||
</template>
|
||||
|
||||
<script setup lang="ts">
|
||||
import { onMounted, ref } from 'vue'
|
||||
import { RefreshCw, RotateCcw } from 'lucide-vue-next'
|
||||
|
||||
import { getPoolStatus, clearPoolCooldown, resetPoolCost } from '@/api/endpoints/pool'
|
||||
import type { PoolStatusResponse } from '@/api/endpoints/pool'
|
||||
import { parseApiError } from '@/utils/errorParser'
|
||||
import { useToast } from '@/composables/useToast'
|
||||
|
||||
import Card from '@/components/ui/card.vue'
|
||||
import Badge from '@/components/ui/badge.vue'
|
||||
import Button from '@/components/ui/button.vue'
|
||||
import RefreshButton from '@/components/ui/refresh-button.vue'
|
||||
|
||||
const props = defineProps<{
|
||||
providerId: string
|
||||
poolEnabled: boolean
|
||||
}>()
|
||||
|
||||
const { error: showError, success } = useToast()
|
||||
|
||||
const poolStatus = ref<PoolStatusResponse | null>(null)
|
||||
const initialLoading = ref(true)
|
||||
const refreshing = ref(false)
|
||||
const actionLoading = ref<string | null>(null)
|
||||
|
||||
async function loadPoolStatus() {
|
||||
try {
|
||||
poolStatus.value = await getPoolStatus(props.providerId)
|
||||
} catch (err) {
|
||||
showError(parseApiError(err))
|
||||
}
|
||||
}
|
||||
|
||||
async function refresh() {
|
||||
refreshing.value = true
|
||||
try {
|
||||
await loadPoolStatus()
|
||||
} finally {
|
||||
refreshing.value = false
|
||||
}
|
||||
}
|
||||
|
||||
async function handleClearCooldown(keyId: string) {
|
||||
actionLoading.value = keyId
|
||||
try {
|
||||
const res = await clearPoolCooldown(props.providerId, keyId)
|
||||
success(res.message)
|
||||
await loadPoolStatus()
|
||||
} catch (err) {
|
||||
showError(parseApiError(err))
|
||||
} finally {
|
||||
actionLoading.value = null
|
||||
}
|
||||
}
|
||||
|
||||
async function handleResetCost(keyId: string) {
|
||||
actionLoading.value = keyId
|
||||
try {
|
||||
const res = await resetPoolCost(props.providerId, keyId)
|
||||
success(res.message)
|
||||
await loadPoolStatus()
|
||||
} catch (err) {
|
||||
showError(parseApiError(err))
|
||||
} finally {
|
||||
actionLoading.value = null
|
||||
}
|
||||
}
|
||||
|
||||
const COOLDOWN_REASON_MAP: Record<string, string> = {
|
||||
rate_limited_429: '429 限流',
|
||||
forbidden_403: '403 禁止',
|
||||
overloaded_529: '529 过载',
|
||||
auth_failed_401: '401 认证失败',
|
||||
payment_required_402: '402 欠费',
|
||||
server_error_500: '500 错误',
|
||||
}
|
||||
|
||||
function formatCooldownReason(reason: string): string {
|
||||
return COOLDOWN_REASON_MAP[reason] || reason
|
||||
}
|
||||
|
||||
function formatTTL(seconds: number): string {
|
||||
if (seconds <= 0) return ''
|
||||
const m = Math.floor(seconds / 60)
|
||||
const s = seconds % 60
|
||||
return m > 0 ? `${m}m ${s}s` : `${s}s`
|
||||
}
|
||||
|
||||
function formatTokens(tokens: number): string {
|
||||
if (tokens >= 1_000_000) return `${(tokens / 1_000_000).toFixed(1)}M`
|
||||
if (tokens >= 1_000) return `${(tokens / 1_000).toFixed(1)}K`
|
||||
return String(tokens)
|
||||
}
|
||||
|
||||
function getCostBarColor(usage: number, limit: number): string {
|
||||
const ratio = usage / limit
|
||||
if (ratio >= 0.9) return 'bg-red-500'
|
||||
if (ratio >= 0.7) return 'bg-yellow-500'
|
||||
return 'bg-green-500'
|
||||
}
|
||||
|
||||
function formatLruScore(score: number): string {
|
||||
const now = Date.now() / 1000
|
||||
const diff = now - score
|
||||
if (diff < 60) return '刚刚'
|
||||
if (diff < 3600) return `${Math.floor(diff / 60)}m 前`
|
||||
if (diff < 86400) return `${Math.floor(diff / 3600)}h 前`
|
||||
return `${Math.floor(diff / 86400)}d 前`
|
||||
}
|
||||
|
||||
onMounted(async () => {
|
||||
await loadPoolStatus()
|
||||
initialLoading.value = false
|
||||
})
|
||||
</script>
|
||||
70
frontend/src/utils/oauthPlanType.ts
Normal file
70
frontend/src/utils/oauthPlanType.ts
Normal file
@@ -0,0 +1,70 @@
|
||||
const PLAN_TYPE_LABELS: Record<string, string> = {
|
||||
free: 'Free',
|
||||
plus: 'Plus',
|
||||
team: 'Team',
|
||||
enterprise: 'Enterprise',
|
||||
paid: 'Paid',
|
||||
pro: 'Pro',
|
||||
'pro+': 'Pro+',
|
||||
power: 'Power',
|
||||
ultra: 'Ultra',
|
||||
}
|
||||
|
||||
const PLAN_TYPE_CLASS_NAMES: Record<string, string> = {
|
||||
plus: 'border-green-500/50 text-green-600 dark:text-green-400',
|
||||
pro: 'border-blue-500/50 text-blue-600 dark:text-blue-400',
|
||||
free: 'border-primary/50 text-primary',
|
||||
paid: 'border-blue-500/50 text-blue-600 dark:text-blue-400',
|
||||
team: 'border-purple-500/50 text-purple-600 dark:text-purple-400',
|
||||
enterprise: 'border-amber-500/50 text-amber-600 dark:text-amber-400',
|
||||
ultra: 'border-amber-500/50 text-amber-600 dark:text-amber-400',
|
||||
'pro+': 'border-purple-500/50 text-purple-600 dark:text-purple-400',
|
||||
power: 'border-amber-500/50 text-amber-600 dark:text-amber-400',
|
||||
}
|
||||
|
||||
export function normalizeOAuthPlanType(planType?: string | null): string | null {
|
||||
if (typeof planType !== 'string') {
|
||||
return null
|
||||
}
|
||||
|
||||
const normalized = planType.trim().toLowerCase()
|
||||
if (!normalized) {
|
||||
return null
|
||||
}
|
||||
return normalized
|
||||
}
|
||||
|
||||
export function formatOAuthPlanType(planType?: string | null): string {
|
||||
const normalized = normalizeOAuthPlanType(planType)
|
||||
if (!normalized) {
|
||||
return ''
|
||||
}
|
||||
|
||||
const knownLabel = PLAN_TYPE_LABELS[normalized]
|
||||
if (knownLabel) {
|
||||
return knownLabel
|
||||
}
|
||||
|
||||
return normalized
|
||||
.replace(/[_-]+/g, ' ')
|
||||
.split(/\s+/)
|
||||
.filter(Boolean)
|
||||
.map(part => part[0].toUpperCase() + part.slice(1))
|
||||
.join(' ')
|
||||
}
|
||||
|
||||
export function getOAuthPlanTypeClass(planType?: string | null): string {
|
||||
const normalized = normalizeOAuthPlanType(planType)
|
||||
if (!normalized) {
|
||||
return ''
|
||||
}
|
||||
return PLAN_TYPE_CLASS_NAMES[normalized] || ''
|
||||
}
|
||||
|
||||
export function isNonFreeOAuthPlan(planType?: string | null): boolean {
|
||||
const normalized = normalizeOAuthPlanType(planType)
|
||||
if (!normalized) {
|
||||
return false
|
||||
}
|
||||
return normalized !== 'free'
|
||||
}
|
||||
855
frontend/src/views/admin/PoolManagement.vue
Normal file
855
frontend/src/views/admin/PoolManagement.vue
Normal file
@@ -0,0 +1,855 @@
|
||||
<template>
|
||||
<div class="space-y-6 pb-8">
|
||||
<Card
|
||||
variant="default"
|
||||
class="overflow-hidden"
|
||||
>
|
||||
<!-- Header -->
|
||||
<div class="px-4 sm:px-6 py-3 sm:py-3.5 border-b border-border/60">
|
||||
<!-- Mobile -->
|
||||
<div class="flex flex-col gap-3 sm:hidden">
|
||||
<div class="flex items-center justify-between">
|
||||
<h3 class="text-base font-semibold">
|
||||
号池管理
|
||||
</h3>
|
||||
<div class="flex items-center gap-1.5">
|
||||
<Button
|
||||
v-if="selectedProviderId"
|
||||
variant="ghost"
|
||||
size="icon"
|
||||
class="h-8 w-8"
|
||||
title="批量导入"
|
||||
@click="showImportDialog = true"
|
||||
>
|
||||
<Upload class="w-3.5 h-3.5" />
|
||||
</Button>
|
||||
<Button
|
||||
v-if="selectedProviderId"
|
||||
variant="ghost"
|
||||
size="icon"
|
||||
class="h-8 w-8"
|
||||
title="号池配置"
|
||||
@click="showConfigDialog = true"
|
||||
>
|
||||
<Settings class="w-3.5 h-3.5" />
|
||||
</Button>
|
||||
<RefreshButton
|
||||
:loading="overviewLoading || keysLoading"
|
||||
@click="refresh"
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
<!-- Filters (mobile) -->
|
||||
<div class="flex items-center gap-2">
|
||||
<Select v-model="selectedProviderIdProxy">
|
||||
<SelectTrigger class="flex-1 h-8 text-xs border-border/60">
|
||||
<SelectValue placeholder="选择 Provider" />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
<SelectItem
|
||||
v-for="item in poolProviders"
|
||||
:key="item.provider_id"
|
||||
:value="item.provider_id"
|
||||
>
|
||||
{{ item.provider_name }}
|
||||
<span class="text-muted-foreground ml-1">({{ item.total_keys }})</span>
|
||||
</SelectItem>
|
||||
</SelectContent>
|
||||
</Select>
|
||||
<Select v-model="statusFilter">
|
||||
<SelectTrigger class="w-24 h-8 text-xs border-border/60">
|
||||
<SelectValue placeholder="状态" />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
<SelectItem value="all">
|
||||
全部
|
||||
</SelectItem>
|
||||
<SelectItem value="active">
|
||||
活跃
|
||||
</SelectItem>
|
||||
<SelectItem value="cooldown">
|
||||
冷却中
|
||||
</SelectItem>
|
||||
<SelectItem value="inactive">
|
||||
已禁用
|
||||
</SelectItem>
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</div>
|
||||
<div
|
||||
v-if="selectedProviderId"
|
||||
class="relative"
|
||||
>
|
||||
<Search class="absolute left-2.5 top-1/2 -translate-y-1/2 h-3.5 w-3.5 text-muted-foreground z-10 pointer-events-none" />
|
||||
<Input
|
||||
v-model="searchQuery"
|
||||
type="text"
|
||||
placeholder="搜索账号..."
|
||||
class="w-full pl-8 pr-3 h-8 text-sm bg-background/50 border-border/60"
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Desktop -->
|
||||
<div class="hidden sm:flex items-center justify-between gap-4">
|
||||
<h3 class="text-base font-semibold">
|
||||
号池管理
|
||||
</h3>
|
||||
<div class="flex items-center gap-2">
|
||||
<Select v-model="selectedProviderIdProxy">
|
||||
<SelectTrigger class="w-44 h-8 text-xs border-border/60">
|
||||
<SelectValue placeholder="选择 Provider" />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
<SelectItem
|
||||
v-for="item in poolProviders"
|
||||
:key="item.provider_id"
|
||||
:value="item.provider_id"
|
||||
>
|
||||
{{ item.provider_name }}
|
||||
<span class="text-muted-foreground ml-1">({{ item.total_keys }})</span>
|
||||
</SelectItem>
|
||||
</SelectContent>
|
||||
</Select>
|
||||
<div class="h-4 w-px bg-border" />
|
||||
<div
|
||||
v-if="selectedProviderId"
|
||||
class="relative"
|
||||
>
|
||||
<Search class="absolute left-2.5 top-1/2 -translate-y-1/2 h-3.5 w-3.5 text-muted-foreground z-10 pointer-events-none" />
|
||||
<Input
|
||||
v-model="searchQuery"
|
||||
type="text"
|
||||
placeholder="搜索账号..."
|
||||
class="w-40 pl-8 pr-2 h-8 text-xs bg-background/50 border-border/60"
|
||||
/>
|
||||
</div>
|
||||
<Select v-model="statusFilter">
|
||||
<SelectTrigger class="w-28 h-8 text-xs border-border/60">
|
||||
<SelectValue placeholder="全部状态" />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
<SelectItem value="all">
|
||||
全部状态
|
||||
</SelectItem>
|
||||
<SelectItem value="active">
|
||||
活跃
|
||||
</SelectItem>
|
||||
<SelectItem value="cooldown">
|
||||
冷却中
|
||||
</SelectItem>
|
||||
<SelectItem value="inactive">
|
||||
已禁用
|
||||
</SelectItem>
|
||||
</SelectContent>
|
||||
</Select>
|
||||
<div
|
||||
v-if="selectedProviderId"
|
||||
class="h-4 w-px bg-border"
|
||||
/>
|
||||
<Button
|
||||
v-if="selectedProviderId"
|
||||
variant="ghost"
|
||||
size="icon"
|
||||
class="h-8 w-8"
|
||||
title="批量导入"
|
||||
@click="showImportDialog = true"
|
||||
>
|
||||
<Upload class="w-3.5 h-3.5" />
|
||||
</Button>
|
||||
<Button
|
||||
v-if="selectedProviderId"
|
||||
variant="ghost"
|
||||
size="icon"
|
||||
class="h-8 w-8"
|
||||
title="号池配置"
|
||||
@click="showConfigDialog = true"
|
||||
>
|
||||
<Settings class="w-3.5 h-3.5" />
|
||||
</Button>
|
||||
<RefreshButton
|
||||
:loading="overviewLoading || keysLoading"
|
||||
@click="refresh"
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Loading (initial) -->
|
||||
<div
|
||||
v-if="overviewLoading"
|
||||
class="flex items-center justify-center py-16"
|
||||
>
|
||||
<div class="animate-spin rounded-full h-8 w-8 border-b-2 border-primary" />
|
||||
</div>
|
||||
|
||||
<!-- No providers -->
|
||||
<div
|
||||
v-else-if="poolProviders.length === 0"
|
||||
class="flex flex-col items-center justify-center py-16 text-center"
|
||||
>
|
||||
<div class="mx-auto flex h-16 w-16 items-center justify-center rounded-full bg-muted">
|
||||
<Database class="h-8 w-8 text-muted-foreground" />
|
||||
</div>
|
||||
<p class="text-sm text-muted-foreground mt-4">
|
||||
暂无 Provider
|
||||
</p>
|
||||
<p class="text-xs text-muted-foreground mt-1">
|
||||
请先添加 Provider
|
||||
</p>
|
||||
</div>
|
||||
|
||||
<!-- No provider selected -->
|
||||
<div
|
||||
v-else-if="!selectedProviderId"
|
||||
class="flex flex-col items-center justify-center py-16 text-center"
|
||||
>
|
||||
<div class="mx-auto flex h-16 w-16 items-center justify-center rounded-full bg-muted">
|
||||
<Database class="h-8 w-8 text-muted-foreground" />
|
||||
</div>
|
||||
<p class="text-sm text-muted-foreground mt-4">
|
||||
请选择一个 Provider 查看账号
|
||||
</p>
|
||||
</div>
|
||||
|
||||
<!-- Loading keys -->
|
||||
<div
|
||||
v-else-if="keysLoading && keyPage.keys.length === 0"
|
||||
class="flex items-center justify-center py-16"
|
||||
>
|
||||
<div class="animate-spin rounded-full h-8 w-8 border-b-2 border-primary" />
|
||||
</div>
|
||||
|
||||
<template v-else>
|
||||
<!-- Batch action bar -->
|
||||
<div
|
||||
v-if="selectedKeys.size > 0"
|
||||
class="flex items-center gap-2 px-4 sm:px-6 py-2.5 bg-muted/40 border-b border-border/40"
|
||||
>
|
||||
<span class="text-xs font-medium text-muted-foreground mr-1">
|
||||
已选 {{ selectedKeys.size }} 个
|
||||
</span>
|
||||
<Button
|
||||
variant="outline"
|
||||
size="sm"
|
||||
class="h-7 text-xs"
|
||||
@click="batchAction('enable')"
|
||||
>
|
||||
启用
|
||||
</Button>
|
||||
<Button
|
||||
variant="outline"
|
||||
size="sm"
|
||||
class="h-7 text-xs"
|
||||
@click="batchAction('disable')"
|
||||
>
|
||||
禁用
|
||||
</Button>
|
||||
<Button
|
||||
variant="outline"
|
||||
size="sm"
|
||||
class="h-7 text-xs"
|
||||
@click="batchAction('clear_cooldown')"
|
||||
>
|
||||
清除冷却
|
||||
</Button>
|
||||
<Button
|
||||
variant="outline"
|
||||
size="sm"
|
||||
class="h-7 text-xs"
|
||||
@click="batchAction('reset_cost')"
|
||||
>
|
||||
重置成本
|
||||
</Button>
|
||||
<Button
|
||||
variant="destructive"
|
||||
size="sm"
|
||||
class="h-7 text-xs"
|
||||
@click="batchAction('delete')"
|
||||
>
|
||||
删除
|
||||
</Button>
|
||||
</div>
|
||||
|
||||
<!-- Desktop table -->
|
||||
<div
|
||||
v-if="keyPage.keys.length > 0"
|
||||
class="hidden xl:block overflow-x-auto"
|
||||
>
|
||||
<Table>
|
||||
<TableHeader>
|
||||
<TableRow class="border-b border-border/60 hover:bg-transparent">
|
||||
<TableHead class="w-10 h-12">
|
||||
<input
|
||||
type="checkbox"
|
||||
:checked="allSelected"
|
||||
class="rounded"
|
||||
@change="toggleSelectAll"
|
||||
>
|
||||
</TableHead>
|
||||
<TableHead class="font-semibold">
|
||||
名称
|
||||
</TableHead>
|
||||
<TableHead class="w-20 font-semibold">
|
||||
状态
|
||||
</TableHead>
|
||||
<TableHead class="w-32 font-semibold">
|
||||
冷却
|
||||
</TableHead>
|
||||
<TableHead class="w-40 font-semibold">
|
||||
成本
|
||||
</TableHead>
|
||||
<TableHead class="w-16 font-semibold text-center">
|
||||
会话
|
||||
</TableHead>
|
||||
<TableHead class="w-24 font-semibold">
|
||||
最后使用
|
||||
</TableHead>
|
||||
<TableHead class="w-20 font-semibold text-center">
|
||||
操作
|
||||
</TableHead>
|
||||
</TableRow>
|
||||
</TableHeader>
|
||||
<TableBody>
|
||||
<TableRow
|
||||
v-for="key in keyPage.keys"
|
||||
:key="key.key_id"
|
||||
class="border-b border-border/40 last:border-b-0 hover:bg-muted/30 transition-colors"
|
||||
:class="{ 'opacity-50': !key.is_active }"
|
||||
>
|
||||
<TableCell class="py-3">
|
||||
<input
|
||||
type="checkbox"
|
||||
:checked="selectedKeys.has(key.key_id)"
|
||||
class="rounded"
|
||||
@change="toggleSelect(key.key_id)"
|
||||
>
|
||||
</TableCell>
|
||||
<TableCell class="py-3">
|
||||
<span class="text-sm truncate max-w-[200px] block">
|
||||
{{ key.key_name || '未命名' }}
|
||||
</span>
|
||||
</TableCell>
|
||||
<TableCell class="py-3">
|
||||
<Badge
|
||||
:variant="key.is_active ? (key.cooldown_reason ? 'destructive' : 'default') : 'secondary'"
|
||||
class="text-[10px]"
|
||||
>
|
||||
{{ key.is_active ? (key.cooldown_reason ? '冷却' : '活跃') : '禁用' }}
|
||||
</Badge>
|
||||
</TableCell>
|
||||
<TableCell class="py-3">
|
||||
<template v-if="key.cooldown_reason">
|
||||
<div class="flex items-center gap-1">
|
||||
<span class="text-xs text-destructive">
|
||||
{{ formatCooldownReason(key.cooldown_reason) }}
|
||||
</span>
|
||||
<span
|
||||
v-if="key.cooldown_ttl_seconds"
|
||||
class="text-[10px] text-muted-foreground"
|
||||
>
|
||||
{{ formatTTL(key.cooldown_ttl_seconds) }}
|
||||
</span>
|
||||
</div>
|
||||
</template>
|
||||
<span
|
||||
v-else
|
||||
class="text-xs text-muted-foreground"
|
||||
>-</span>
|
||||
</TableCell>
|
||||
<TableCell class="py-3">
|
||||
<div
|
||||
v-if="key.cost_limit != null"
|
||||
class="flex items-center gap-2"
|
||||
>
|
||||
<div class="flex-1 h-1.5 bg-border rounded-full overflow-hidden max-w-[80px]">
|
||||
<div
|
||||
class="h-full transition-all duration-300 rounded-full"
|
||||
:class="getCostBarColor(key.cost_window_usage, key.cost_limit)"
|
||||
:style="{ width: `${Math.min((key.cost_window_usage / key.cost_limit) * 100, 100)}%` }"
|
||||
/>
|
||||
</div>
|
||||
<span class="text-[10px] tabular-nums text-muted-foreground whitespace-nowrap">
|
||||
{{ formatTokens(key.cost_window_usage) }}/{{ formatTokens(key.cost_limit) }}
|
||||
</span>
|
||||
</div>
|
||||
<span
|
||||
v-else-if="key.cost_window_usage > 0"
|
||||
class="text-[10px] tabular-nums text-muted-foreground"
|
||||
>
|
||||
{{ formatTokens(key.cost_window_usage) }}
|
||||
</span>
|
||||
<span
|
||||
v-else
|
||||
class="text-xs text-muted-foreground"
|
||||
>-</span>
|
||||
</TableCell>
|
||||
<TableCell class="py-3 text-center">
|
||||
<span
|
||||
v-if="key.sticky_sessions > 0"
|
||||
class="text-xs tabular-nums"
|
||||
>
|
||||
{{ key.sticky_sessions }}
|
||||
</span>
|
||||
<span
|
||||
v-else
|
||||
class="text-xs text-muted-foreground"
|
||||
>-</span>
|
||||
</TableCell>
|
||||
<TableCell class="py-3">
|
||||
<span class="text-[10px] text-muted-foreground whitespace-nowrap">
|
||||
{{ key.last_used_at ? formatRelativeTime(key.last_used_at) : '-' }}
|
||||
</span>
|
||||
</TableCell>
|
||||
<TableCell class="py-3">
|
||||
<div class="flex justify-center gap-0.5">
|
||||
<Button
|
||||
v-if="key.cooldown_reason"
|
||||
variant="ghost"
|
||||
size="icon"
|
||||
class="h-7 w-7 text-muted-foreground hover:text-green-600"
|
||||
title="清除冷却"
|
||||
@click="clearCooldown(key.key_id)"
|
||||
>
|
||||
<RefreshCw class="w-3.5 h-3.5" />
|
||||
</Button>
|
||||
<Button
|
||||
variant="ghost"
|
||||
size="icon"
|
||||
class="h-7 w-7 text-muted-foreground hover:text-foreground"
|
||||
:title="key.is_active ? '禁用' : '启用'"
|
||||
@click="toggleKeyActive(key)"
|
||||
>
|
||||
<component
|
||||
:is="key.is_active ? Ban : Check"
|
||||
class="w-3.5 h-3.5"
|
||||
/>
|
||||
</Button>
|
||||
</div>
|
||||
</TableCell>
|
||||
</TableRow>
|
||||
</TableBody>
|
||||
</Table>
|
||||
</div>
|
||||
|
||||
<!-- Mobile card list -->
|
||||
<div
|
||||
v-if="keyPage.keys.length > 0"
|
||||
class="xl:hidden divide-y divide-border/40"
|
||||
>
|
||||
<div
|
||||
v-for="key in keyPage.keys"
|
||||
:key="key.key_id"
|
||||
class="p-4 sm:p-5 hover:bg-muted/30 transition-colors"
|
||||
:class="{ 'opacity-50': !key.is_active }"
|
||||
>
|
||||
<div class="flex items-center gap-3">
|
||||
<input
|
||||
type="checkbox"
|
||||
:checked="selectedKeys.has(key.key_id)"
|
||||
class="rounded shrink-0"
|
||||
@change="toggleSelect(key.key_id)"
|
||||
>
|
||||
<div class="flex-1 min-w-0">
|
||||
<div class="flex items-center gap-2">
|
||||
<span class="text-sm font-medium truncate">
|
||||
{{ key.key_name || '未命名' }}
|
||||
</span>
|
||||
<Badge
|
||||
:variant="key.is_active ? (key.cooldown_reason ? 'destructive' : 'default') : 'secondary'"
|
||||
class="text-[10px] shrink-0"
|
||||
>
|
||||
{{ key.is_active ? (key.cooldown_reason ? '冷却' : '活跃') : '禁用' }}
|
||||
</Badge>
|
||||
</div>
|
||||
</div>
|
||||
<div class="flex items-center gap-0.5 shrink-0">
|
||||
<Button
|
||||
v-if="key.cooldown_reason"
|
||||
variant="ghost"
|
||||
size="icon"
|
||||
class="h-7 w-7 text-muted-foreground hover:text-green-600"
|
||||
title="清除冷却"
|
||||
@click="clearCooldown(key.key_id)"
|
||||
>
|
||||
<RefreshCw class="w-3.5 h-3.5" />
|
||||
</Button>
|
||||
<Button
|
||||
variant="ghost"
|
||||
size="icon"
|
||||
class="h-7 w-7 text-muted-foreground hover:text-foreground"
|
||||
:title="key.is_active ? '禁用' : '启用'"
|
||||
@click="toggleKeyActive(key)"
|
||||
>
|
||||
<component
|
||||
:is="key.is_active ? Ban : Check"
|
||||
class="w-3.5 h-3.5"
|
||||
/>
|
||||
</Button>
|
||||
</div>
|
||||
</div>
|
||||
<div class="mt-2.5 ml-7 grid grid-cols-3 gap-2">
|
||||
<div class="p-2 bg-muted/50 rounded-lg text-xs">
|
||||
<div class="text-muted-foreground mb-0.5">
|
||||
冷却
|
||||
</div>
|
||||
<div
|
||||
v-if="key.cooldown_reason"
|
||||
class="font-medium text-destructive text-[11px]"
|
||||
>
|
||||
{{ formatCooldownReason(key.cooldown_reason) }}
|
||||
<span
|
||||
v-if="key.cooldown_ttl_seconds"
|
||||
class="text-muted-foreground font-normal"
|
||||
>
|
||||
{{ formatTTL(key.cooldown_ttl_seconds) }}
|
||||
</span>
|
||||
</div>
|
||||
<div
|
||||
v-else
|
||||
class="text-muted-foreground"
|
||||
>
|
||||
-
|
||||
</div>
|
||||
</div>
|
||||
<div class="p-2 bg-muted/50 rounded-lg text-xs">
|
||||
<div class="text-muted-foreground mb-0.5">
|
||||
成本
|
||||
</div>
|
||||
<div
|
||||
v-if="key.cost_limit != null"
|
||||
class="font-medium tabular-nums text-[11px]"
|
||||
>
|
||||
{{ formatTokens(key.cost_window_usage) }}/{{ formatTokens(key.cost_limit) }}
|
||||
</div>
|
||||
<div
|
||||
v-else-if="key.cost_window_usage > 0"
|
||||
class="tabular-nums text-[11px]"
|
||||
>
|
||||
{{ formatTokens(key.cost_window_usage) }}
|
||||
</div>
|
||||
<div
|
||||
v-else
|
||||
class="text-muted-foreground"
|
||||
>
|
||||
-
|
||||
</div>
|
||||
</div>
|
||||
<div class="p-2 bg-muted/50 rounded-lg text-xs">
|
||||
<div class="text-muted-foreground mb-0.5">
|
||||
最后使用
|
||||
</div>
|
||||
<div class="text-[11px]">
|
||||
{{ key.last_used_at ? formatRelativeTime(key.last_used_at) : '-' }}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Empty keys -->
|
||||
<div
|
||||
v-if="keyPage.keys.length === 0 && !keysLoading"
|
||||
class="flex flex-col items-center justify-center py-16 text-center"
|
||||
>
|
||||
<div class="mx-auto flex h-16 w-16 items-center justify-center rounded-full bg-muted">
|
||||
<KeyRound class="h-8 w-8 text-muted-foreground" />
|
||||
</div>
|
||||
<p class="text-sm text-muted-foreground mt-4">
|
||||
暂无账号
|
||||
</p>
|
||||
<Button
|
||||
variant="outline"
|
||||
size="sm"
|
||||
class="mt-3"
|
||||
@click="showImportDialog = true"
|
||||
>
|
||||
<Upload class="w-3.5 h-3.5 mr-1.5" />
|
||||
批量导入
|
||||
</Button>
|
||||
</div>
|
||||
|
||||
<!-- Pagination -->
|
||||
<Pagination
|
||||
v-if="keyPage.keys.length > 0"
|
||||
:current="currentPage"
|
||||
:total="keyPage.total"
|
||||
:page-size="pageSize"
|
||||
cache-key="pool-keys-page-size"
|
||||
@update:current="currentPage = $event"
|
||||
@update:page-size="pageSize = $event"
|
||||
/>
|
||||
</template>
|
||||
</Card>
|
||||
|
||||
<!-- Dialogs -->
|
||||
<PoolImportDialog
|
||||
v-if="selectedProviderId"
|
||||
v-model="showImportDialog"
|
||||
:provider-id="selectedProviderId"
|
||||
@imported="loadKeys"
|
||||
/>
|
||||
<PoolConfigDialog
|
||||
v-if="selectedProviderId"
|
||||
v-model="showConfigDialog"
|
||||
:provider-id="selectedProviderId"
|
||||
:provider-type="selectedProviderData?.provider_type"
|
||||
:current-config="selectedProviderConfig"
|
||||
:current-claude-config="selectedProviderData?.claude_code_advanced"
|
||||
@saved="loadOverview"
|
||||
/>
|
||||
</div>
|
||||
</template>
|
||||
|
||||
<script setup lang="ts">
|
||||
import { ref, computed, watch, onMounted } from 'vue'
|
||||
import { Search, Upload, Settings, RefreshCw, Ban, Check, Database, KeyRound } from 'lucide-vue-next'
|
||||
|
||||
import {
|
||||
Card,
|
||||
Badge,
|
||||
Button,
|
||||
Input,
|
||||
Select,
|
||||
SelectTrigger,
|
||||
SelectValue,
|
||||
SelectContent,
|
||||
SelectItem,
|
||||
Table,
|
||||
TableHeader,
|
||||
TableBody,
|
||||
TableRow,
|
||||
TableHead,
|
||||
TableCell,
|
||||
Pagination,
|
||||
} from '@/components/ui'
|
||||
import RefreshButton from '@/components/ui/refresh-button.vue'
|
||||
import { useToast } from '@/composables/useToast'
|
||||
import { parseApiError } from '@/utils/errorParser'
|
||||
import {
|
||||
getPoolOverview,
|
||||
listPoolKeys,
|
||||
clearPoolCooldown,
|
||||
batchActionPoolKeys,
|
||||
} from '@/api/endpoints/pool'
|
||||
import type {
|
||||
PoolOverviewItem,
|
||||
PoolKeyDetail,
|
||||
PoolKeysPageResponse,
|
||||
} from '@/api/endpoints/pool'
|
||||
import type { PoolAdvancedConfig, ProviderWithEndpointsSummary } from '@/api/endpoints/types/provider'
|
||||
import { getProvider } from '@/api/endpoints'
|
||||
import PoolImportDialog from '@/features/pool/components/PoolImportDialog.vue'
|
||||
import PoolConfigDialog from '@/features/pool/components/PoolConfigDialog.vue'
|
||||
|
||||
const { success, error: showError } = useToast()
|
||||
|
||||
// --- Overview ---
|
||||
const poolProviders = ref<PoolOverviewItem[]>([])
|
||||
const overviewLoading = ref(true)
|
||||
|
||||
async function loadOverview() {
|
||||
overviewLoading.value = true
|
||||
try {
|
||||
const res = await getPoolOverview()
|
||||
poolProviders.value = res.items
|
||||
// Auto-select first provider if none selected
|
||||
if (!selectedProviderId.value && res.items.length > 0) {
|
||||
await selectProvider(res.items[0].provider_id)
|
||||
}
|
||||
} catch (err) {
|
||||
showError(parseApiError(err))
|
||||
} finally {
|
||||
overviewLoading.value = false
|
||||
}
|
||||
}
|
||||
|
||||
// --- Provider Selection ---
|
||||
const selectedProviderId = ref<string | null>(null)
|
||||
const selectedProviderData = ref<ProviderWithEndpointsSummary | null>(null)
|
||||
|
||||
// Proxy for Select v-model (string, not string|null)
|
||||
const selectedProviderIdProxy = computed({
|
||||
get: () => selectedProviderId.value ?? '',
|
||||
set: (val: string) => {
|
||||
if (val && val !== selectedProviderId.value) {
|
||||
selectProvider(val)
|
||||
}
|
||||
},
|
||||
})
|
||||
|
||||
const selectedProviderConfig = computed<PoolAdvancedConfig | null>(() => {
|
||||
return (selectedProviderData.value as Record<string, unknown> | null)?.pool_advanced as PoolAdvancedConfig | null ?? null
|
||||
})
|
||||
|
||||
async function selectProvider(id: string) {
|
||||
selectedProviderId.value = id
|
||||
selectedKeys.value.clear()
|
||||
currentPage.value = 1
|
||||
searchQuery.value = ''
|
||||
statusFilter.value = 'all'
|
||||
await Promise.all([loadKeys(), loadProviderData(id)])
|
||||
}
|
||||
|
||||
async function loadProviderData(id: string) {
|
||||
try {
|
||||
selectedProviderData.value = await getProvider(id)
|
||||
} catch {
|
||||
selectedProviderData.value = null
|
||||
}
|
||||
}
|
||||
|
||||
async function refresh() {
|
||||
await loadOverview()
|
||||
if (selectedProviderId.value) {
|
||||
await loadKeys()
|
||||
}
|
||||
}
|
||||
|
||||
// --- Keys ---
|
||||
const keyPage = ref<PoolKeysPageResponse>({ total: 0, page: 1, page_size: 50, keys: [] })
|
||||
const keysLoading = ref(false)
|
||||
const searchQuery = ref('')
|
||||
const statusFilter = ref('all')
|
||||
const currentPage = ref(1)
|
||||
const pageSize = ref(50)
|
||||
|
||||
async function loadKeys() {
|
||||
if (!selectedProviderId.value) return
|
||||
keysLoading.value = true
|
||||
try {
|
||||
keyPage.value = await listPoolKeys(selectedProviderId.value, {
|
||||
page: currentPage.value,
|
||||
page_size: pageSize.value,
|
||||
search: searchQuery.value || undefined,
|
||||
status: statusFilter.value as 'all' | 'active' | 'cooldown' | 'inactive',
|
||||
})
|
||||
} catch (err) {
|
||||
showError(parseApiError(err))
|
||||
} finally {
|
||||
keysLoading.value = false
|
||||
}
|
||||
}
|
||||
|
||||
watch([currentPage, pageSize], () => loadKeys())
|
||||
watch([searchQuery, statusFilter], () => {
|
||||
currentPage.value = 1
|
||||
loadKeys()
|
||||
})
|
||||
|
||||
// --- Key Selection ---
|
||||
const selectedKeys = ref(new Set<string>())
|
||||
|
||||
const allSelected = computed(() => {
|
||||
if (keyPage.value.keys.length === 0) return false
|
||||
return keyPage.value.keys.every(k => selectedKeys.value.has(k.key_id))
|
||||
})
|
||||
|
||||
function toggleSelectAll() {
|
||||
if (allSelected.value) {
|
||||
selectedKeys.value.clear()
|
||||
} else {
|
||||
keyPage.value.keys.forEach(k => selectedKeys.value.add(k.key_id))
|
||||
}
|
||||
}
|
||||
|
||||
function toggleSelect(id: string) {
|
||||
if (selectedKeys.value.has(id)) {
|
||||
selectedKeys.value.delete(id)
|
||||
} else {
|
||||
selectedKeys.value.add(id)
|
||||
}
|
||||
}
|
||||
|
||||
// --- Actions ---
|
||||
async function clearCooldown(keyId: string) {
|
||||
if (!selectedProviderId.value) return
|
||||
try {
|
||||
const res = await clearPoolCooldown(selectedProviderId.value, keyId)
|
||||
success(res.message)
|
||||
await loadKeys()
|
||||
} catch (err) {
|
||||
showError(parseApiError(err))
|
||||
}
|
||||
}
|
||||
|
||||
async function toggleKeyActive(key: PoolKeyDetail) {
|
||||
if (!selectedProviderId.value) return
|
||||
try {
|
||||
const action = key.is_active ? 'disable' : 'enable'
|
||||
await batchActionPoolKeys(selectedProviderId.value, {
|
||||
key_ids: [key.key_id],
|
||||
action,
|
||||
})
|
||||
await loadKeys()
|
||||
} catch (err) {
|
||||
showError(parseApiError(err))
|
||||
}
|
||||
}
|
||||
|
||||
async function batchAction(action: 'enable' | 'disable' | 'delete' | 'clear_cooldown' | 'reset_cost') {
|
||||
if (!selectedProviderId.value || selectedKeys.value.size === 0) return
|
||||
try {
|
||||
const res = await batchActionPoolKeys(selectedProviderId.value, {
|
||||
key_ids: Array.from(selectedKeys.value),
|
||||
action,
|
||||
})
|
||||
success(res.message)
|
||||
selectedKeys.value.clear()
|
||||
await loadKeys()
|
||||
} catch (err) {
|
||||
showError(parseApiError(err))
|
||||
}
|
||||
}
|
||||
|
||||
// --- Dialogs ---
|
||||
const showImportDialog = ref(false)
|
||||
const showConfigDialog = ref(false)
|
||||
|
||||
// --- Formatting ---
|
||||
const COOLDOWN_REASON_MAP: Record<string, string> = {
|
||||
rate_limited_429: '429 限流',
|
||||
forbidden_403: '403 禁止',
|
||||
overloaded_529: '529 过载',
|
||||
auth_failed_401: '401 认证失败',
|
||||
payment_required_402: '402 欠费',
|
||||
server_error_500: '500 错误',
|
||||
}
|
||||
|
||||
function formatCooldownReason(reason: string): string {
|
||||
return COOLDOWN_REASON_MAP[reason] || reason
|
||||
}
|
||||
|
||||
function formatTTL(seconds: number): string {
|
||||
if (seconds <= 0) return ''
|
||||
const m = Math.floor(seconds / 60)
|
||||
const s = seconds % 60
|
||||
return m > 0 ? `${m}m ${s}s` : `${s}s`
|
||||
}
|
||||
|
||||
function formatTokens(tokens: number): string {
|
||||
if (tokens >= 1_000_000) return `${(tokens / 1_000_000).toFixed(1)}M`
|
||||
if (tokens >= 1_000) return `${(tokens / 1_000).toFixed(1)}K`
|
||||
return String(tokens)
|
||||
}
|
||||
|
||||
function getCostBarColor(usage: number, limit: number): string {
|
||||
const ratio = usage / limit
|
||||
if (ratio >= 0.9) return 'bg-red-500'
|
||||
if (ratio >= 0.7) return 'bg-yellow-500'
|
||||
return 'bg-green-500'
|
||||
}
|
||||
|
||||
function formatRelativeTime(isoStr: string): string {
|
||||
const diff = (Date.now() - new Date(isoStr).getTime()) / 1000
|
||||
if (diff < 60) return '刚刚'
|
||||
if (diff < 3600) return `${Math.floor(diff / 60)}m 前`
|
||||
if (diff < 86400) return `${Math.floor(diff / 3600)}h 前`
|
||||
return `${Math.floor(diff / 86400)}d 前`
|
||||
}
|
||||
|
||||
// --- Init ---
|
||||
onMounted(async () => {
|
||||
await loadOverview()
|
||||
})
|
||||
</script>
|
||||
5
src/api/admin/pool/__init__.py
Normal file
5
src/api/admin/pool/__init__.py
Normal file
@@ -0,0 +1,5 @@
|
||||
"""Pool management admin API."""
|
||||
|
||||
from .routes import router
|
||||
|
||||
__all__ = ["router"]
|
||||
450
src/api/admin/pool/routes.py
Normal file
450
src/api/admin/pool/routes.py
Normal file
@@ -0,0 +1,450 @@
|
||||
"""Pool management admin API routes.
|
||||
|
||||
Provides endpoints for managing account pools at scale:
|
||||
- Overview of all pool-enabled providers
|
||||
- Paginated key listing with search/filter
|
||||
- Batch import / batch actions
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
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.orm import Session
|
||||
|
||||
from src.api.base.admin_adapter import AdminApiAdapter
|
||||
from src.api.base.context import ApiRequestContext
|
||||
from src.api.base.pipeline import ApiRequestPipeline
|
||||
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.services.provider.pool import redis_ops as pool_redis
|
||||
from src.services.provider.pool.config import parse_pool_config
|
||||
|
||||
from .schemas import (
|
||||
BatchActionRequest,
|
||||
BatchActionResponse,
|
||||
BatchImportError,
|
||||
BatchImportRequest,
|
||||
BatchImportResponse,
|
||||
PoolKeyDetail,
|
||||
PoolKeysPageResponse,
|
||||
PoolOverviewItem,
|
||||
PoolOverviewResponse,
|
||||
)
|
||||
|
||||
router = APIRouter(prefix="/api/admin/pool", tags=["pool-management"])
|
||||
pipeline = ApiRequestPipeline()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GET /api/admin/pool/overview
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@router.get("/overview", response_model=PoolOverviewResponse)
|
||||
async def pool_overview(
|
||||
request: Request,
|
||||
db: Session = Depends(get_db),
|
||||
) -> PoolOverviewResponse:
|
||||
"""Return all pool-enabled providers with summary stats."""
|
||||
adapter = AdminPoolOverviewAdapter()
|
||||
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GET /api/admin/pool/{provider_id}/keys
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@router.get("/{provider_id}/keys", response_model=PoolKeysPageResponse)
|
||||
async def list_pool_keys(
|
||||
provider_id: str,
|
||||
request: Request,
|
||||
page: int = Query(1, ge=1),
|
||||
page_size: int = Query(50, ge=1, le=200),
|
||||
search: str = Query("", description="Search by key name"),
|
||||
status: str = Query("all", description="all/active/cooldown/inactive"),
|
||||
db: Session = Depends(get_db),
|
||||
) -> PoolKeysPageResponse:
|
||||
"""Server-side paginated account list for a pool-enabled provider."""
|
||||
adapter = AdminListPoolKeysAdapter(
|
||||
provider_id=provider_id,
|
||||
page=page,
|
||||
page_size=page_size,
|
||||
search=search,
|
||||
status=status,
|
||||
)
|
||||
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# POST /api/admin/pool/{provider_id}/keys/batch-import
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@router.post("/{provider_id}/keys/batch-import", response_model=BatchImportResponse)
|
||||
async def batch_import_keys(
|
||||
provider_id: str,
|
||||
body: BatchImportRequest,
|
||||
request: Request,
|
||||
db: Session = Depends(get_db),
|
||||
) -> BatchImportResponse:
|
||||
"""Batch import keys into a provider's pool."""
|
||||
adapter = AdminBatchImportKeysAdapter(provider_id=provider_id, body=body)
|
||||
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# POST /api/admin/pool/{provider_id}/keys/batch-action
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
ALLOWED_ACTIONS = {"enable", "disable", "delete", "clear_cooldown", "reset_cost"}
|
||||
|
||||
|
||||
@router.post("/{provider_id}/keys/batch-action", response_model=BatchActionResponse)
|
||||
async def batch_action_keys(
|
||||
provider_id: str,
|
||||
body: BatchActionRequest,
|
||||
request: Request,
|
||||
db: Session = Depends(get_db),
|
||||
) -> BatchActionResponse:
|
||||
"""Batch enable/disable/delete/clear_cooldown/reset_cost on pool keys."""
|
||||
adapter = AdminBatchActionKeysAdapter(provider_id=provider_id, body=body)
|
||||
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Adapters
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class AdminPoolOverviewAdapter(AdminApiAdapter):
|
||||
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
|
||||
db = context.db
|
||||
providers = (
|
||||
db.query(Provider)
|
||||
.filter(Provider.is_active.is_(True))
|
||||
.order_by(Provider.provider_priority.asc())
|
||||
.all()
|
||||
)
|
||||
|
||||
items: list[PoolOverviewItem] = []
|
||||
for p in providers:
|
||||
pid = str(p.id)
|
||||
pcfg = parse_pool_config(getattr(p, "config", None))
|
||||
|
||||
# Non-pool providers: skip Redis + key queries entirely.
|
||||
if pcfg is None:
|
||||
items.append(
|
||||
PoolOverviewItem(
|
||||
provider_id=pid,
|
||||
provider_name=p.name,
|
||||
provider_type=str(getattr(p, "provider_type", "custom") or "custom"),
|
||||
pool_enabled=False,
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
keys = db.query(ProviderAPIKey).filter(ProviderAPIKey.provider_id == pid).all()
|
||||
key_ids = [str(k.id) for k in keys]
|
||||
|
||||
cooldown_count = 0
|
||||
if key_ids:
|
||||
cooldowns = await pool_redis.batch_get_cooldowns(pid, key_ids)
|
||||
cooldown_count = sum(1 for v in cooldowns.values() if v is not None)
|
||||
|
||||
items.append(
|
||||
PoolOverviewItem(
|
||||
provider_id=pid,
|
||||
provider_name=p.name,
|
||||
provider_type=str(getattr(p, "provider_type", "custom") or "custom"),
|
||||
total_keys=len(keys),
|
||||
active_keys=sum(1 for k in keys if k.is_active),
|
||||
cooldown_count=cooldown_count,
|
||||
pool_enabled=True,
|
||||
)
|
||||
)
|
||||
|
||||
return PoolOverviewResponse(items=items)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AdminListPoolKeysAdapter(AdminApiAdapter):
|
||||
provider_id: str = ""
|
||||
page: int = 1
|
||||
page_size: int = 50
|
||||
search: str = ""
|
||||
status: str = "all"
|
||||
|
||||
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")
|
||||
|
||||
pcfg = parse_pool_config(getattr(provider, "config", None))
|
||||
pid = str(provider.id)
|
||||
|
||||
# Base query
|
||||
q = db.query(ProviderAPIKey).filter(ProviderAPIKey.provider_id == pid)
|
||||
|
||||
if self.search:
|
||||
escaped = self.search.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
|
||||
q = q.filter(ProviderAPIKey.name.ilike(f"%{escaped}%"))
|
||||
|
||||
if self.status == "active":
|
||||
q = q.filter(ProviderAPIKey.is_active.is_(True))
|
||||
elif self.status == "inactive":
|
||||
q = q.filter(ProviderAPIKey.is_active.is_(False))
|
||||
# "cooldown" filtering is done post-query (Redis state)
|
||||
|
||||
total = q.count()
|
||||
|
||||
# For cooldown filtering we need to fetch all, then filter, then paginate.
|
||||
# Limit scan range to avoid loading the entire table into memory.
|
||||
if self.status == "cooldown":
|
||||
_max_scan = 2000
|
||||
all_keys = q.order_by(ProviderAPIKey.created_at.desc()).limit(_max_scan).all()
|
||||
key_ids = [str(k.id) for k in all_keys]
|
||||
cooldowns = await pool_redis.batch_get_cooldowns(pid, key_ids) if key_ids else {}
|
||||
all_keys = [k for k in all_keys if cooldowns.get(str(k.id)) is not None]
|
||||
total = len(all_keys)
|
||||
offset = (self.page - 1) * self.page_size
|
||||
keys = all_keys[offset : offset + self.page_size]
|
||||
else:
|
||||
offset = (self.page - 1) * self.page_size
|
||||
keys = (
|
||||
q.order_by(ProviderAPIKey.created_at.desc())
|
||||
.offset(offset)
|
||||
.limit(self.page_size)
|
||||
.all()
|
||||
)
|
||||
|
||||
# Batch fetch Redis state (parallel where possible)
|
||||
key_ids = [str(k.id) for k in keys]
|
||||
if key_ids:
|
||||
_lru_coro = (
|
||||
pool_redis.get_lru_scores(pid, key_ids)
|
||||
if pcfg and pcfg.lru_enabled
|
||||
else asyncio.sleep(0, result={})
|
||||
)
|
||||
_cost_coro = (
|
||||
pool_redis.batch_get_cost_totals(pid, key_ids, pcfg.cost_window_seconds)
|
||||
if pcfg
|
||||
else asyncio.sleep(0, result={})
|
||||
)
|
||||
cooldowns, cooldown_ttls, lru_scores, cost_totals = await asyncio.gather(
|
||||
pool_redis.batch_get_cooldowns(pid, key_ids),
|
||||
pool_redis.batch_get_cooldown_ttls(pid, key_ids),
|
||||
_lru_coro,
|
||||
_cost_coro,
|
||||
)
|
||||
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)
|
||||
)
|
||||
sticky_counts = dict(zip(key_ids, counts))
|
||||
|
||||
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
|
||||
|
||||
key_details.append(
|
||||
PoolKeyDetail(
|
||||
key_id=kid,
|
||||
key_name=k.name or "",
|
||||
is_active=bool(k.is_active),
|
||||
auth_type=str(getattr(k, "auth_type", "api_key") or "api_key"),
|
||||
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,
|
||||
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
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
return PoolKeysPageResponse(
|
||||
total=total,
|
||||
page=self.page,
|
||||
page_size=self.page_size,
|
||||
keys=key_details,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AdminBatchImportKeysAdapter(AdminApiAdapter):
|
||||
provider_id: str = ""
|
||||
body: BatchImportRequest = field(default_factory=lambda: BatchImportRequest(keys=[]))
|
||||
|
||||
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")
|
||||
|
||||
imported = 0
|
||||
skipped = 0
|
||||
errors: list[BatchImportError] = []
|
||||
now = datetime.now(timezone.utc)
|
||||
|
||||
for idx, item in enumerate(self.body.keys):
|
||||
if not item.api_key.strip():
|
||||
errors.append(BatchImportError(index=idx, reason="api_key is empty"))
|
||||
continue
|
||||
|
||||
try:
|
||||
encrypted_key = crypto_service.encrypt(item.api_key)
|
||||
new_key = ProviderAPIKey(
|
||||
id=str(uuid.uuid4()),
|
||||
provider_id=self.provider_id,
|
||||
name=item.name or f"imported-{idx}",
|
||||
api_key=encrypted_key,
|
||||
auth_type=item.auth_type or "api_key",
|
||||
is_active=True,
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
)
|
||||
db.add(new_key)
|
||||
imported += 1
|
||||
except Exception as exc:
|
||||
logger.warning("batch import key #{} failed: {}", idx, exc)
|
||||
errors.append(BatchImportError(index=idx, reason=str(exc)))
|
||||
|
||||
if imported > 0:
|
||||
try:
|
||||
db.commit()
|
||||
except Exception as exc:
|
||||
db.rollback()
|
||||
logger.error("batch import commit failed: {}", exc)
|
||||
return BatchImportResponse(
|
||||
imported=0,
|
||||
skipped=skipped,
|
||||
errors=[BatchImportError(index=-1, reason=f"commit failed: {exc}")],
|
||||
)
|
||||
|
||||
admin_name = context.user.username if context.user else "admin"
|
||||
logger.info(
|
||||
"Pool batch import by {}: provider={}, imported={}, skipped={}, errors={}",
|
||||
admin_name,
|
||||
self.provider_id[:8],
|
||||
imported,
|
||||
skipped,
|
||||
len(errors),
|
||||
)
|
||||
|
||||
return BatchImportResponse(imported=imported, skipped=skipped, errors=errors)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AdminBatchActionKeysAdapter(AdminApiAdapter):
|
||||
provider_id: str = ""
|
||||
body: BatchActionRequest = field(
|
||||
default_factory=lambda: BatchActionRequest(key_ids=[], action="")
|
||||
)
|
||||
|
||||
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
|
||||
from fastapi import HTTPException
|
||||
|
||||
if self.body.action not in ALLOWED_ACTIONS:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=(
|
||||
f"Invalid action: {self.body.action}. "
|
||||
f"Allowed: {', '.join(sorted(ALLOWED_ACTIONS))}"
|
||||
),
|
||||
)
|
||||
|
||||
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)
|
||||
affected = 0
|
||||
|
||||
keys = (
|
||||
db.query(ProviderAPIKey)
|
||||
.filter(
|
||||
ProviderAPIKey.provider_id == pid,
|
||||
ProviderAPIKey.id.in_(self.body.key_ids),
|
||||
)
|
||||
.all()
|
||||
)
|
||||
|
||||
for key in keys:
|
||||
kid = str(key.id)
|
||||
|
||||
if self.body.action == "enable":
|
||||
key.is_active = True
|
||||
affected += 1
|
||||
|
||||
elif self.body.action == "disable":
|
||||
key.is_active = False
|
||||
affected += 1
|
||||
|
||||
elif self.body.action == "delete":
|
||||
db.delete(key)
|
||||
affected += 1
|
||||
|
||||
elif self.body.action == "clear_cooldown":
|
||||
await pool_redis.clear_cooldown(pid, kid)
|
||||
affected += 1
|
||||
|
||||
elif self.body.action == "reset_cost":
|
||||
await pool_redis.clear_cost(pid, kid)
|
||||
affected += 1
|
||||
|
||||
if self.body.action in {"enable", "disable", "delete"}:
|
||||
try:
|
||||
db.commit()
|
||||
except Exception as exc:
|
||||
db.rollback()
|
||||
logger.error("batch action commit failed: {}", exc)
|
||||
return BatchActionResponse(affected=0, message=f"commit failed: {exc}")
|
||||
|
||||
action_labels = {
|
||||
"enable": "enabled",
|
||||
"disable": "disabled",
|
||||
"delete": "deleted",
|
||||
"clear_cooldown": "cooldown cleared",
|
||||
"reset_cost": "cost reset",
|
||||
}
|
||||
|
||||
admin_name = context.user.username if context.user else "admin"
|
||||
affected_ids = [str(k.id)[:8] for k in keys]
|
||||
logger.info(
|
||||
"Pool batch action by {}: provider={}, action={}, affected={}, key_ids={}",
|
||||
admin_name,
|
||||
self.provider_id[:8],
|
||||
self.body.action,
|
||||
affected,
|
||||
affected_ids,
|
||||
)
|
||||
|
||||
return BatchActionResponse(
|
||||
affected=affected,
|
||||
message=f"{affected} keys {action_labels.get(self.body.action, self.body.action)}",
|
||||
)
|
||||
103
src/api/admin/pool/schemas.py
Normal file
103
src/api/admin/pool/schemas.py
Normal file
@@ -0,0 +1,103 @@
|
||||
"""Pydantic schemas for Pool management API."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Overview
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class PoolOverviewItem(BaseModel):
|
||||
"""One Provider in the overview list."""
|
||||
|
||||
provider_id: str
|
||||
provider_name: str
|
||||
provider_type: str = "custom"
|
||||
total_keys: int = 0
|
||||
active_keys: int = 0
|
||||
cooldown_count: int = 0
|
||||
pool_enabled: bool = False
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
class PoolOverviewResponse(BaseModel):
|
||||
items: list[PoolOverviewItem] = Field(default_factory=list)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Paginated key list
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class PoolKeyDetail(BaseModel):
|
||||
"""Detailed status of a single pool key."""
|
||||
|
||||
key_id: str
|
||||
key_name: str
|
||||
is_active: bool
|
||||
auth_type: str = "api_key"
|
||||
cooldown_reason: str | None = None
|
||||
cooldown_ttl_seconds: int | None = None
|
||||
cost_window_usage: int = 0
|
||||
cost_limit: int | None = None
|
||||
sticky_sessions: int = 0
|
||||
lru_score: float | None = None
|
||||
created_at: str | None = None
|
||||
last_used_at: str | None = None
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
class PoolKeysPageResponse(BaseModel):
|
||||
"""Server-side paginated key list."""
|
||||
|
||||
total: int
|
||||
page: int
|
||||
page_size: int
|
||||
keys: list[PoolKeyDetail] = Field(default_factory=list)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Batch import
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class PoolKeyImportItem(BaseModel):
|
||||
"""Single key to import."""
|
||||
|
||||
name: str
|
||||
api_key: str
|
||||
auth_type: str = "api_key"
|
||||
|
||||
|
||||
class BatchImportRequest(BaseModel):
|
||||
keys: list[PoolKeyImportItem] = Field(..., max_length=500)
|
||||
|
||||
|
||||
class BatchImportError(BaseModel):
|
||||
index: int
|
||||
reason: str
|
||||
|
||||
|
||||
class BatchImportResponse(BaseModel):
|
||||
imported: int = 0
|
||||
skipped: int = 0
|
||||
errors: list[BatchImportError] = Field(default_factory=list)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Batch action
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class BatchActionRequest(BaseModel):
|
||||
key_ids: list[str] = Field(..., max_length=500)
|
||||
action: str # enable / disable / delete / clear_cooldown / reset_cost
|
||||
|
||||
|
||||
class BatchActionResponse(BaseModel):
|
||||
affected: int = 0
|
||||
message: str = ""
|
||||
166
src/core/oauth_plan.py
Normal file
166
src/core/oauth_plan.py
Normal file
@@ -0,0 +1,166 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
from src.core.crypto import crypto_service
|
||||
|
||||
|
||||
def normalize_oauth_plan_type(plan_type: Any) -> str | None:
|
||||
if not isinstance(plan_type, str):
|
||||
return None
|
||||
normalized = plan_type.strip().lower()
|
||||
if not normalized:
|
||||
return None
|
||||
return normalized
|
||||
|
||||
|
||||
def extract_oauth_plan_type_from_auth_config_data(auth_config: Any) -> str | None:
|
||||
if not isinstance(auth_config, dict):
|
||||
return None
|
||||
|
||||
# Codex: plan_type (free/plus/team/enterprise)
|
||||
plan_type = normalize_oauth_plan_type(auth_config.get("plan_type"))
|
||||
if plan_type:
|
||||
return plan_type
|
||||
|
||||
# Antigravity: tier (PAID/FREE/...)
|
||||
tier = normalize_oauth_plan_type(auth_config.get("tier"))
|
||||
if tier:
|
||||
return tier
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def decrypt_auth_config_to_dict(
|
||||
encrypted_auth_config: str | None,
|
||||
*,
|
||||
silent: bool = True,
|
||||
) -> dict[str, Any] | None:
|
||||
if not encrypted_auth_config:
|
||||
return None
|
||||
try:
|
||||
decrypted = crypto_service.decrypt(encrypted_auth_config, silent=silent)
|
||||
parsed = json.loads(decrypted)
|
||||
if isinstance(parsed, dict):
|
||||
return parsed
|
||||
except Exception:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _strip_provider_prefix(value: str, provider_type: str | None = None) -> str:
|
||||
normalized = value.strip()
|
||||
if not normalized:
|
||||
return ""
|
||||
|
||||
prefixes: list[str] = []
|
||||
if isinstance(provider_type, str) and provider_type.strip():
|
||||
prefixes.append(provider_type.strip())
|
||||
# Kiro 目前将套餐信息记录为 "KIRO FREE" / "KIRO PRO+"。
|
||||
if "kiro" not in {p.lower() for p in prefixes}:
|
||||
prefixes.append("kiro")
|
||||
|
||||
upper = normalized.upper()
|
||||
for prefix in prefixes:
|
||||
prefix_upper = prefix.upper()
|
||||
if upper == prefix_upper:
|
||||
return ""
|
||||
if upper.startswith(f"{prefix_upper} "):
|
||||
return normalized[len(prefix) :].strip()
|
||||
|
||||
return normalized
|
||||
|
||||
|
||||
def extract_oauth_plan_type_from_upstream_metadata(
|
||||
upstream_metadata: Any,
|
||||
*,
|
||||
provider_type: str | None = None,
|
||||
) -> str | None:
|
||||
if not isinstance(upstream_metadata, dict):
|
||||
return None
|
||||
|
||||
kiro_meta = upstream_metadata.get("kiro")
|
||||
if isinstance(kiro_meta, dict):
|
||||
subscription_title = kiro_meta.get("subscription_title")
|
||||
if isinstance(subscription_title, str):
|
||||
normalized = _strip_provider_prefix(subscription_title, provider_type=provider_type)
|
||||
return normalize_oauth_plan_type(normalized)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def extract_oauth_plan_type(
|
||||
encrypted_auth_config: str | None,
|
||||
*,
|
||||
upstream_metadata: Any = None,
|
||||
provider_type: str | None = None,
|
||||
silent: bool = True,
|
||||
) -> str | None:
|
||||
auth_config = decrypt_auth_config_to_dict(encrypted_auth_config, silent=silent)
|
||||
plan_type = extract_oauth_plan_type_from_auth_config_data(auth_config)
|
||||
if plan_type:
|
||||
return plan_type
|
||||
return extract_oauth_plan_type_from_upstream_metadata(
|
||||
upstream_metadata, provider_type=provider_type
|
||||
)
|
||||
|
||||
|
||||
def normalize_antigravity_tier(raw_tier: Any) -> str | None:
|
||||
normalized = normalize_oauth_plan_type(raw_tier)
|
||||
if not normalized:
|
||||
return None
|
||||
if "ultra" in normalized:
|
||||
return "ultra"
|
||||
if "pro" in normalized or "paid" in normalized:
|
||||
return "pro"
|
||||
if "free" in normalized or "legacy" in normalized:
|
||||
return "free"
|
||||
return normalized
|
||||
|
||||
|
||||
def _extract_antigravity_tier_raw(tier_obj: Any) -> str | None:
|
||||
if isinstance(tier_obj, str):
|
||||
stripped = tier_obj.strip()
|
||||
return stripped or None
|
||||
if isinstance(tier_obj, dict):
|
||||
for key in ("id", "tierType"):
|
||||
value = tier_obj.get(key)
|
||||
if isinstance(value, str):
|
||||
stripped = value.strip()
|
||||
if stripped:
|
||||
return stripped
|
||||
return None
|
||||
|
||||
|
||||
def _format_antigravity_tier_label(
|
||||
normalized_tier: str,
|
||||
*,
|
||||
fallback_raw: str | None = None,
|
||||
) -> str:
|
||||
if normalized_tier == "ultra":
|
||||
return "Ultra"
|
||||
if normalized_tier == "pro":
|
||||
return "Pro"
|
||||
if normalized_tier == "free":
|
||||
return "Free"
|
||||
if fallback_raw:
|
||||
return fallback_raw
|
||||
return normalized_tier
|
||||
|
||||
|
||||
def extract_antigravity_tier_from_code_assist(code_assist: Any) -> str:
|
||||
if not isinstance(code_assist, dict):
|
||||
return "Free"
|
||||
|
||||
paid_tier_raw = _extract_antigravity_tier_raw(code_assist.get("paidTier"))
|
||||
paid_tier = normalize_antigravity_tier(paid_tier_raw)
|
||||
if paid_tier:
|
||||
return _format_antigravity_tier_label(paid_tier, fallback_raw=paid_tier_raw)
|
||||
|
||||
current_tier_raw = _extract_antigravity_tier_raw(code_assist.get("currentTier"))
|
||||
current_tier = normalize_antigravity_tier(current_tier_raw)
|
||||
if current_tier:
|
||||
return _format_antigravity_tier_label(current_tier, fallback_raw=current_tier_raw)
|
||||
|
||||
return "Free"
|
||||
@@ -0,0 +1,6 @@
|
||||
"""Backward compat shim -- canonical definitions moved to src.services.provider.pool.config."""
|
||||
|
||||
from src.services.provider.pool.config import * # noqa: F401,F403
|
||||
from src.services.provider.pool.config import PoolConfig, UnschedulableRule, parse_pool_config
|
||||
|
||||
__all__ = ["PoolConfig", "UnschedulableRule", "parse_pool_config"]
|
||||
@@ -0,0 +1,10 @@
|
||||
"""Backward compat shim -- canonical definitions moved to src.services.provider.pool.cost_tracker."""
|
||||
|
||||
from src.services.provider.pool.cost_tracker import ( # noqa: F401
|
||||
get_window_usage,
|
||||
is_approaching_limit,
|
||||
is_at_limit,
|
||||
record_usage,
|
||||
)
|
||||
|
||||
__all__ = ["record_usage", "get_window_usage", "is_at_limit", "is_approaching_limit"]
|
||||
@@ -0,0 +1,6 @@
|
||||
"""Backward compat shim -- canonical definitions moved to src.services.provider.pool.health_policy."""
|
||||
|
||||
from src.services.provider.pool.health_policy import * # noqa: F401,F403
|
||||
from src.services.provider.pool.health_policy import apply_health_policy # noqa: F811
|
||||
|
||||
__all__ = ["apply_health_policy"]
|
||||
27
src/services/provider/adapters/claude_code/pool_hook.py
Normal file
27
src/services/provider/adapters/claude_code/pool_hook.py
Normal file
@@ -0,0 +1,27 @@
|
||||
"""Claude Code pool scheduling hook."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
class ClaudeCodePoolHook:
|
||||
"""Pool scheduling hook for Claude Code providers.
|
||||
|
||||
Extracts the session UUID from ``metadata.user_id`` which follows the
|
||||
pattern ``<user>_session_<uuid>``.
|
||||
"""
|
||||
|
||||
name = "claude_code"
|
||||
|
||||
def extract_session_uuid(self, request_body: dict[str, Any]) -> str | None:
|
||||
metadata = request_body.get("metadata")
|
||||
if isinstance(metadata, dict):
|
||||
user_id = metadata.get("user_id")
|
||||
if isinstance(user_id, str) and "_session_" in user_id:
|
||||
idx = user_id.rfind("_session_")
|
||||
return user_id[idx + len("_session_") :].strip() or None
|
||||
return None
|
||||
|
||||
|
||||
claude_code_pool_hook = ClaudeCodePoolHook()
|
||||
@@ -0,0 +1,8 @@
|
||||
"""Backward compat shim -- canonical definitions moved to src.services.provider.pool.manager."""
|
||||
|
||||
from src.services.provider.pool.manager import * # noqa: F401,F403
|
||||
from src.services.provider.pool.manager import PoolManager
|
||||
|
||||
ClaudeCodePoolManager = PoolManager # noqa: F811
|
||||
|
||||
__all__ = ["PoolManager", "ClaudeCodePoolManager"]
|
||||
9
src/services/provider/adapters/claude_code/pool_oauth.py
Normal file
9
src/services/provider/adapters/claude_code/pool_oauth.py
Normal file
@@ -0,0 +1,9 @@
|
||||
"""Backward compat shim -- canonical definitions moved to src.services.provider.pool.oauth_cache."""
|
||||
|
||||
from src.services.provider.pool.oauth_cache import ( # noqa: F401
|
||||
cache_token,
|
||||
get_cached_token,
|
||||
invalidate_token,
|
||||
)
|
||||
|
||||
__all__ = ["get_cached_token", "cache_token", "invalidate_token"]
|
||||
23
src/services/provider/adapters/claude_code/pool_redis.py
Normal file
23
src/services/provider/adapters/claude_code/pool_redis.py
Normal file
@@ -0,0 +1,23 @@
|
||||
"""Backward compat shim -- canonical definitions moved to src.services.provider.pool.redis_ops."""
|
||||
|
||||
from src.services.provider.pool.redis_ops import * # noqa: F401,F403
|
||||
from src.services.provider.pool.redis_ops import (
|
||||
add_cost_entry,
|
||||
batch_get_cooldowns,
|
||||
cache_oauth_token,
|
||||
clear_cooldown,
|
||||
clear_cost,
|
||||
delete_sticky_binding,
|
||||
get_cached_oauth_token,
|
||||
get_cooldown,
|
||||
get_cooldown_ttl,
|
||||
get_cost_window_total,
|
||||
get_key_sticky_count,
|
||||
get_lru_scores,
|
||||
get_sticky_binding,
|
||||
get_sticky_session_count,
|
||||
invalidate_oauth_token_cache,
|
||||
set_cooldown,
|
||||
set_sticky_binding,
|
||||
touch_lru,
|
||||
)
|
||||
14
src/services/provider/pool/__init__.py
Normal file
14
src/services/provider/pool/__init__.py
Normal file
@@ -0,0 +1,14 @@
|
||||
"""Generic Account Pool management for any Provider type.
|
||||
|
||||
Re-exports the main public API for convenience.
|
||||
"""
|
||||
|
||||
from src.services.provider.pool.config import PoolConfig, UnschedulableRule, parse_pool_config
|
||||
from src.services.provider.pool.manager import PoolManager
|
||||
|
||||
__all__ = [
|
||||
"PoolConfig",
|
||||
"PoolManager",
|
||||
"UnschedulableRule",
|
||||
"parse_pool_config",
|
||||
]
|
||||
138
src/services/provider/pool/config.py
Normal file
138
src/services/provider/pool/config.py
Normal file
@@ -0,0 +1,138 @@
|
||||
"""Account Pool configuration (provider-agnostic)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from src.core.logger import logger
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UnschedulableRule:
|
||||
"""Keyword-based temporary unschedule rule."""
|
||||
|
||||
keyword: str
|
||||
duration_minutes: int = 5
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PoolConfig:
|
||||
"""Parsed pool configuration for any Provider.
|
||||
|
||||
All transient state lives in Redis; this dataclass only holds
|
||||
the *configuration* that controls pool behaviour.
|
||||
"""
|
||||
|
||||
# -- Sticky Session -------------------------------------------------------
|
||||
sticky_session_ttl_seconds: int = 3600 # 1 hour
|
||||
|
||||
# -- Load-Aware Selection -------------------------------------------------
|
||||
load_threshold_percent: int = 80
|
||||
|
||||
# -- LRU ------------------------------------------------------------------
|
||||
lru_enabled: bool = True
|
||||
|
||||
# -- Rolling-Window Cost Tracking -----------------------------------------
|
||||
cost_window_seconds: int = 18000 # 5 hours
|
||||
cost_limit_per_key_tokens: int | None = None # None = unlimited
|
||||
cost_soft_threshold_percent: int = 80
|
||||
|
||||
# -- Cooldown Defaults ----------------------------------------------------
|
||||
rate_limit_cooldown_seconds: int = 300 # 429
|
||||
overload_cooldown_seconds: int = 30 # 529
|
||||
|
||||
# -- OAuth Proactive Refresh ----------------------------------------------
|
||||
proactive_refresh_seconds: int = 180 # 3 minutes before expiry
|
||||
|
||||
# -- Health Policy --------------------------------------------------------
|
||||
health_policy_enabled: bool = True
|
||||
|
||||
# -- Temporary Unschedulable Rules ----------------------------------------
|
||||
unschedulable_rules: list[UnschedulableRule] = field(default_factory=list)
|
||||
|
||||
# -- Pluggable Strategies -------------------------------------------------
|
||||
strategies: tuple[str, ...] = ()
|
||||
|
||||
|
||||
def parse_pool_config(provider_config: Any) -> PoolConfig | None:
|
||||
"""Parse PoolConfig from ``Provider.config``.
|
||||
|
||||
Only looks for the explicit ``pool_advanced`` key. Returns ``None``
|
||||
when the provider has no pool section configured, meaning the caller
|
||||
should use the normal (non-pool) scheduling path.
|
||||
"""
|
||||
config_dict = provider_config if isinstance(provider_config, dict) else {}
|
||||
|
||||
raw_advanced = config_dict.get("pool_advanced")
|
||||
if raw_advanced is None:
|
||||
return None
|
||||
|
||||
if not isinstance(raw_advanced, dict):
|
||||
# Could be a pre-validated Pydantic model; grab its dict.
|
||||
try:
|
||||
raw_advanced = raw_advanced.model_dump() # type: ignore[union-attr]
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"PoolConfig: advanced config type invalid ({}), falling back to defaults",
|
||||
type(raw_advanced).__name__,
|
||||
)
|
||||
return PoolConfig()
|
||||
|
||||
rules: list[UnschedulableRule] = []
|
||||
raw_rules = raw_advanced.get("unschedulable_rules")
|
||||
if isinstance(raw_rules, list):
|
||||
for r in raw_rules:
|
||||
if isinstance(r, dict) and isinstance(r.get("keyword"), str):
|
||||
rules.append(
|
||||
UnschedulableRule(
|
||||
keyword=r["keyword"],
|
||||
duration_minutes=int(r.get("duration_minutes", 5)),
|
||||
)
|
||||
)
|
||||
|
||||
def _int_or(key: str, default: int) -> int:
|
||||
v = raw_advanced.get(key)
|
||||
if v is None:
|
||||
return default
|
||||
try:
|
||||
return int(v)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
def _bool_or(key: str, default: bool) -> bool:
|
||||
v = raw_advanced.get(key)
|
||||
if v is None:
|
||||
return default
|
||||
return bool(v)
|
||||
|
||||
def _opt_int(key: str) -> int | None:
|
||||
v = raw_advanced.get(key)
|
||||
if v is None:
|
||||
return None
|
||||
try:
|
||||
return int(v)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
return PoolConfig(
|
||||
sticky_session_ttl_seconds=_int_or("sticky_session_ttl_seconds", 3600),
|
||||
load_threshold_percent=_int_or("load_threshold_percent", 80),
|
||||
lru_enabled=_bool_or("lru_enabled", True),
|
||||
cost_window_seconds=_int_or("cost_window_seconds", 18000),
|
||||
cost_limit_per_key_tokens=_opt_int("cost_limit_per_key_tokens"),
|
||||
cost_soft_threshold_percent=_int_or("cost_soft_threshold_percent", 80),
|
||||
rate_limit_cooldown_seconds=_int_or("rate_limit_cooldown_seconds", 300),
|
||||
overload_cooldown_seconds=_int_or("overload_cooldown_seconds", 30),
|
||||
proactive_refresh_seconds=_int_or("proactive_refresh_seconds", 180),
|
||||
health_policy_enabled=_bool_or("health_policy_enabled", True),
|
||||
unschedulable_rules=rules,
|
||||
strategies=_parse_strategies(raw_advanced.get("strategies")),
|
||||
)
|
||||
|
||||
|
||||
def _parse_strategies(raw: Any) -> tuple[str, ...]:
|
||||
"""Parse strategy names from config (list[str] -> tuple[str, ...])."""
|
||||
if not isinstance(raw, list):
|
||||
return ()
|
||||
return tuple(str(s) for s in raw if isinstance(s, str) and s)
|
||||
66
src/services/provider/pool/cost_tracker.py
Normal file
66
src/services/provider/pool/cost_tracker.py
Normal file
@@ -0,0 +1,66 @@
|
||||
"""Rolling-window cost tracking for the Account Pool.
|
||||
|
||||
Each key has a configurable token budget per rolling window (e.g. 5 hours).
|
||||
When the budget is exhausted the key is marked as unschedulable by the pool
|
||||
manager. A "soft threshold" (default 80 %) causes the pool to *prefer*
|
||||
other keys but still allows traffic if no alternatives exist.
|
||||
|
||||
All state is stored in Redis sorted sets via :mod:`redis_ops`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from src.services.provider.pool import redis_ops
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from src.services.provider.pool.config import PoolConfig
|
||||
|
||||
|
||||
async def record_usage(
|
||||
provider_id: str,
|
||||
key_id: str,
|
||||
tokens: int,
|
||||
config: PoolConfig,
|
||||
) -> None:
|
||||
"""Record *tokens* used by *key_id* in the rolling cost window."""
|
||||
if tokens <= 0:
|
||||
return
|
||||
if config.cost_limit_per_key_tokens is None:
|
||||
return # cost tracking disabled
|
||||
await redis_ops.add_cost_entry(provider_id, key_id, tokens, config.cost_window_seconds)
|
||||
|
||||
|
||||
async def get_window_usage(
|
||||
provider_id: str,
|
||||
key_id: str,
|
||||
config: PoolConfig,
|
||||
) -> int:
|
||||
"""Return total tokens used by *key_id* within the current window."""
|
||||
return await redis_ops.get_cost_window_total(provider_id, key_id, config.cost_window_seconds)
|
||||
|
||||
|
||||
async def is_at_limit(
|
||||
provider_id: str,
|
||||
key_id: str,
|
||||
config: PoolConfig,
|
||||
) -> bool:
|
||||
"""Return ``True`` if the key has exhausted its budget."""
|
||||
if config.cost_limit_per_key_tokens is None:
|
||||
return False
|
||||
total = await get_window_usage(provider_id, key_id, config)
|
||||
return total >= config.cost_limit_per_key_tokens
|
||||
|
||||
|
||||
async def is_approaching_limit(
|
||||
provider_id: str,
|
||||
key_id: str,
|
||||
config: PoolConfig,
|
||||
) -> bool:
|
||||
"""Return ``True`` if the key is above the soft threshold."""
|
||||
if config.cost_limit_per_key_tokens is None:
|
||||
return False
|
||||
total = await get_window_usage(provider_id, key_id, config)
|
||||
threshold = config.cost_limit_per_key_tokens * config.cost_soft_threshold_percent / 100
|
||||
return total >= threshold
|
||||
204
src/services/provider/pool/health_policy.py
Normal file
204
src/services/provider/pool/health_policy.py
Normal file
@@ -0,0 +1,204 @@
|
||||
"""Account Pool health policy: error code classification and key state management.
|
||||
|
||||
Maps upstream HTTP status codes to pool-level actions:
|
||||
|
||||
| Code | Action |
|
||||
|------|--------------------------------------------------------------|
|
||||
| 401 | Invalidate OAuth token cache -> attempt refresh -> disable |
|
||||
| 402 | Auto-disable key (payment issue) |
|
||||
| 403 | Auto-disable key (suspended/banned) |
|
||||
| 400 | Check body for "organization has been disabled" -> disable |
|
||||
| 429 | Set cooldown (retry-after header or config default) |
|
||||
| 529 | Set cooldown (config default) |
|
||||
| * | Check unschedulable_rules keyword matching |
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.services.provider.pool import redis_ops
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from src.services.provider.pool.config import PoolConfig
|
||||
|
||||
# Patterns in 400 error body that indicate account-level issues.
|
||||
_ACCOUNT_DISABLE_PATTERNS = (
|
||||
"organization has been disabled",
|
||||
"organization_disabled",
|
||||
"account has been disabled",
|
||||
"account_disabled",
|
||||
)
|
||||
|
||||
|
||||
def _parse_retry_after(headers: dict[str, str] | None) -> int | None:
|
||||
"""Extract retry-after seconds from response headers."""
|
||||
if not headers:
|
||||
return None
|
||||
raw = headers.get("retry-after") or headers.get("Retry-After")
|
||||
if not raw:
|
||||
return None
|
||||
try:
|
||||
val = int(raw)
|
||||
return max(1, min(val, 3600))
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
|
||||
|
||||
def _extract_error_message(error_body: str | None) -> str:
|
||||
"""Best-effort extraction of error message from JSON body."""
|
||||
if not error_body:
|
||||
return ""
|
||||
try:
|
||||
data = json.loads(error_body)
|
||||
if isinstance(data, dict):
|
||||
error_obj = data.get("error")
|
||||
if isinstance(error_obj, dict):
|
||||
return str(error_obj.get("message", ""))
|
||||
if isinstance(error_obj, str):
|
||||
return error_obj
|
||||
return str(data.get("message", ""))
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
return error_body[:500]
|
||||
|
||||
|
||||
async def apply_health_policy(
|
||||
*,
|
||||
provider_id: str,
|
||||
key_id: str,
|
||||
status_code: int,
|
||||
error_body: str | None,
|
||||
response_headers: dict[str, str] | None,
|
||||
config: PoolConfig,
|
||||
) -> None:
|
||||
"""Apply health policy for an upstream error.
|
||||
|
||||
This is fire-and-forget; exceptions are caught and logged.
|
||||
"""
|
||||
if not config.health_policy_enabled:
|
||||
return
|
||||
|
||||
try:
|
||||
await _apply(
|
||||
provider_id=provider_id,
|
||||
key_id=key_id,
|
||||
status_code=status_code,
|
||||
error_body=error_body,
|
||||
response_headers=response_headers,
|
||||
config=config,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Pool health policy failed for key {}: {}",
|
||||
key_id[:8],
|
||||
str(exc),
|
||||
)
|
||||
|
||||
|
||||
async def _apply(
|
||||
*,
|
||||
provider_id: str,
|
||||
key_id: str,
|
||||
status_code: int,
|
||||
error_body: str | None,
|
||||
response_headers: dict[str, str] | None,
|
||||
config: PoolConfig,
|
||||
) -> None:
|
||||
error_msg = _extract_error_message(error_body)
|
||||
|
||||
# --- 401 Unauthorized ---------------------------------------------------
|
||||
if status_code == 401:
|
||||
await redis_ops.invalidate_oauth_token_cache(key_id)
|
||||
# Set a short cooldown to avoid hammering while refresh happens.
|
||||
await redis_ops.set_cooldown(provider_id, key_id, "auth_failed_401", ttl=60)
|
||||
logger.info(
|
||||
"Pool[{}]: key {} got 401, token cache invalidated + 60s cooldown",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
)
|
||||
return
|
||||
|
||||
# --- 402 Payment Required ------------------------------------------------
|
||||
if status_code == 402:
|
||||
await redis_ops.set_cooldown(provider_id, key_id, "payment_required_402", ttl=3600)
|
||||
logger.warning(
|
||||
"Pool[{}]: key {} got 402 (payment required), cooldown 1h",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
)
|
||||
return
|
||||
|
||||
# --- 403 Forbidden -------------------------------------------------------
|
||||
if status_code == 403:
|
||||
await redis_ops.set_cooldown(provider_id, key_id, "forbidden_403", ttl=3600)
|
||||
logger.warning(
|
||||
"Pool[{}]: key {} got 403 (forbidden/suspended), cooldown 1h",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
)
|
||||
return
|
||||
|
||||
# --- 400 with account-disable pattern ------------------------------------
|
||||
if status_code == 400:
|
||||
error_lower = error_msg.lower()
|
||||
for pattern in _ACCOUNT_DISABLE_PATTERNS:
|
||||
if pattern in error_lower:
|
||||
await redis_ops.set_cooldown(
|
||||
provider_id, key_id, f"account_disabled_400:{pattern}", ttl=3600
|
||||
)
|
||||
logger.warning(
|
||||
"Pool[{}]: key {} got 400 with '{}', cooldown 1h",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
pattern,
|
||||
)
|
||||
return
|
||||
|
||||
# --- 429 Rate Limited ----------------------------------------------------
|
||||
if status_code == 429:
|
||||
retry_after = _parse_retry_after(response_headers)
|
||||
ttl = retry_after or config.rate_limit_cooldown_seconds
|
||||
await redis_ops.set_cooldown(provider_id, key_id, "rate_limited_429", ttl=ttl)
|
||||
logger.info(
|
||||
"Pool[{}]: key {} got 429, cooldown {}s",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
ttl,
|
||||
)
|
||||
return
|
||||
|
||||
# --- 529 Overloaded ------------------------------------------------------
|
||||
if status_code == 529:
|
||||
ttl = config.overload_cooldown_seconds
|
||||
await redis_ops.set_cooldown(provider_id, key_id, "overloaded_529", ttl=ttl)
|
||||
logger.info(
|
||||
"Pool[{}]: key {} got 529, cooldown {}s",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
ttl,
|
||||
)
|
||||
return
|
||||
|
||||
# --- Keyword-based unschedulable rules -----------------------------------
|
||||
if config.unschedulable_rules and error_msg:
|
||||
error_lower = error_msg.lower()
|
||||
for rule in config.unschedulable_rules:
|
||||
if rule.keyword.lower() in error_lower:
|
||||
ttl = max(60, rule.duration_minutes * 60)
|
||||
await redis_ops.set_cooldown(
|
||||
provider_id,
|
||||
key_id,
|
||||
f"rule:{rule.keyword}",
|
||||
ttl=ttl,
|
||||
)
|
||||
logger.info(
|
||||
"Pool[{}]: key {} matched rule '{}', cooldown {}m",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
rule.keyword,
|
||||
rule.duration_minutes,
|
||||
)
|
||||
return
|
||||
91
src/services/provider/pool/hooks.py
Normal file
91
src/services/provider/pool/hooks.py
Normal file
@@ -0,0 +1,91 @@
|
||||
"""Pool scheduling hooks -- provider-type-specific pool behaviour.
|
||||
|
||||
Some provider types need custom logic during pool scheduling (e.g. extracting
|
||||
a session UUID for sticky binding). This module provides a small Protocol +
|
||||
registry so the pool layer stays generic while provider-specific behaviour
|
||||
lives alongside each adapter.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from typing import Any, Protocol, runtime_checkable
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Protocol
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class PoolSchedulingHook(Protocol):
|
||||
"""Provider-type-specific pool scheduling behaviour.
|
||||
|
||||
Each provider type can optionally register a hook to customize:
|
||||
- Session UUID extraction (for sticky sessions)
|
||||
- Post-success / post-error callbacks
|
||||
|
||||
Optional methods (checked via ``hasattr`` by callers):
|
||||
- ``on_pool_success``
|
||||
- ``on_pool_error``
|
||||
"""
|
||||
|
||||
name: str
|
||||
|
||||
def extract_session_uuid(self, request_body: dict[str, Any]) -> str | None:
|
||||
"""Extract a session UUID for sticky binding from the request body."""
|
||||
...
|
||||
|
||||
# -- Optional lifecycle callbacks -----------------------------------------
|
||||
# These are checked via ``hasattr`` so existing implementations that
|
||||
# don't define them will continue to work.
|
||||
|
||||
def on_pool_success(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
session_uuid: str | None,
|
||||
context: dict[str, Any],
|
||||
) -> None:
|
||||
"""Called after a successful pool request (provider-specific logic)."""
|
||||
...
|
||||
|
||||
def on_pool_error(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
status_code: int,
|
||||
context: dict[str, Any],
|
||||
) -> None:
|
||||
"""Called after a failed pool request (provider-specific logic)."""
|
||||
...
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Registry
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_hook_registry: dict[str, PoolSchedulingHook] = {}
|
||||
_registry_lock = threading.Lock()
|
||||
|
||||
|
||||
def register_pool_hook(provider_type: str, hook: PoolSchedulingHook) -> None:
|
||||
"""Register a pool scheduling hook for a provider type."""
|
||||
from src.core.provider_types import normalize_provider_type
|
||||
|
||||
pt = normalize_provider_type(provider_type)
|
||||
with _registry_lock:
|
||||
_hook_registry[pt] = hook
|
||||
|
||||
|
||||
def get_pool_hook(provider_type: str | None) -> PoolSchedulingHook | None:
|
||||
"""Return the pool scheduling hook for a provider type, or ``None``."""
|
||||
if not provider_type:
|
||||
return None
|
||||
from src.services.provider.envelope import ensure_providers_bootstrapped
|
||||
|
||||
ensure_providers_bootstrapped()
|
||||
|
||||
from src.core.provider_types import normalize_provider_type
|
||||
|
||||
pt = normalize_provider_type(provider_type)
|
||||
return _hook_registry.get(pt)
|
||||
545
src/services/provider/pool/manager.py
Normal file
545
src/services/provider/pool/manager.py
Normal file
@@ -0,0 +1,545 @@
|
||||
"""Account Pool Manager (provider-agnostic).
|
||||
|
||||
Stateless facade that coordinates pool operations for any Provider with
|
||||
pool configuration enabled. All state lives in Redis via :mod:`redis_ops`.
|
||||
|
||||
Usage::
|
||||
|
||||
mgr = PoolManager(provider_id, pool_config)
|
||||
reordered = await mgr.reorder_candidates(session_uuid, candidates)
|
||||
# ... execute request ...
|
||||
await mgr.on_request_success(session_uuid=..., key_id=..., tokens_used=...)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import random
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, Any, TypeVar
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.services.provider.pool import redis_ops
|
||||
from src.services.provider.pool.config import PoolConfig
|
||||
from src.services.provider.pool.trace import PoolCandidateTrace, PoolSchedulingTrace
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from src.models.database import ProviderAPIKey
|
||||
from src.services.scheduling.schemas import ProviderCandidate
|
||||
|
||||
|
||||
class PoolManager:
|
||||
"""Coordinate pool-level scheduling for a single Provider."""
|
||||
|
||||
__slots__ = ("provider_id", "config")
|
||||
|
||||
def __init__(self, provider_id: str, config: PoolConfig) -> None:
|
||||
self.provider_id = provider_id
|
||||
self.config = config
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Core scheduling: reorder candidate list for pool-aware selection
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def reorder_candidates(
|
||||
self,
|
||||
session_uuid: str | None,
|
||||
candidates: list[ProviderCandidate],
|
||||
) -> list[ProviderCandidate]:
|
||||
"""Reorder *candidates* according to pool rules.
|
||||
|
||||
The returned list keeps the same elements but in a new order:
|
||||
|
||||
1. **Sticky session hit** -- if the session is already bound to a key
|
||||
and that key appears in *candidates* and is not in cooldown, move it
|
||||
to position 0.
|
||||
2. **Filter** out keys in cooldown or cost-exhausted state (mark
|
||||
``is_skipped``).
|
||||
3. **LRU sort** -- among remaining candidates at the same priority
|
||||
level, sort by least-recently-used.
|
||||
4. **Random tiebreak** -- among candidates with identical LRU score.
|
||||
|
||||
Also builds a :class:`PoolSchedulingTrace` and attaches per-candidate
|
||||
trace data via ``_pool_extra_data`` / ``_pool_scheduling_trace``
|
||||
attributes on candidate objects.
|
||||
"""
|
||||
if not candidates:
|
||||
return candidates
|
||||
|
||||
pid = self.provider_id
|
||||
|
||||
# Build trace
|
||||
trace = PoolSchedulingTrace(
|
||||
provider_id=pid,
|
||||
total_keys=len(candidates),
|
||||
session_uuid=session_uuid[:8] if session_uuid else None,
|
||||
)
|
||||
|
||||
# --- Strategy: before_select ----------------------------------
|
||||
strategies = _get_active_strategies(self.config)
|
||||
key_ids = [str(c.key.id) for c in candidates]
|
||||
strategy_context: dict[str, Any] = {"session_uuid": session_uuid}
|
||||
for strategy in strategies:
|
||||
if hasattr(strategy, "on_before_select"):
|
||||
try:
|
||||
filtered = strategy.on_before_select(
|
||||
provider_id=pid,
|
||||
key_ids=key_ids,
|
||||
config=self.config,
|
||||
context=strategy_context,
|
||||
)
|
||||
if filtered is not None:
|
||||
key_ids = filtered
|
||||
except Exception:
|
||||
logger.opt(exception=True).debug(
|
||||
"Pool[{}]: strategy before_select failed", pid[:8]
|
||||
)
|
||||
|
||||
# --- 1. Sticky session ----------------------------------------
|
||||
sticky_key_id: str | None = None
|
||||
if session_uuid and self.config.sticky_session_ttl_seconds > 0:
|
||||
sticky_key_id = await redis_ops.get_sticky_binding(
|
||||
pid, session_uuid, self.config.sticky_session_ttl_seconds
|
||||
)
|
||||
|
||||
# --- 2. Batch fetch pool state (parallel) ---------------------
|
||||
all_key_ids = [str(c.key.id) for c in candidates]
|
||||
|
||||
# Fire independent Redis queries concurrently.
|
||||
_cooldown_coro = redis_ops.batch_get_cooldowns(pid, all_key_ids, include_ttl=True)
|
||||
_cost_coro = (
|
||||
redis_ops.batch_get_cost_totals(pid, all_key_ids, self.config.cost_window_seconds)
|
||||
if self.config.cost_limit_per_key_tokens is not None
|
||||
else None
|
||||
)
|
||||
_lru_coro = redis_ops.get_lru_scores(pid, all_key_ids) if self.config.lru_enabled else None
|
||||
|
||||
# Gather all non-None coroutines in parallel.
|
||||
coros: list[Any] = [_cooldown_coro]
|
||||
_cost_idx = -1
|
||||
_lru_idx = -1
|
||||
if _cost_coro is not None:
|
||||
_cost_idx = len(coros)
|
||||
coros.append(_cost_coro)
|
||||
if _lru_coro is not None:
|
||||
_lru_idx = len(coros)
|
||||
coros.append(_lru_coro)
|
||||
|
||||
gathered = await asyncio.gather(*coros)
|
||||
|
||||
cooldowns_raw = gathered[0]
|
||||
# cooldowns_raw: dict[str, tuple[str | None, int | None]]
|
||||
cooldowns: dict[str, str | None] = {}
|
||||
cooldown_ttls: dict[str, int | None] = {}
|
||||
for kid, val in cooldowns_raw.items():
|
||||
if isinstance(val, tuple):
|
||||
cooldowns[kid] = val[0]
|
||||
cooldown_ttls[kid] = val[1]
|
||||
else:
|
||||
cooldowns[kid] = val
|
||||
cooldown_ttls[kid] = None
|
||||
|
||||
# Cost check
|
||||
cost_exhausted: set[str] = set()
|
||||
cost_soft: set[str] = set()
|
||||
cost_totals: dict[str, int] = {}
|
||||
if _cost_idx >= 0:
|
||||
cost_totals = gathered[_cost_idx]
|
||||
limit = self.config.cost_limit_per_key_tokens
|
||||
assert limit is not None # guarded by _cost_idx >= 0
|
||||
for kid, total in cost_totals.items():
|
||||
if total >= limit:
|
||||
cost_exhausted.add(kid)
|
||||
elif total >= limit * self.config.cost_soft_threshold_percent / 100:
|
||||
cost_soft.add(kid)
|
||||
|
||||
# LRU scores
|
||||
lru_scores: dict[str, float] = {}
|
||||
if _lru_idx >= 0:
|
||||
lru_scores = gathered[_lru_idx]
|
||||
|
||||
# --- Strategy: compute_score ----------------------------------
|
||||
for strategy in strategies:
|
||||
if hasattr(strategy, "compute_score"):
|
||||
for kid in all_key_ids:
|
||||
try:
|
||||
custom = strategy.compute_score(
|
||||
key_id=kid,
|
||||
config=self.config,
|
||||
context=strategy_context,
|
||||
)
|
||||
if custom is not None:
|
||||
lru_scores[kid] = custom
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# --- 3. Classify candidates -----------------------------------
|
||||
sticky_candidate: ProviderCandidate | None = None
|
||||
available: list[ProviderCandidate] = []
|
||||
skipped: list[ProviderCandidate] = []
|
||||
|
||||
for c in candidates:
|
||||
kid = str(c.key.id)
|
||||
ct = PoolCandidateTrace(key_id=kid)
|
||||
|
||||
# Already skipped upstream?
|
||||
if c.is_skipped:
|
||||
skipped.append(c)
|
||||
ct.skipped = True
|
||||
ct.skip_type = "upstream"
|
||||
trace.candidate_traces[kid] = ct
|
||||
continue
|
||||
|
||||
# Cooldown?
|
||||
cd_reason = cooldowns.get(kid)
|
||||
if cd_reason is not None:
|
||||
c.is_skipped = True
|
||||
c.skip_reason = f"pool cooldown: {cd_reason}"
|
||||
skipped.append(c)
|
||||
ct.skipped = True
|
||||
ct.skip_type = "cooldown"
|
||||
ct.cooldown_reason = cd_reason
|
||||
ct.cooldown_ttl = cooldown_ttls.get(kid)
|
||||
_attach_pool_extra(c, ct)
|
||||
trace.candidate_traces[kid] = ct
|
||||
continue
|
||||
|
||||
# Cost exhausted?
|
||||
if kid in cost_exhausted:
|
||||
c.is_skipped = True
|
||||
c.skip_reason = "pool cost limit reached"
|
||||
skipped.append(c)
|
||||
ct.skipped = True
|
||||
ct.skip_type = "cost_exhausted"
|
||||
ct.cost_window_usage = cost_totals.get(kid, 0)
|
||||
ct.cost_limit = self.config.cost_limit_per_key_tokens
|
||||
_attach_pool_extra(c, ct)
|
||||
trace.candidate_traces[kid] = ct
|
||||
continue
|
||||
|
||||
# Sticky hit?
|
||||
if sticky_key_id and kid == sticky_key_id:
|
||||
sticky_candidate = c
|
||||
ct.reason = "sticky"
|
||||
ct.sticky_hit = True
|
||||
trace.sticky_session_used = True
|
||||
else:
|
||||
available.append(c)
|
||||
ct.reason = "lru" if lru_scores.get(kid, 0) > 0 else "random"
|
||||
|
||||
ct.lru_score = lru_scores.get(kid, 0.0)
|
||||
ct.cost_window_usage = cost_totals.get(kid, 0)
|
||||
ct.cost_limit = self.config.cost_limit_per_key_tokens
|
||||
if kid in cost_soft:
|
||||
ct.cost_soft_threshold = True
|
||||
_attach_pool_extra(c, ct)
|
||||
trace.candidate_traces[kid] = ct
|
||||
|
||||
# --- 4. Sort available by LRU ---------------------------------
|
||||
if lru_scores and available:
|
||||
available.sort(key=lambda c: lru_scores.get(str(c.key.id), 0.0))
|
||||
|
||||
# Random tiebreak among candidates with the same LRU score
|
||||
if len(available) > 1 and lru_scores:
|
||||
_shuffle_same_score_groups(available, lru_scores)
|
||||
|
||||
# --- 5. Assemble final order ----------------------------------
|
||||
result: list[ProviderCandidate] = []
|
||||
if sticky_candidate is not None:
|
||||
result.append(sticky_candidate)
|
||||
result.extend(available)
|
||||
result.extend(skipped)
|
||||
|
||||
if sticky_candidate:
|
||||
logger.debug(
|
||||
"Pool[{}]: sticky hit key={}",
|
||||
pid[:8],
|
||||
sticky_key_id and sticky_key_id[:8],
|
||||
)
|
||||
|
||||
# --- Strategy: after_select -----------------------------------
|
||||
if result:
|
||||
first_kid = str(result[0].key.id)
|
||||
first_trace = trace.candidate_traces.get(first_kid)
|
||||
for strategy in strategies:
|
||||
if hasattr(strategy, "on_after_select") and first_trace:
|
||||
try:
|
||||
strategy.on_after_select(
|
||||
provider_id=pid,
|
||||
selected_key_id=first_kid,
|
||||
trace=first_trace,
|
||||
config=self.config,
|
||||
context=strategy_context,
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Attach the full trace to the first candidate for downstream use.
|
||||
if result:
|
||||
setattr(result[0], "_pool_scheduling_trace", trace)
|
||||
|
||||
return result
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Single-key selection (used by CandidateBuilder for pooled providers)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def select_key(
|
||||
self,
|
||||
session_uuid: str | None,
|
||||
keys: list[ProviderAPIKey],
|
||||
) -> ProviderAPIKey | None:
|
||||
"""Select the best key from *keys* according to pool rules.
|
||||
|
||||
Same logic as :meth:`reorder_candidates` but operates directly on
|
||||
:class:`ProviderAPIKey` objects instead of candidates:
|
||||
|
||||
1. Sticky session hit (if bound and still healthy).
|
||||
2. Filter out keys in cooldown or cost-exhausted.
|
||||
3. LRU sort among remaining keys.
|
||||
4. Random tiebreak for identical LRU scores.
|
||||
5. Return the first available key, or ``None``.
|
||||
"""
|
||||
if not keys:
|
||||
return None
|
||||
|
||||
pid = self.provider_id
|
||||
|
||||
# --- 1. Sticky session ------------------------------------------------
|
||||
sticky_key_id: str | None = None
|
||||
if session_uuid and self.config.sticky_session_ttl_seconds > 0:
|
||||
sticky_key_id = await redis_ops.get_sticky_binding(
|
||||
pid, session_uuid, self.config.sticky_session_ttl_seconds
|
||||
)
|
||||
|
||||
# --- 2. Batch fetch pool state (parallel) -----------------------------
|
||||
key_ids = [str(k.id) for k in keys]
|
||||
|
||||
_cooldown_coro = redis_ops.batch_get_cooldowns(pid, key_ids)
|
||||
_cost_coro = (
|
||||
redis_ops.batch_get_cost_totals(pid, key_ids, self.config.cost_window_seconds)
|
||||
if self.config.cost_limit_per_key_tokens is not None
|
||||
else None
|
||||
)
|
||||
_lru_coro = redis_ops.get_lru_scores(pid, key_ids) if self.config.lru_enabled else None
|
||||
|
||||
coros_sk: list[Any] = [_cooldown_coro]
|
||||
_cost_idx_sk = -1
|
||||
_lru_idx_sk = -1
|
||||
if _cost_coro is not None:
|
||||
_cost_idx_sk = len(coros_sk)
|
||||
coros_sk.append(_cost_coro)
|
||||
if _lru_coro is not None:
|
||||
_lru_idx_sk = len(coros_sk)
|
||||
coros_sk.append(_lru_coro)
|
||||
|
||||
gathered_sk = await asyncio.gather(*coros_sk)
|
||||
|
||||
cooldowns = gathered_sk[0]
|
||||
|
||||
cost_exhausted: set[str] = set()
|
||||
if _cost_idx_sk >= 0:
|
||||
cost_totals = gathered_sk[_cost_idx_sk]
|
||||
for kid, total in cost_totals.items():
|
||||
if total >= self.config.cost_limit_per_key_tokens: # type: ignore[operator]
|
||||
cost_exhausted.add(kid)
|
||||
|
||||
lru_scores: dict[str, float] = {}
|
||||
if _lru_idx_sk >= 0:
|
||||
lru_scores = gathered_sk[_lru_idx_sk]
|
||||
|
||||
# --- Strategy: compute_score ------------------------------------------
|
||||
strategies = _get_active_strategies(self.config)
|
||||
strategy_context: dict[str, Any] = {"session_uuid": session_uuid}
|
||||
for strategy in strategies:
|
||||
if hasattr(strategy, "compute_score"):
|
||||
for kid in key_ids:
|
||||
try:
|
||||
custom = strategy.compute_score(
|
||||
key_id=kid,
|
||||
config=self.config,
|
||||
context=strategy_context,
|
||||
)
|
||||
if custom is not None:
|
||||
lru_scores[kid] = custom
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# --- 3. Classify keys -------------------------------------------------
|
||||
sticky_key: ProviderAPIKey | None = None
|
||||
available: list[ProviderAPIKey] = []
|
||||
|
||||
for k in keys:
|
||||
kid = str(k.id)
|
||||
|
||||
if cooldowns.get(kid) is not None:
|
||||
continue
|
||||
if kid in cost_exhausted:
|
||||
continue
|
||||
|
||||
if sticky_key_id and kid == sticky_key_id:
|
||||
sticky_key = k
|
||||
continue
|
||||
|
||||
available.append(k)
|
||||
|
||||
# --- 4. Sort by LRU ---------------------------------------------------
|
||||
if lru_scores and available:
|
||||
available.sort(key=lambda k: lru_scores.get(str(k.id), 0.0))
|
||||
|
||||
# Random tiebreak within same-score groups
|
||||
if len(available) > 1 and lru_scores:
|
||||
_shuffle_same_score_keys(available, lru_scores)
|
||||
|
||||
# --- 5. Pick the winner -----------------------------------------------
|
||||
if sticky_key is not None:
|
||||
logger.debug(
|
||||
"Pool[{}]: sticky select key={}",
|
||||
pid[:8],
|
||||
sticky_key_id and sticky_key_id[:8],
|
||||
)
|
||||
return sticky_key
|
||||
|
||||
return available[0] if available else None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Post-request hooks
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def on_request_success(
|
||||
self,
|
||||
*,
|
||||
session_uuid: str | None,
|
||||
key_id: str,
|
||||
tokens_used: int = 0,
|
||||
) -> None:
|
||||
"""Called after a successful upstream request."""
|
||||
pid = self.provider_id
|
||||
|
||||
# Bind sticky session
|
||||
if session_uuid and self.config.sticky_session_ttl_seconds > 0:
|
||||
await redis_ops.set_sticky_binding(
|
||||
pid, session_uuid, key_id, self.config.sticky_session_ttl_seconds
|
||||
)
|
||||
|
||||
# Touch LRU
|
||||
if self.config.lru_enabled:
|
||||
await redis_ops.touch_lru(pid, key_id)
|
||||
|
||||
# Record cost
|
||||
if tokens_used > 0 and self.config.cost_limit_per_key_tokens is not None:
|
||||
await redis_ops.add_cost_entry(
|
||||
pid, key_id, tokens_used, self.config.cost_window_seconds
|
||||
)
|
||||
|
||||
async def on_request_error(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
status_code: int,
|
||||
error_body: str | None = None,
|
||||
response_headers: dict[str, str] | None = None,
|
||||
) -> None:
|
||||
"""Called after an upstream error. Delegates to health policy."""
|
||||
# Import lazily to avoid circular deps
|
||||
from src.services.provider.pool.health_policy import apply_health_policy
|
||||
|
||||
await apply_health_policy(
|
||||
provider_id=self.provider_id,
|
||||
key_id=key_id,
|
||||
status_code=status_code,
|
||||
error_body=error_body,
|
||||
response_headers=response_headers,
|
||||
config=self.config,
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Key schedulability check (used by candidate_builder)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def is_key_schedulable(self, key_id: str) -> tuple[bool, str | None]:
|
||||
"""Check if *key_id* is currently schedulable (not in cooldown, not
|
||||
cost-exhausted). Returns ``(True, None)`` or ``(False, reason)``.
|
||||
"""
|
||||
pid = self.provider_id
|
||||
|
||||
# Cooldown check
|
||||
cd = await redis_ops.get_cooldown(pid, key_id)
|
||||
if cd is not None:
|
||||
return False, f"pool cooldown: {cd}"
|
||||
|
||||
# Cost check
|
||||
if self.config.cost_limit_per_key_tokens is not None:
|
||||
total = await redis_ops.get_cost_window_total(
|
||||
pid, key_id, self.config.cost_window_seconds
|
||||
)
|
||||
if total >= self.config.cost_limit_per_key_tokens:
|
||||
return False, "pool cost limit reached"
|
||||
|
||||
return True, None
|
||||
|
||||
|
||||
# Backward-compatible alias
|
||||
ClaudeCodePoolManager = PoolManager
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_T = TypeVar("_T")
|
||||
|
||||
|
||||
def _attach_pool_extra(candidate: Any, ct: PoolCandidateTrace) -> None:
|
||||
"""Attach pool trace extra_data onto a candidate object."""
|
||||
existing = getattr(candidate, "_pool_extra_data", None) or {}
|
||||
existing.update(ct.to_extra_data())
|
||||
setattr(candidate, "_pool_extra_data", existing)
|
||||
|
||||
|
||||
def _get_active_strategies(config: PoolConfig) -> list[Any]:
|
||||
"""Get active strategies for the given config (lazy import)."""
|
||||
if not config.strategies:
|
||||
return []
|
||||
try:
|
||||
from src.services.provider.pool.strategy import get_active_strategies
|
||||
|
||||
return get_active_strategies(config.strategies)
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
||||
def _shuffle_same_score(
|
||||
items: list[_T],
|
||||
lru_scores: dict[str, float],
|
||||
key_fn: Callable[[_T], str],
|
||||
) -> None:
|
||||
"""In-place random shuffle within groups that share the same LRU score."""
|
||||
if len(items) <= 1:
|
||||
return
|
||||
|
||||
i = 0
|
||||
while i < len(items):
|
||||
score_i = lru_scores.get(key_fn(items[i]), 0.0)
|
||||
j = i + 1
|
||||
while j < len(items) and lru_scores.get(key_fn(items[j]), 0.0) == score_i:
|
||||
j += 1
|
||||
if j - i > 1:
|
||||
group = items[i:j]
|
||||
random.shuffle(group)
|
||||
items[i:j] = group
|
||||
i = j
|
||||
|
||||
|
||||
def _shuffle_same_score_groups(
|
||||
candidates: list[ProviderCandidate],
|
||||
lru_scores: dict[str, float],
|
||||
) -> None:
|
||||
_shuffle_same_score(candidates, lru_scores, lambda c: str(c.key.id))
|
||||
|
||||
|
||||
def _shuffle_same_score_keys(
|
||||
keys: list[ProviderAPIKey],
|
||||
lru_scores: dict[str, float],
|
||||
) -> None:
|
||||
_shuffle_same_score(keys, lru_scores, lambda k: str(k.id))
|
||||
42
src/services/provider/pool/oauth_cache.py
Normal file
42
src/services/provider/pool/oauth_cache.py
Normal file
@@ -0,0 +1,42 @@
|
||||
"""OAuth token Redis cache for the Account Pool.
|
||||
|
||||
Additions over the base ``auth.py`` refresh flow:
|
||||
|
||||
- **Redis token cache**: Avoids repeated DB decryption for hot keys.
|
||||
Cache key: ``provider_oauth_token_cache:{key_id}``
|
||||
- **Configurable proactive refresh skew**: Default 180 s (3 min) instead
|
||||
of the base 120 s, configurable via ``PoolConfig.proactive_refresh_seconds``.
|
||||
- **401 immediate invalidation**: Clears the Redis cache so the next request
|
||||
triggers a fresh refresh.
|
||||
|
||||
This module does NOT replace ``auth.py``; it adds a caching layer that
|
||||
``auth.py`` can consult before decrypting from DB.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.services.provider.pool import redis_ops
|
||||
|
||||
|
||||
async def get_cached_token(key_id: str) -> str | None:
|
||||
"""Return cached access token from Redis, or None."""
|
||||
return await redis_ops.get_cached_oauth_token(key_id)
|
||||
|
||||
|
||||
async def cache_token(key_id: str, token: str, expires_in_seconds: int) -> None:
|
||||
"""Cache an access token in Redis.
|
||||
|
||||
*expires_in_seconds* is the remaining lifetime of the token. We shave
|
||||
off 60 s so the cache expires slightly before the token itself, giving
|
||||
the refresh flow time to act.
|
||||
"""
|
||||
ttl = max(1, expires_in_seconds - 60)
|
||||
await redis_ops.cache_oauth_token(key_id, token, ttl)
|
||||
logger.debug("Pool OAuth: cached token for key {} (TTL={}s)", key_id[:8], ttl)
|
||||
|
||||
|
||||
async def invalidate_token(key_id: str) -> None:
|
||||
"""Invalidate the cached token (e.g. after a 401)."""
|
||||
await redis_ops.invalidate_oauth_token_cache(key_id)
|
||||
logger.debug("Pool OAuth: invalidated token cache for key {}", key_id[:8])
|
||||
470
src/services/provider/pool/redis_ops.py
Normal file
470
src/services/provider/pool/redis_ops.py
Normal file
@@ -0,0 +1,470 @@
|
||||
"""Redis operations for the Account Pool (provider-agnostic).
|
||||
|
||||
All pool transient state is stored in Redis. This module centralises key
|
||||
naming, Lua scripts, and graceful fallbacks so that the rest of the pool
|
||||
layer is free of Redis specifics.
|
||||
|
||||
Key schema
|
||||
----------
|
||||
ap:{pid}:sticky:{session_uuid} STRING -> key_id (TTL: config)
|
||||
ap:{pid}:lru ZSET member=key_id, score=unix_ts
|
||||
ap:{pid}:cooldown:{key_id} STRING -> reason (TTL: error-specific)
|
||||
ap:{pid}:cost:{key_id} ZSET member=req_id, score=unix_ts
|
||||
provider_oauth_token_cache:{key_id} STRING -> access_token (TTL: expires - 60)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
import uuid
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from src.clients.redis_client import get_redis_client
|
||||
from src.core.logger import logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import redis.asyncio as aioredis
|
||||
|
||||
PREFIX = "ap"
|
||||
|
||||
|
||||
def _sticky_key(provider_id: str, session_uuid: str) -> str:
|
||||
return f"{PREFIX}:{provider_id}:sticky:{session_uuid}"
|
||||
|
||||
|
||||
def _lru_key(provider_id: str) -> str:
|
||||
return f"{PREFIX}:{provider_id}:lru"
|
||||
|
||||
|
||||
def _cooldown_key(provider_id: str, key_id: str) -> str:
|
||||
return f"{PREFIX}:{provider_id}:cooldown:{key_id}"
|
||||
|
||||
|
||||
def _cost_key(provider_id: str, key_id: str) -> str:
|
||||
return f"{PREFIX}:{provider_id}:cost:{key_id}"
|
||||
|
||||
|
||||
def _oauth_cache_key(key_id: str) -> str:
|
||||
return f"provider_oauth_token_cache:{key_id}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Lua scripts
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Sticky-select: GET binding, verify it's not in cooldown, refresh TTL.
|
||||
# KEYS[1] = sticky key, KEYS[2] = cooldown key prefix (ap:{pid}:cooldown:)
|
||||
# ARGV[1] = ttl
|
||||
# Returns: key_id or nil
|
||||
_STICKY_SELECT_LUA = """
|
||||
local binding = redis.call("GET", KEYS[1])
|
||||
if not binding then
|
||||
return nil
|
||||
end
|
||||
-- Check cooldown for the bound key
|
||||
local cooldown_key = KEYS[2] .. binding
|
||||
local in_cooldown = redis.call("EXISTS", cooldown_key)
|
||||
if in_cooldown == 1 then
|
||||
redis.call("DEL", KEYS[1])
|
||||
return nil
|
||||
end
|
||||
redis.call("EXPIRE", KEYS[1], tonumber(ARGV[1]))
|
||||
return binding
|
||||
"""
|
||||
|
||||
# Cost window cleanup + sum tokens in a single round-trip.
|
||||
# KEYS[1] = cost zset key, ARGV[1] = window_start timestamp
|
||||
# Returns total token count within the window.
|
||||
_COST_WINDOW_SUM_LUA = """
|
||||
local key = KEYS[1]
|
||||
local window_start = tonumber(ARGV[1])
|
||||
redis.call("ZREMRANGEBYSCORE", key, "-inf", window_start)
|
||||
local members = redis.call("ZRANGEBYSCORE", key, window_start, "+inf")
|
||||
local total = 0
|
||||
for _, m in ipairs(members) do
|
||||
local colon = string.find(m, ":", 1, true)
|
||||
if colon then
|
||||
local n = tonumber(string.sub(m, colon + 1))
|
||||
if n then total = total + n end
|
||||
end
|
||||
end
|
||||
return total
|
||||
"""
|
||||
|
||||
|
||||
async def _get_redis() -> "aioredis.Redis | None":
|
||||
return await get_redis_client(require_redis=False)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Sticky session
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def get_sticky_binding(provider_id: str, session_uuid: str, ttl: int) -> str | None:
|
||||
"""Get and refresh sticky session binding. Returns key_id or None."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return None
|
||||
try:
|
||||
result = await redis.eval(
|
||||
_STICKY_SELECT_LUA,
|
||||
2,
|
||||
_sticky_key(provider_id, session_uuid),
|
||||
f"{PREFIX}:{provider_id}:cooldown:",
|
||||
str(ttl),
|
||||
)
|
||||
if result:
|
||||
return result.decode() if isinstance(result, bytes) else str(result)
|
||||
return None
|
||||
except Exception:
|
||||
logger.debug("Pool: sticky GET failed for session {}", session_uuid[:8])
|
||||
return None
|
||||
|
||||
|
||||
async def set_sticky_binding(provider_id: str, session_uuid: str, key_id: str, ttl: int) -> None:
|
||||
"""Create or update sticky session binding."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
await redis.setex(_sticky_key(provider_id, session_uuid), ttl, key_id)
|
||||
except Exception:
|
||||
logger.debug("Pool: sticky SET failed for session {}", session_uuid[:8])
|
||||
|
||||
|
||||
async def delete_sticky_binding(provider_id: str, session_uuid: str) -> None:
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
await redis.delete(_sticky_key(provider_id, session_uuid))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# LRU
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def get_lru_scores(provider_id: str, key_ids: list[str]) -> dict[str, float]:
|
||||
"""Batch-fetch LRU timestamps. Missing keys get score 0 (highest priority)."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return {}
|
||||
try:
|
||||
lru_k = _lru_key(provider_id)
|
||||
scores = await redis.zmscore(lru_k, key_ids)
|
||||
result: dict[str, float] = {}
|
||||
for kid, score in zip(key_ids, scores):
|
||||
result[kid] = float(score) if score is not None else 0.0
|
||||
return result
|
||||
except Exception:
|
||||
logger.debug("Pool: LRU ZMSCORE failed for provider {}", provider_id[:8])
|
||||
return {}
|
||||
|
||||
|
||||
async def touch_lru(provider_id: str, key_id: str) -> None:
|
||||
"""Update last-used timestamp."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
await redis.zadd(_lru_key(provider_id), {key_id: time.time()})
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cooldown
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def set_cooldown(provider_id: str, key_id: str, reason: str, ttl: int) -> None:
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
await redis.setex(_cooldown_key(provider_id, key_id), ttl, reason)
|
||||
logger.info(
|
||||
"Pool: key {} cooldown set: {} ({}s)",
|
||||
key_id[:8],
|
||||
reason,
|
||||
ttl,
|
||||
)
|
||||
except Exception:
|
||||
logger.debug("Pool: cooldown SET failed for key {}", key_id[:8])
|
||||
|
||||
|
||||
async def get_cooldown(provider_id: str, key_id: str) -> str | None:
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return None
|
||||
try:
|
||||
val = await redis.get(_cooldown_key(provider_id, key_id))
|
||||
if val:
|
||||
return val.decode() if isinstance(val, bytes) else str(val)
|
||||
return None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
async def clear_cooldown(provider_id: str, key_id: str) -> None:
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
await redis.delete(_cooldown_key(provider_id, key_id))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
async def batch_get_cooldowns(
|
||||
provider_id: str,
|
||||
key_ids: list[str],
|
||||
*,
|
||||
include_ttl: bool = False,
|
||||
) -> dict[str, str | None] | dict[str, tuple[str | None, int | None]]:
|
||||
"""Batch check cooldown status for multiple keys.
|
||||
|
||||
When *include_ttl* is ``True``, each value is a ``(reason, ttl_seconds)``
|
||||
tuple instead of a plain reason string. The TTL commands are batched in
|
||||
the same pipeline so there is no extra round-trip.
|
||||
"""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
if include_ttl:
|
||||
return {k: (None, None) for k in key_ids}
|
||||
return {k: None for k in key_ids}
|
||||
try:
|
||||
pipe = redis.pipeline()
|
||||
for kid in key_ids:
|
||||
ck = _cooldown_key(provider_id, kid)
|
||||
pipe.get(ck)
|
||||
if include_ttl:
|
||||
pipe.ttl(ck)
|
||||
results = await pipe.execute()
|
||||
|
||||
if include_ttl:
|
||||
out_ttl: dict[str, tuple[str | None, int | None]] = {}
|
||||
# results interleave GET/TTL: [val0, ttl0, val1, ttl1, ...]
|
||||
for i, kid in enumerate(key_ids):
|
||||
val = results[i * 2]
|
||||
ttl_val = results[i * 2 + 1]
|
||||
reason: str | None = None
|
||||
if val:
|
||||
reason = val.decode() if isinstance(val, bytes) else str(val)
|
||||
ttl_sec: int | None = None
|
||||
if isinstance(ttl_val, int) and ttl_val > 0:
|
||||
ttl_sec = ttl_val
|
||||
out_ttl[kid] = (reason, ttl_sec)
|
||||
return out_ttl
|
||||
|
||||
out: dict[str, str | None] = {}
|
||||
for kid, val in zip(key_ids, results):
|
||||
if val:
|
||||
out[kid] = val.decode() if isinstance(val, bytes) else str(val)
|
||||
else:
|
||||
out[kid] = None
|
||||
return out
|
||||
except Exception:
|
||||
if include_ttl:
|
||||
return {k: (None, None) for k in key_ids}
|
||||
return {k: None for k in key_ids}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cost tracking
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def add_cost_entry(provider_id: str, key_id: str, tokens: int, window_seconds: int) -> None:
|
||||
"""Record a cost entry (tokens used) with automatic window expiry."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
now = time.time()
|
||||
cost_k = _cost_key(provider_id, key_id)
|
||||
member = f"{uuid.uuid4().hex}:{tokens}"
|
||||
pipe = redis.pipeline()
|
||||
pipe.zadd(cost_k, {member: now})
|
||||
# Set a TTL slightly larger than the window to auto-clean abandoned keys.
|
||||
pipe.expire(cost_k, window_seconds + 600)
|
||||
await pipe.execute()
|
||||
except Exception:
|
||||
logger.debug("Pool: cost ADD failed for key {}", key_id[:8])
|
||||
|
||||
|
||||
async def get_cost_window_total(provider_id: str, key_id: str, window_seconds: int) -> int:
|
||||
"""Sum tokens used within the rolling window (single key)."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return 0
|
||||
try:
|
||||
now = time.time()
|
||||
window_start = now - window_seconds
|
||||
cost_k = _cost_key(provider_id, key_id)
|
||||
result = await redis.eval(_COST_WINDOW_SUM_LUA, 1, cost_k, str(window_start))
|
||||
return int(result) if result else 0
|
||||
except Exception:
|
||||
logger.debug("Pool: cost SUM failed for key {}", key_id[:8])
|
||||
return 0
|
||||
|
||||
|
||||
async def batch_get_cost_totals(
|
||||
provider_id: str, key_ids: list[str], window_seconds: int
|
||||
) -> dict[str, int]:
|
||||
"""Batch-fetch cost totals for multiple keys using pipeline + Lua."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return {k: 0 for k in key_ids}
|
||||
try:
|
||||
now = time.time()
|
||||
window_start = now - window_seconds
|
||||
pipe = redis.pipeline()
|
||||
for kid in key_ids:
|
||||
cost_k = _cost_key(provider_id, kid)
|
||||
pipe.eval(_COST_WINDOW_SUM_LUA, 1, cost_k, str(window_start))
|
||||
results = await pipe.execute()
|
||||
out: dict[str, int] = {}
|
||||
for kid, val in zip(key_ids, results):
|
||||
out[kid] = int(val) if val else 0
|
||||
return out
|
||||
except Exception:
|
||||
logger.debug("Pool: batch cost SUM failed for provider {}", provider_id[:8])
|
||||
return {k: 0 for k in key_ids}
|
||||
|
||||
|
||||
async def clear_cost(provider_id: str, key_id: str) -> None:
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
await redis.delete(_cost_key(provider_id, key_id))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# OAuth token cache
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def cache_oauth_token(key_id: str, token: str, ttl: int) -> None:
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
if ttl > 0:
|
||||
await redis.setex(_oauth_cache_key(key_id), ttl, token)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
async def get_cached_oauth_token(key_id: str) -> str | None:
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return None
|
||||
try:
|
||||
val = await redis.get(_oauth_cache_key(key_id))
|
||||
if val:
|
||||
return val.decode() if isinstance(val, bytes) else str(val)
|
||||
return None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
async def invalidate_oauth_token_cache(key_id: str) -> None:
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
await redis.delete(_oauth_cache_key(key_id))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pool status query (admin)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def get_sticky_session_count(provider_id: str) -> int:
|
||||
"""Approximate count of active sticky sessions (via SCAN, for admin display only)."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return 0
|
||||
try:
|
||||
pattern = f"{PREFIX}:{provider_id}:sticky:*"
|
||||
count = 0
|
||||
async for _ in redis.scan_iter(match=pattern, count=100):
|
||||
count += 1
|
||||
return count
|
||||
except Exception:
|
||||
return 0
|
||||
|
||||
|
||||
async def get_key_sticky_count(provider_id: str, key_id: str) -> int:
|
||||
"""Count sticky sessions bound to a specific key (admin only).
|
||||
|
||||
Uses batched SCAN + pipeline MGET to reduce Redis round-trips.
|
||||
"""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return 0
|
||||
try:
|
||||
pattern = f"{PREFIX}:{provider_id}:sticky:*"
|
||||
count = 0
|
||||
batch: list[bytes | str] = []
|
||||
async for k in redis.scan_iter(match=pattern, count=200):
|
||||
batch.append(k)
|
||||
if len(batch) >= 200:
|
||||
vals = await redis.mget(batch)
|
||||
for val in vals:
|
||||
if val:
|
||||
bound_id = val.decode() if isinstance(val, bytes) else str(val)
|
||||
if bound_id == key_id:
|
||||
count += 1
|
||||
batch.clear()
|
||||
if batch:
|
||||
vals = await redis.mget(batch)
|
||||
for val in vals:
|
||||
if val:
|
||||
bound_id = val.decode() if isinstance(val, bytes) else str(val)
|
||||
if bound_id == key_id:
|
||||
count += 1
|
||||
return count
|
||||
except Exception:
|
||||
return 0
|
||||
|
||||
|
||||
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()
|
||||
if redis is None:
|
||||
return None
|
||||
try:
|
||||
ttl = await redis.ttl(_cooldown_key(provider_id, key_id))
|
||||
return ttl if ttl > 0 else None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
async def batch_get_cooldown_ttls(provider_id: str, key_ids: list[str]) -> dict[str, int | None]:
|
||||
"""Batch-fetch cooldown TTLs for multiple keys using pipeline."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return {k: None for k in key_ids}
|
||||
try:
|
||||
pipe = redis.pipeline()
|
||||
for kid in key_ids:
|
||||
pipe.ttl(_cooldown_key(provider_id, kid))
|
||||
results = await pipe.execute()
|
||||
out: dict[str, int | None] = {}
|
||||
for kid, ttl in zip(key_ids, results):
|
||||
out[kid] = int(ttl) if isinstance(ttl, int) and ttl > 0 else None
|
||||
return out
|
||||
except Exception:
|
||||
return {k: None for k in key_ids}
|
||||
107
src/services/provider/pool/strategy.py
Normal file
107
src/services/provider/pool/strategy.py
Normal file
@@ -0,0 +1,107 @@
|
||||
"""Pluggable pool scheduling strategies.
|
||||
|
||||
Strategies allow customising pool-level candidate selection without
|
||||
modifying the core :class:`PoolManager`. Each strategy is an object
|
||||
that implements one or more optional methods defined by the
|
||||
:class:`PoolSchedulingStrategy` protocol.
|
||||
|
||||
Registration uses a thread-safe global registry (same pattern as
|
||||
:mod:`~src.services.provider.pool.hooks`).
|
||||
|
||||
Usage::
|
||||
|
||||
from src.services.provider.pool.strategy import register_pool_strategy
|
||||
|
||||
class MyStrategy:
|
||||
name = "usage_weight"
|
||||
|
||||
def compute_score(self, *, key_id, config, context):
|
||||
...
|
||||
|
||||
register_pool_strategy("usage_weight", MyStrategy())
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from src.services.provider.pool.config import PoolConfig
|
||||
from src.services.provider.pool.trace import PoolCandidateTrace
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class PoolSchedulingStrategy(Protocol):
|
||||
"""Pluggable pool scheduling strategy.
|
||||
|
||||
All methods are optional -- callers check via ``hasattr``.
|
||||
Strategies are activated per-provider through ``PoolConfig.strategies``.
|
||||
"""
|
||||
|
||||
name: str
|
||||
|
||||
def on_before_select(
|
||||
self,
|
||||
*,
|
||||
provider_id: str,
|
||||
key_ids: list[str],
|
||||
config: PoolConfig,
|
||||
context: dict[str, Any],
|
||||
) -> list[str] | None:
|
||||
"""Filter / reorder *key_ids* before selection.
|
||||
|
||||
Return ``None`` to leave the list unchanged.
|
||||
"""
|
||||
...
|
||||
|
||||
def on_after_select(
|
||||
self,
|
||||
*,
|
||||
provider_id: str,
|
||||
selected_key_id: str,
|
||||
trace: PoolCandidateTrace,
|
||||
config: PoolConfig,
|
||||
context: dict[str, Any],
|
||||
) -> None:
|
||||
"""Called after a key has been selected (for logging / metrics)."""
|
||||
...
|
||||
|
||||
def compute_score(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
config: PoolConfig,
|
||||
context: dict[str, Any],
|
||||
) -> float | None:
|
||||
"""Return a custom sort score. ``None`` means "do not override"."""
|
||||
...
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Registry
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_strategy_registry: dict[str, PoolSchedulingStrategy] = {}
|
||||
_strategy_lock = threading.Lock()
|
||||
|
||||
|
||||
def register_pool_strategy(name: str, strategy: PoolSchedulingStrategy) -> None:
|
||||
"""Register a pool scheduling strategy globally."""
|
||||
with _strategy_lock:
|
||||
_strategy_registry[name] = strategy
|
||||
|
||||
|
||||
def get_pool_strategy(name: str) -> PoolSchedulingStrategy | None:
|
||||
"""Return a registered strategy by *name*, or ``None``."""
|
||||
return _strategy_registry.get(name)
|
||||
|
||||
|
||||
def get_active_strategies(names: tuple[str, ...] | list[str]) -> list[PoolSchedulingStrategy]:
|
||||
"""Return registered strategies whose names appear in *names*."""
|
||||
result: list[PoolSchedulingStrategy] = []
|
||||
for n in names:
|
||||
s = _strategy_registry.get(n)
|
||||
if s is not None:
|
||||
result.append(s)
|
||||
return result
|
||||
96
src/services/provider/pool/trace.py
Normal file
96
src/services/provider/pool/trace.py
Normal file
@@ -0,0 +1,96 @@
|
||||
"""Pool scheduling trace -- per-candidate decision records.
|
||||
|
||||
Collects scheduling decisions made during pool-level candidate selection
|
||||
without adding any extra Redis round-trips. Trace data is later written
|
||||
to ``RequestCandidate.extra_data`` and ``Usage.request_metadata``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class PoolCandidateTrace:
|
||||
"""Single candidate scheduling decision in a pool context."""
|
||||
|
||||
key_id: str
|
||||
reason: str = "" # sticky / lru / random / tiebreak
|
||||
sticky_hit: bool = False
|
||||
lru_score: float = 0.0
|
||||
cost_window_usage: int = 0
|
||||
cost_limit: int | None = None
|
||||
cost_soft_threshold: bool = False
|
||||
skipped: bool = False
|
||||
skip_type: str | None = None # cooldown / cost_exhausted
|
||||
cooldown_reason: str | None = None
|
||||
cooldown_ttl: int | None = None
|
||||
|
||||
def to_extra_data(self) -> dict[str, Any]:
|
||||
"""Build dict to merge into ``RequestCandidate.extra_data``."""
|
||||
if self.skipped:
|
||||
skip_info: dict[str, Any] = {"type": self.skip_type}
|
||||
if self.cooldown_reason is not None:
|
||||
skip_info["cooldown_reason"] = self.cooldown_reason
|
||||
if self.cooldown_ttl is not None:
|
||||
skip_info["cooldown_ttl"] = self.cooldown_ttl
|
||||
if self.cost_window_usage:
|
||||
skip_info["cost_window_usage"] = self.cost_window_usage
|
||||
return {"pool_skip": skip_info}
|
||||
|
||||
sel: dict[str, Any] = {"reason": self.reason}
|
||||
if self.sticky_hit:
|
||||
sel["sticky_hit"] = True
|
||||
if self.lru_score:
|
||||
sel["lru_score"] = self.lru_score
|
||||
if self.cost_window_usage:
|
||||
sel["cost_window_usage"] = self.cost_window_usage
|
||||
if self.cost_limit is not None:
|
||||
sel["cost_limit"] = self.cost_limit
|
||||
if self.cost_soft_threshold:
|
||||
sel["cost_soft_threshold"] = True
|
||||
return {"pool_selection": sel}
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class PoolSchedulingTrace:
|
||||
"""Aggregated scheduling trace for one pool-provider dispatch."""
|
||||
|
||||
provider_id: str
|
||||
total_keys: int = 0
|
||||
sticky_session_used: bool = False
|
||||
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]:
|
||||
"""Build compact dict for ``Usage.request_metadata["pool_summary"]``."""
|
||||
skipped_cooldown = 0
|
||||
skipped_cost = 0
|
||||
attempted = 0
|
||||
for t in self.candidate_traces.values():
|
||||
if t.skipped:
|
||||
if t.skip_type == "cooldown":
|
||||
skipped_cooldown += 1
|
||||
elif t.skip_type == "cost_exhausted":
|
||||
skipped_cost += 1
|
||||
else:
|
||||
attempted += 1
|
||||
|
||||
success_reason: str | None = None
|
||||
if success_key_id and success_key_id in self.candidate_traces:
|
||||
success_reason = self.candidate_traces[success_key_id].reason
|
||||
|
||||
summary: dict[str, Any] = {
|
||||
"enabled": True,
|
||||
"total_keys": self.total_keys,
|
||||
"attempted": attempted,
|
||||
"skipped_cooldown": skipped_cooldown,
|
||||
"skipped_cost": skipped_cost,
|
||||
"sticky_session": self.sticky_session_used,
|
||||
}
|
||||
if success_key_id:
|
||||
summary["success_key_id"] = success_key_id[:8]
|
||||
if success_reason:
|
||||
summary["success_reason"] = success_reason
|
||||
return summary
|
||||
106
tests/services/test_pool_config.py
Normal file
106
tests/services/test_pool_config.py
Normal file
@@ -0,0 +1,106 @@
|
||||
"""Tests for pool_config.py — configuration parsing and defaults."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from src.services.provider.pool.config import (
|
||||
PoolConfig,
|
||||
UnschedulableRule,
|
||||
parse_pool_config,
|
||||
)
|
||||
|
||||
|
||||
def test_parse_pool_config_returns_none_when_no_advanced_section() -> None:
|
||||
assert parse_pool_config({}) is None
|
||||
assert parse_pool_config(None) is None
|
||||
assert parse_pool_config({"other_key": 1}) is None
|
||||
|
||||
|
||||
def test_parse_pool_config_returns_defaults_for_empty_advanced() -> None:
|
||||
cfg = parse_pool_config({"pool_advanced": {}})
|
||||
assert cfg is not None
|
||||
assert cfg.sticky_session_ttl_seconds == 3600
|
||||
assert cfg.load_threshold_percent == 80
|
||||
assert cfg.lru_enabled is True
|
||||
assert cfg.cost_window_seconds == 18000
|
||||
assert cfg.cost_limit_per_key_tokens is None
|
||||
assert cfg.cost_soft_threshold_percent == 80
|
||||
assert cfg.rate_limit_cooldown_seconds == 300
|
||||
assert cfg.overload_cooldown_seconds == 30
|
||||
assert cfg.proactive_refresh_seconds == 180
|
||||
assert cfg.health_policy_enabled is True
|
||||
assert cfg.unschedulable_rules == []
|
||||
|
||||
|
||||
def test_parse_pool_config_overrides_values() -> None:
|
||||
cfg = parse_pool_config(
|
||||
{
|
||||
"pool_advanced": {
|
||||
"sticky_session_ttl_seconds": 7200,
|
||||
"load_threshold_percent": 90,
|
||||
"lru_enabled": False,
|
||||
"cost_window_seconds": 36000,
|
||||
"cost_limit_per_key_tokens": 100000,
|
||||
"cost_soft_threshold_percent": 70,
|
||||
"rate_limit_cooldown_seconds": 600,
|
||||
"overload_cooldown_seconds": 60,
|
||||
"proactive_refresh_seconds": 300,
|
||||
"health_policy_enabled": False,
|
||||
}
|
||||
}
|
||||
)
|
||||
assert cfg is not None
|
||||
assert cfg.sticky_session_ttl_seconds == 7200
|
||||
assert cfg.load_threshold_percent == 90
|
||||
assert cfg.lru_enabled is False
|
||||
assert cfg.cost_window_seconds == 36000
|
||||
assert cfg.cost_limit_per_key_tokens == 100000
|
||||
assert cfg.cost_soft_threshold_percent == 70
|
||||
assert cfg.rate_limit_cooldown_seconds == 600
|
||||
assert cfg.overload_cooldown_seconds == 60
|
||||
assert cfg.proactive_refresh_seconds == 300
|
||||
assert cfg.health_policy_enabled is False
|
||||
|
||||
|
||||
def test_parse_pool_config_parses_unschedulable_rules() -> None:
|
||||
cfg = parse_pool_config(
|
||||
{
|
||||
"pool_advanced": {
|
||||
"unschedulable_rules": [
|
||||
{"keyword": "rate_limit", "duration_minutes": 10},
|
||||
{"keyword": "overloaded"},
|
||||
{"invalid": "entry"}, # should be skipped
|
||||
"not_a_dict", # should be skipped
|
||||
]
|
||||
}
|
||||
}
|
||||
)
|
||||
assert cfg is not None
|
||||
assert len(cfg.unschedulable_rules) == 2
|
||||
assert cfg.unschedulable_rules[0] == UnschedulableRule(
|
||||
keyword="rate_limit", duration_minutes=10
|
||||
)
|
||||
assert cfg.unschedulable_rules[1] == UnschedulableRule(keyword="overloaded", duration_minutes=5)
|
||||
|
||||
|
||||
def test_parse_pool_config_handles_invalid_types_gracefully() -> None:
|
||||
# Invalid int values should fall back to defaults
|
||||
cfg = parse_pool_config(
|
||||
{
|
||||
"pool_advanced": {
|
||||
"sticky_session_ttl_seconds": "not_a_number",
|
||||
"cost_limit_per_key_tokens": "bad",
|
||||
}
|
||||
}
|
||||
)
|
||||
assert cfg is not None
|
||||
assert cfg.sticky_session_ttl_seconds == 3600 # default
|
||||
assert cfg.cost_limit_per_key_tokens is None # default for opt_int
|
||||
|
||||
|
||||
def test_pool_config_is_frozen() -> None:
|
||||
cfg = PoolConfig()
|
||||
try:
|
||||
cfg.sticky_session_ttl_seconds = 999 # type: ignore[misc]
|
||||
assert False, "Should have raised FrozenInstanceError"
|
||||
except AttributeError:
|
||||
pass
|
||||
54
tests/services/test_pool_config_backward_compat.py
Normal file
54
tests/services/test_pool_config_backward_compat.py
Normal file
@@ -0,0 +1,54 @@
|
||||
"""Tests for pool config backward compatibility.
|
||||
|
||||
Verifies that:
|
||||
1. Only ``pool_advanced`` key activates pool mode
|
||||
2. ``claude_code_advanced`` alone does NOT activate pool mode
|
||||
3. Old import paths via shim modules still resolve correctly
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from src.services.provider.pool.config import PoolConfig, parse_pool_config
|
||||
|
||||
|
||||
def test_pool_advanced_key_activates_pool() -> None:
|
||||
"""pool_advanced key activates pool mode."""
|
||||
cfg = parse_pool_config({"pool_advanced": {"sticky_session_ttl_seconds": 999}})
|
||||
assert cfg is not None
|
||||
assert cfg.sticky_session_ttl_seconds == 999
|
||||
|
||||
|
||||
def test_claude_code_advanced_alone_does_not_activate_pool() -> None:
|
||||
"""claude_code_advanced alone does NOT activate pool mode."""
|
||||
cfg = parse_pool_config({"claude_code_advanced": {"lru_enabled": False}})
|
||||
assert cfg is None
|
||||
|
||||
|
||||
def test_no_config_returns_none() -> None:
|
||||
cfg = parse_pool_config({"other": 1})
|
||||
assert cfg is None
|
||||
|
||||
|
||||
def test_empty_pool_advanced_returns_defaults() -> None:
|
||||
cfg = parse_pool_config({"pool_advanced": {}})
|
||||
assert cfg is not None
|
||||
defaults = PoolConfig()
|
||||
assert cfg.sticky_session_ttl_seconds == defaults.sticky_session_ttl_seconds
|
||||
assert cfg.lru_enabled is True
|
||||
|
||||
|
||||
def test_shim_imports_resolve() -> None:
|
||||
"""Old import paths through shim modules still work."""
|
||||
from src.services.provider.adapters.claude_code.pool_config import PoolConfig as ShimPoolConfig
|
||||
from src.services.provider.adapters.claude_code.pool_config import (
|
||||
parse_pool_config as shim_parse,
|
||||
)
|
||||
from src.services.provider.adapters.claude_code.pool_manager import (
|
||||
ClaudeCodePoolManager,
|
||||
)
|
||||
from src.services.provider.pool.manager import PoolManager
|
||||
|
||||
# ShimPoolConfig should be the same class
|
||||
assert ShimPoolConfig is PoolConfig
|
||||
assert shim_parse is parse_pool_config
|
||||
assert ClaudeCodePoolManager is PoolManager
|
||||
141
tests/services/test_pool_cost_tracker.py
Normal file
141
tests/services/test_pool_cost_tracker.py
Normal file
@@ -0,0 +1,141 @@
|
||||
"""Tests for pool_cost_tracker.py — rolling window cost tracking."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from src.services.provider.pool.config import PoolConfig
|
||||
from src.services.provider.pool.cost_tracker import (
|
||||
get_window_usage,
|
||||
is_approaching_limit,
|
||||
is_at_limit,
|
||||
record_usage,
|
||||
)
|
||||
|
||||
PID = "provider-test"
|
||||
KID = "key-test"
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def config() -> PoolConfig:
|
||||
return PoolConfig(
|
||||
cost_limit_per_key_tokens=10000,
|
||||
cost_soft_threshold_percent=80,
|
||||
cost_window_seconds=18000,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def config_no_limit() -> PoolConfig:
|
||||
return PoolConfig(cost_limit_per_key_tokens=None)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# record_usage
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_usage_calls_redis(config: PoolConfig) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.add_cost_entry",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_add:
|
||||
await record_usage(PID, KID, 500, config)
|
||||
|
||||
mock_add.assert_called_once_with(PID, KID, 500, 18000)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_usage_skips_zero_tokens(config: PoolConfig) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.add_cost_entry",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_add:
|
||||
await record_usage(PID, KID, 0, config)
|
||||
|
||||
mock_add.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_usage_skips_when_no_limit(config_no_limit: PoolConfig) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.add_cost_entry",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_add:
|
||||
await record_usage(PID, KID, 500, config_no_limit)
|
||||
|
||||
mock_add.assert_not_called()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# is_at_limit
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_at_limit_true(config: PoolConfig) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.get_cost_window_total",
|
||||
new_callable=AsyncMock,
|
||||
return_value=10000,
|
||||
):
|
||||
assert await is_at_limit(PID, KID, config) is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_at_limit_false(config: PoolConfig) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.get_cost_window_total",
|
||||
new_callable=AsyncMock,
|
||||
return_value=5000,
|
||||
):
|
||||
assert await is_at_limit(PID, KID, config) is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_at_limit_always_false_when_no_limit(config_no_limit: PoolConfig) -> None:
|
||||
assert await is_at_limit(PID, KID, config_no_limit) is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# is_approaching_limit
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_approaching_limit_true(config: PoolConfig) -> None:
|
||||
# 80% of 10000 = 8000
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.get_cost_window_total",
|
||||
new_callable=AsyncMock,
|
||||
return_value=8500,
|
||||
):
|
||||
assert await is_approaching_limit(PID, KID, config) is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_approaching_limit_false(config: PoolConfig) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.get_cost_window_total",
|
||||
new_callable=AsyncMock,
|
||||
return_value=5000,
|
||||
):
|
||||
assert await is_approaching_limit(PID, KID, config) is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_window_usage
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_window_usage(config: PoolConfig) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.get_cost_window_total",
|
||||
new_callable=AsyncMock,
|
||||
return_value=4200,
|
||||
):
|
||||
assert await get_window_usage(PID, KID, config) == 4200
|
||||
279
tests/services/test_pool_health_policy.py
Normal file
279
tests/services/test_pool_health_policy.py
Normal file
@@ -0,0 +1,279 @@
|
||||
"""Tests for pool_health_policy.py — error code classification."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from src.services.provider.pool.config import (
|
||||
PoolConfig,
|
||||
UnschedulableRule,
|
||||
)
|
||||
from src.services.provider.pool.health_policy import (
|
||||
_extract_error_message,
|
||||
_parse_retry_after,
|
||||
apply_health_policy,
|
||||
)
|
||||
|
||||
PID = "provider-test"
|
||||
KID = "key-test"
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def config() -> PoolConfig:
|
||||
return PoolConfig()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helper tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_parse_retry_after_none() -> None:
|
||||
assert _parse_retry_after(None) is None
|
||||
assert _parse_retry_after({}) is None
|
||||
|
||||
|
||||
def test_parse_retry_after_valid() -> None:
|
||||
assert _parse_retry_after({"retry-after": "60"}) == 60
|
||||
assert _parse_retry_after({"Retry-After": "120"}) == 120
|
||||
|
||||
|
||||
def test_parse_retry_after_clamped() -> None:
|
||||
assert _parse_retry_after({"retry-after": "0"}) == 1
|
||||
assert _parse_retry_after({"retry-after": "9999"}) == 3600
|
||||
|
||||
|
||||
def test_parse_retry_after_invalid() -> None:
|
||||
assert _parse_retry_after({"retry-after": "not-a-number"}) is None
|
||||
|
||||
|
||||
def test_extract_error_message_json_error_object() -> None:
|
||||
body = json.dumps({"error": {"message": "bad request"}})
|
||||
assert _extract_error_message(body) == "bad request"
|
||||
|
||||
|
||||
def test_extract_error_message_json_error_string() -> None:
|
||||
body = json.dumps({"error": "something went wrong"})
|
||||
assert _extract_error_message(body) == "something went wrong"
|
||||
|
||||
|
||||
def test_extract_error_message_plain_text() -> None:
|
||||
assert _extract_error_message("plain text error") == "plain text error"
|
||||
|
||||
|
||||
def test_extract_error_message_empty() -> None:
|
||||
assert _extract_error_message(None) == ""
|
||||
assert _extract_error_message("") == ""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Status code handling
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_401_invalidates_oauth_cache_and_sets_cooldown(config: PoolConfig) -> None:
|
||||
with (
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.invalidate_oauth_token_cache",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_inv,
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.set_cooldown",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_cd,
|
||||
):
|
||||
await apply_health_policy(
|
||||
provider_id=PID,
|
||||
key_id=KID,
|
||||
status_code=401,
|
||||
error_body=None,
|
||||
response_headers=None,
|
||||
config=config,
|
||||
)
|
||||
|
||||
mock_inv.assert_called_once_with(KID)
|
||||
mock_cd.assert_called_once_with(PID, KID, "auth_failed_401", ttl=60)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_402_sets_long_cooldown(config: PoolConfig) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.set_cooldown",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_cd:
|
||||
await apply_health_policy(
|
||||
provider_id=PID,
|
||||
key_id=KID,
|
||||
status_code=402,
|
||||
error_body=None,
|
||||
response_headers=None,
|
||||
config=config,
|
||||
)
|
||||
|
||||
mock_cd.assert_called_once_with(PID, KID, "payment_required_402", ttl=3600)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_403_sets_long_cooldown(config: PoolConfig) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.set_cooldown",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_cd:
|
||||
await apply_health_policy(
|
||||
provider_id=PID,
|
||||
key_id=KID,
|
||||
status_code=403,
|
||||
error_body=None,
|
||||
response_headers=None,
|
||||
config=config,
|
||||
)
|
||||
|
||||
mock_cd.assert_called_once_with(PID, KID, "forbidden_403", ttl=3600)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_400_with_org_disabled_pattern(config: PoolConfig) -> None:
|
||||
body = json.dumps({"error": {"message": "Your organization has been disabled"}})
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.set_cooldown",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_cd:
|
||||
await apply_health_policy(
|
||||
provider_id=PID,
|
||||
key_id=KID,
|
||||
status_code=400,
|
||||
error_body=body,
|
||||
response_headers=None,
|
||||
config=config,
|
||||
)
|
||||
|
||||
mock_cd.assert_called_once()
|
||||
assert "account_disabled_400" in mock_cd.call_args.args[2]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_400_without_pattern_does_nothing(config: PoolConfig) -> None:
|
||||
body = json.dumps({"error": {"message": "invalid json field"}})
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.set_cooldown",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_cd:
|
||||
await apply_health_policy(
|
||||
provider_id=PID,
|
||||
key_id=KID,
|
||||
status_code=400,
|
||||
error_body=body,
|
||||
response_headers=None,
|
||||
config=config,
|
||||
)
|
||||
|
||||
mock_cd.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_429_uses_retry_after_header(config: PoolConfig) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.set_cooldown",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_cd:
|
||||
await apply_health_policy(
|
||||
provider_id=PID,
|
||||
key_id=KID,
|
||||
status_code=429,
|
||||
error_body=None,
|
||||
response_headers={"retry-after": "120"},
|
||||
config=config,
|
||||
)
|
||||
|
||||
mock_cd.assert_called_once_with(PID, KID, "rate_limited_429", ttl=120)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_429_falls_back_to_config_default(config: PoolConfig) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.set_cooldown",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_cd:
|
||||
await apply_health_policy(
|
||||
provider_id=PID,
|
||||
key_id=KID,
|
||||
status_code=429,
|
||||
error_body=None,
|
||||
response_headers=None,
|
||||
config=config,
|
||||
)
|
||||
|
||||
mock_cd.assert_called_once_with(PID, KID, "rate_limited_429", ttl=300)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_529_uses_overload_cooldown(config: PoolConfig) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.set_cooldown",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_cd:
|
||||
await apply_health_policy(
|
||||
provider_id=PID,
|
||||
key_id=KID,
|
||||
status_code=529,
|
||||
error_body=None,
|
||||
response_headers=None,
|
||||
config=config,
|
||||
)
|
||||
|
||||
mock_cd.assert_called_once_with(PID, KID, "overloaded_529", ttl=30)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unschedulable keyword rules
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_keyword_rule_matches_and_sets_cooldown() -> None:
|
||||
cfg = PoolConfig(
|
||||
unschedulable_rules=[UnschedulableRule(keyword="capacity", duration_minutes=10)]
|
||||
)
|
||||
body = json.dumps({"error": {"message": "Server at capacity, try later"}})
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.set_cooldown",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_cd:
|
||||
await apply_health_policy(
|
||||
provider_id=PID,
|
||||
key_id=KID,
|
||||
status_code=500,
|
||||
error_body=body,
|
||||
response_headers=None,
|
||||
config=cfg,
|
||||
)
|
||||
|
||||
mock_cd.assert_called_once_with(PID, KID, "rule:capacity", ttl=600)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# health_policy_enabled = False
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disabled_health_policy_does_nothing() -> None:
|
||||
cfg = PoolConfig(health_policy_enabled=False)
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.set_cooldown",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_cd:
|
||||
await apply_health_policy(
|
||||
provider_id=PID,
|
||||
key_id=KID,
|
||||
status_code=429,
|
||||
error_body=None,
|
||||
response_headers=None,
|
||||
config=cfg,
|
||||
)
|
||||
|
||||
mock_cd.assert_not_called()
|
||||
274
tests/services/test_pool_manager.py
Normal file
274
tests/services/test_pool_manager.py
Normal file
@@ -0,0 +1,274 @@
|
||||
"""Tests for pool_manager.py — candidate reordering, success/error hooks."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from src.services.provider.pool.config import PoolConfig
|
||||
from src.services.provider.pool.manager import PoolManager
|
||||
|
||||
|
||||
def _make_candidate(key_id: str, *, is_skipped: bool = False) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
key=SimpleNamespace(id=key_id),
|
||||
is_skipped=is_skipped,
|
||||
skip_reason=None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def pool() -> PoolManager:
|
||||
return PoolManager("provider-1", PoolConfig())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# reorder_candidates
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reorder_empty_candidates(pool: PoolManager) -> None:
|
||||
result = await pool.reorder_candidates("sess-1", [])
|
||||
assert result == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reorder_sticky_hit_moves_to_front(pool: PoolManager) -> None:
|
||||
c1 = _make_candidate("key-1")
|
||||
c2 = _make_candidate("key-2")
|
||||
c3 = _make_candidate("key-3")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.get_sticky_binding",
|
||||
new_callable=AsyncMock,
|
||||
return_value="key-2",
|
||||
),
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.batch_get_cooldowns",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"key-1": None, "key-2": None, "key-3": None},
|
||||
),
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.get_lru_scores",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"key-1": 100.0, "key-2": 200.0, "key-3": 50.0},
|
||||
),
|
||||
):
|
||||
result = await pool.reorder_candidates("sess-1", [c1, c2, c3])
|
||||
|
||||
assert result[0].key.id == "key-2" # sticky hit first
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reorder_cooldown_keys_are_skipped(pool: PoolManager) -> None:
|
||||
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": "rate_limited_429", "key-2": None},
|
||||
),
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.get_lru_scores",
|
||||
new_callable=AsyncMock,
|
||||
return_value={},
|
||||
),
|
||||
):
|
||||
result = await pool.reorder_candidates(None, [c1, c2])
|
||||
|
||||
# key-1 should be skipped
|
||||
assert c1.is_skipped is True
|
||||
assert "cooldown" in (c1.skip_reason or "")
|
||||
# key-2 first in available
|
||||
assert result[0].key.id == "key-2"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reorder_cost_exhausted_keys_are_skipped() -> None:
|
||||
pool = PoolManager(
|
||||
"provider-1",
|
||||
PoolConfig(cost_limit_per_key_tokens=1000, lru_enabled=False),
|
||||
)
|
||||
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, "key-2": None},
|
||||
),
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.batch_get_cost_totals",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"key-1": 1500, "key-2": 200},
|
||||
),
|
||||
):
|
||||
result = await pool.reorder_candidates(None, [c1, c2])
|
||||
|
||||
assert c1.is_skipped is True
|
||||
assert "cost" in (c1.skip_reason or "")
|
||||
assert result[0].key.id == "key-2"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reorder_lru_sorts_least_recently_used_first(
|
||||
pool: PoolManager,
|
||||
) -> None:
|
||||
c1 = _make_candidate("key-1")
|
||||
c2 = _make_candidate("key-2")
|
||||
c3 = _make_candidate("key-3")
|
||||
|
||||
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, "key-2": None, "key-3": None},
|
||||
),
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.get_lru_scores",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"key-1": 300.0, "key-2": 100.0, "key-3": 200.0},
|
||||
),
|
||||
):
|
||||
result = await pool.reorder_candidates(None, [c1, c2, c3])
|
||||
|
||||
# Least recently used (lowest score) first
|
||||
available = [r for r in result if not r.is_skipped]
|
||||
assert available[0].key.id == "key-2"
|
||||
assert available[1].key.id == "key-3"
|
||||
assert available[2].key.id == "key-1"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# on_request_success
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_on_request_success_binds_sticky_and_touches_lru(
|
||||
pool: PoolManager,
|
||||
) -> None:
|
||||
with (
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.set_sticky_binding",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_sticky,
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.touch_lru",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_lru,
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.add_cost_entry",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_cost,
|
||||
):
|
||||
await pool.on_request_success(session_uuid="sess-1", key_id="key-1", tokens_used=0)
|
||||
|
||||
mock_sticky.assert_called_once_with("provider-1", "sess-1", "key-1", 3600)
|
||||
mock_lru.assert_called_once_with("provider-1", "key-1")
|
||||
mock_cost.assert_not_called() # tokens_used=0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_on_request_success_records_cost_when_configured() -> None:
|
||||
pool = PoolManager(
|
||||
"provider-1",
|
||||
PoolConfig(cost_limit_per_key_tokens=50000),
|
||||
)
|
||||
with (
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.set_sticky_binding",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.touch_lru",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.add_cost_entry",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_cost,
|
||||
):
|
||||
await pool.on_request_success(session_uuid="sess-1", key_id="key-1", tokens_used=500)
|
||||
|
||||
mock_cost.assert_called_once_with("provider-1", "key-1", 500, 18000)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# on_request_error
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_on_request_error_delegates_to_health_policy(
|
||||
pool: PoolManager,
|
||||
) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.health_policy.apply_health_policy",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_hp:
|
||||
await pool.on_request_error(
|
||||
key_id="key-1", status_code=429, error_body=None, response_headers=None
|
||||
)
|
||||
|
||||
mock_hp.assert_called_once()
|
||||
call_kwargs = mock_hp.call_args.kwargs
|
||||
assert call_kwargs["key_id"] == "key-1"
|
||||
assert call_kwargs["status_code"] == 429
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# is_key_schedulable
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_key_schedulable_returns_false_when_in_cooldown(
|
||||
pool: PoolManager,
|
||||
) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.get_cooldown",
|
||||
new_callable=AsyncMock,
|
||||
return_value="rate_limited_429",
|
||||
):
|
||||
ok, reason = await pool.is_key_schedulable("key-1")
|
||||
|
||||
assert ok is False
|
||||
assert "cooldown" in (reason or "")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_key_schedulable_returns_true_when_healthy(
|
||||
pool: PoolManager,
|
||||
) -> None:
|
||||
with patch(
|
||||
"src.services.provider.pool.redis_ops.get_cooldown",
|
||||
new_callable=AsyncMock,
|
||||
return_value=None,
|
||||
):
|
||||
ok, reason = await pool.is_key_schedulable("key-1")
|
||||
|
||||
assert ok is True
|
||||
assert reason is None
|
||||
202
tests/services/test_pool_manager_trace.py
Normal file
202
tests/services/test_pool_manager_trace.py
Normal file
@@ -0,0 +1,202 @@
|
||||
"""Tests for pool manager trace data collection and attachment."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from src.services.provider.pool.config import PoolConfig
|
||||
from src.services.provider.pool.manager import PoolManager
|
||||
from src.services.provider.pool.trace import PoolSchedulingTrace
|
||||
|
||||
|
||||
def _make_candidate(key_id: str, *, is_skipped: bool = False) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
key=SimpleNamespace(id=key_id),
|
||||
is_skipped=is_skipped,
|
||||
skip_reason=None,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Trace attachment to candidates
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trace_attached_to_first_candidate() -> 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={"key-1": 100.0, "key-2": 50.0},
|
||||
),
|
||||
):
|
||||
result = await pool.reorder_candidates("sess-1", [c1, c2])
|
||||
|
||||
# First candidate should have _pool_scheduling_trace
|
||||
trace = getattr(result[0], "_pool_scheduling_trace", None)
|
||||
assert trace is not None
|
||||
assert isinstance(trace, PoolSchedulingTrace)
|
||||
assert trace.total_keys == 2
|
||||
assert trace.provider_id == "prov-1"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pool_extra_data_on_selected_candidate() -> None:
|
||||
pool = PoolManager("prov-1", PoolConfig())
|
||||
c1 = _make_candidate("key-1")
|
||||
|
||||
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)},
|
||||
),
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.get_lru_scores",
|
||||
new_callable=AsyncMock,
|
||||
return_value={},
|
||||
),
|
||||
):
|
||||
result = await pool.reorder_candidates(None, [c1])
|
||||
|
||||
extra = getattr(result[0], "_pool_extra_data", None)
|
||||
assert extra is not None
|
||||
assert "pool_selection" in extra
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pool_extra_data_on_skipped_candidate() -> 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": ("rate_limited_429", 120),
|
||||
"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])
|
||||
|
||||
# key-1 is skipped with pool_skip extra data
|
||||
skipped = [r for r in result if r.is_skipped]
|
||||
assert len(skipped) == 1
|
||||
extra = getattr(skipped[0], "_pool_extra_data", None)
|
||||
assert extra is not None
|
||||
assert extra["pool_skip"]["type"] == "cooldown"
|
||||
assert extra["pool_skip"]["cooldown_reason"] == "rate_limited_429"
|
||||
assert extra["pool_skip"]["cooldown_ttl"] == 120
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trace_build_summary_matches() -> None:
|
||||
pool = PoolManager("prov-1", PoolConfig(cost_limit_per_key_tokens=1000))
|
||||
c1 = _make_candidate("key-1")
|
||||
c2 = _make_candidate("key-2")
|
||||
c3 = _make_candidate("key-3")
|
||||
|
||||
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": ("overloaded_529", 60),
|
||||
"key-2": (None, None),
|
||||
"key-3": (None, None),
|
||||
},
|
||||
),
|
||||
patch(
|
||||
"src.services.provider.pool.redis_ops.batch_get_cost_totals",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"key-2": 1500, "key-3": 200},
|
||||
),
|
||||
):
|
||||
result = await pool.reorder_candidates(None, [c1, c2, c3])
|
||||
|
||||
trace = getattr(result[0], "_pool_scheduling_trace", None)
|
||||
assert trace is not None
|
||||
|
||||
summary = trace.build_summary(success_key_id="key-3")
|
||||
assert summary["total_keys"] == 3
|
||||
assert summary["skipped_cooldown"] == 1 # key-1
|
||||
assert summary["skipped_cost"] == 1 # key-2
|
||||
assert summary["attempted"] == 1 # key-3
|
||||
assert summary["success_key_id"] == "key-3"[:8]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sticky_trace_info() -> None:
|
||||
pool = PoolManager("prov-1", PoolConfig(sticky_session_ttl_seconds=3600))
|
||||
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="key-2",
|
||||
),
|
||||
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("sess-1", [c1, c2])
|
||||
|
||||
trace = getattr(result[0], "_pool_scheduling_trace", None)
|
||||
assert trace is not None
|
||||
assert trace.sticky_session_used is True
|
||||
|
||||
# key-2 (sticky hit) should be first
|
||||
assert result[0].key.id == "key-2"
|
||||
extra = getattr(result[0], "_pool_extra_data", None)
|
||||
assert extra is not None
|
||||
assert extra["pool_selection"]["sticky_hit"] is True
|
||||
65
tests/services/test_pool_strategy.py
Normal file
65
tests/services/test_pool_strategy.py
Normal file
@@ -0,0 +1,65 @@
|
||||
"""Tests for pool scheduling strategy registry."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from src.services.provider.pool.strategy import (
|
||||
_strategy_registry,
|
||||
get_active_strategies,
|
||||
get_pool_strategy,
|
||||
register_pool_strategy,
|
||||
)
|
||||
|
||||
|
||||
class _DummyStrategy:
|
||||
name = "dummy"
|
||||
|
||||
def compute_score(self, *, key_id, config, context):
|
||||
return 42.0
|
||||
|
||||
|
||||
class _AnotherStrategy:
|
||||
name = "another"
|
||||
|
||||
def on_before_select(self, *, provider_id, key_ids, config, context):
|
||||
return key_ids[:1]
|
||||
|
||||
|
||||
def setup_function():
|
||||
_strategy_registry.clear()
|
||||
|
||||
|
||||
def teardown_function():
|
||||
_strategy_registry.clear()
|
||||
|
||||
|
||||
def test_register_and_get() -> None:
|
||||
s = _DummyStrategy()
|
||||
register_pool_strategy("dummy", s)
|
||||
assert get_pool_strategy("dummy") is s
|
||||
|
||||
|
||||
def test_get_nonexistent() -> None:
|
||||
assert get_pool_strategy("nonexistent") is None
|
||||
|
||||
|
||||
def test_get_active_strategies_filters_by_name() -> None:
|
||||
s1 = _DummyStrategy()
|
||||
s2 = _AnotherStrategy()
|
||||
register_pool_strategy("dummy", s1)
|
||||
register_pool_strategy("another", s2)
|
||||
|
||||
active = get_active_strategies(["dummy"])
|
||||
assert len(active) == 1
|
||||
assert active[0] is s1
|
||||
|
||||
active_both = get_active_strategies(["dummy", "another"])
|
||||
assert len(active_both) == 2
|
||||
|
||||
active_none = get_active_strategies(["missing"])
|
||||
assert len(active_none) == 0
|
||||
|
||||
|
||||
def test_get_active_strategies_empty_names() -> None:
|
||||
register_pool_strategy("dummy", _DummyStrategy())
|
||||
assert get_active_strategies([]) == []
|
||||
assert get_active_strategies(()) == []
|
||||
145
tests/services/test_pool_trace.py
Normal file
145
tests/services/test_pool_trace.py
Normal file
@@ -0,0 +1,145 @@
|
||||
"""Tests for pool scheduling trace dataclasses."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from src.services.provider.pool.trace import PoolCandidateTrace, PoolSchedulingTrace
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PoolCandidateTrace.to_extra_data
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestPoolCandidateTraceExtraData:
|
||||
def test_selected_sticky(self) -> None:
|
||||
ct = PoolCandidateTrace(key_id="k1", reason="sticky", sticky_hit=True)
|
||||
data = ct.to_extra_data()
|
||||
assert "pool_selection" in data
|
||||
assert data["pool_selection"]["reason"] == "sticky"
|
||||
assert data["pool_selection"]["sticky_hit"] is True
|
||||
|
||||
def test_selected_lru_with_cost(self) -> None:
|
||||
ct = PoolCandidateTrace(
|
||||
key_id="k2",
|
||||
reason="lru",
|
||||
lru_score=1234.5,
|
||||
cost_window_usage=500,
|
||||
cost_limit=1000,
|
||||
)
|
||||
data = ct.to_extra_data()
|
||||
sel = data["pool_selection"]
|
||||
assert sel["reason"] == "lru"
|
||||
assert sel["lru_score"] == 1234.5
|
||||
assert sel["cost_window_usage"] == 500
|
||||
assert sel["cost_limit"] == 1000
|
||||
assert "sticky_hit" not in sel
|
||||
|
||||
def test_selected_with_soft_threshold(self) -> None:
|
||||
ct = PoolCandidateTrace(
|
||||
key_id="k3",
|
||||
reason="lru",
|
||||
cost_window_usage=850,
|
||||
cost_limit=1000,
|
||||
cost_soft_threshold=True,
|
||||
)
|
||||
data = ct.to_extra_data()
|
||||
assert data["pool_selection"]["cost_soft_threshold"] is True
|
||||
|
||||
def test_selected_random_minimal(self) -> None:
|
||||
ct = PoolCandidateTrace(key_id="k4", reason="random")
|
||||
data = ct.to_extra_data()
|
||||
sel = data["pool_selection"]
|
||||
assert sel == {"reason": "random"}
|
||||
|
||||
def test_skipped_cooldown(self) -> None:
|
||||
ct = PoolCandidateTrace(
|
||||
key_id="k5",
|
||||
skipped=True,
|
||||
skip_type="cooldown",
|
||||
cooldown_reason="rate_limited_429",
|
||||
cooldown_ttl=120,
|
||||
)
|
||||
data = ct.to_extra_data()
|
||||
assert "pool_skip" in data
|
||||
skip = data["pool_skip"]
|
||||
assert skip["type"] == "cooldown"
|
||||
assert skip["cooldown_reason"] == "rate_limited_429"
|
||||
assert skip["cooldown_ttl"] == 120
|
||||
|
||||
def test_skipped_cost_exhausted(self) -> None:
|
||||
ct = PoolCandidateTrace(
|
||||
key_id="k6",
|
||||
skipped=True,
|
||||
skip_type="cost_exhausted",
|
||||
cost_window_usage=2000,
|
||||
)
|
||||
data = ct.to_extra_data()
|
||||
skip = data["pool_skip"]
|
||||
assert skip["type"] == "cost_exhausted"
|
||||
assert skip["cost_window_usage"] == 2000
|
||||
assert "cooldown_reason" not in skip
|
||||
|
||||
def test_skipped_minimal(self) -> None:
|
||||
ct = PoolCandidateTrace(key_id="k7", skipped=True, skip_type="upstream")
|
||||
data = ct.to_extra_data()
|
||||
skip = data["pool_skip"]
|
||||
assert skip == {"type": "upstream"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PoolSchedulingTrace.build_summary
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestPoolSchedulingTraceSummary:
|
||||
def test_basic_summary(self) -> None:
|
||||
trace = PoolSchedulingTrace(provider_id="prov-1", total_keys=5)
|
||||
trace.candidate_traces = {
|
||||
"k1": PoolCandidateTrace(key_id="k1", reason="sticky", sticky_hit=True),
|
||||
"k2": PoolCandidateTrace(key_id="k2", reason="lru"),
|
||||
"k3": PoolCandidateTrace(
|
||||
key_id="k3", skipped=True, skip_type="cooldown", cooldown_reason="429"
|
||||
),
|
||||
"k4": PoolCandidateTrace(key_id="k4", skipped=True, skip_type="cost_exhausted"),
|
||||
"k5": PoolCandidateTrace(
|
||||
key_id="k5", skipped=True, skip_type="cooldown", cooldown_reason="500"
|
||||
),
|
||||
}
|
||||
trace.sticky_session_used = True
|
||||
|
||||
summary = trace.build_summary(success_key_id="k1")
|
||||
|
||||
assert summary["enabled"] is True
|
||||
assert summary["total_keys"] == 5
|
||||
assert summary["attempted"] == 2
|
||||
assert summary["skipped_cooldown"] == 2
|
||||
assert summary["skipped_cost"] == 1
|
||||
assert summary["sticky_session"] is True
|
||||
assert summary["success_key_id"] == "k1"[:8]
|
||||
assert summary["success_reason"] == "sticky"
|
||||
|
||||
def test_summary_no_success_key(self) -> None:
|
||||
trace = PoolSchedulingTrace(provider_id="prov-2", total_keys=2)
|
||||
trace.candidate_traces = {
|
||||
"k1": PoolCandidateTrace(key_id="k1", reason="random"),
|
||||
"k2": PoolCandidateTrace(key_id="k2", reason="lru"),
|
||||
}
|
||||
|
||||
summary = trace.build_summary(success_key_id=None)
|
||||
|
||||
assert summary["attempted"] == 2
|
||||
assert summary["skipped_cooldown"] == 0
|
||||
assert summary["skipped_cost"] == 0
|
||||
assert "success_key_id" not in summary
|
||||
assert "success_reason" not in summary
|
||||
|
||||
def test_summary_all_skipped(self) -> None:
|
||||
trace = PoolSchedulingTrace(provider_id="prov-3", total_keys=2)
|
||||
trace.candidate_traces = {
|
||||
"k1": PoolCandidateTrace(key_id="k1", skipped=True, skip_type="cooldown"),
|
||||
"k2": PoolCandidateTrace(key_id="k2", skipped=True, skip_type="cost_exhausted"),
|
||||
}
|
||||
|
||||
summary = trace.build_summary()
|
||||
assert summary["attempted"] == 0
|
||||
assert summary["skipped_cooldown"] == 1
|
||||
assert summary["skipped_cost"] == 1
|
||||
Reference in New Issue
Block a user