mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +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