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

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

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

Closes #197

Co-authored-by: AAEE86 <33052466+AAEE86@users.noreply.github.com>
This commit is contained in:
fawney19
2026-03-03 09:22:20 +08:00
parent f787b1b02a
commit 11997c024e
40 changed files with 4419 additions and 823 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
src/api/admin/monitoring/trace.py
View File

@@ -163,13 +163,19 @@ class AdminGetRequestTraceAdapter(AdminApiAdapter):
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
db = context.db
# 只查询 candidates
candidates = RequestCandidateService.get_candidates_by_request_id(db, self.request_id)
# 查询所有候选后，默认只展示已发生调度结果的子集：
# - 过滤 available/unused（预创建但未实际参与本次调度）
# - 若过滤后为空（例如请求尚未开始），回退到全量，避免前端空白
all_candidates = RequestCandidateService.get_candidates_by_request_id(db, self.request_id)
# 如果没有数据，返回 404
if not candidates:
if not all_candidates:
raise HTTPException(status_code=404, detail="Request not found")
candidates = [
c for c in all_candidates if c.status not in ("available", "unused")
] or all_candidates
# 计算总延迟（只统计已完成的候选：success, failed, cancelled）
# 使用显式的 is not None 检查，避免过滤掉 0ms 的快速响应
total_latency = sum(

562
src/api/admin/pool/routes.py
View File

@@ -9,12 +9,14 @@ Provides endpoints for managing account pools at scale:
from __future__ import annotations
import asyncio
import json
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any
from fastapi import APIRouter, Depends, Query, Request
from sqlalchemy import func
from sqlalchemy.orm import Session
from src.api.base.admin_adapter import AdminApiAdapter
@@ -24,9 +26,14 @@ from src.core.crypto import crypto_service
from src.core.exceptions import NotFoundException
from src.core.logger import logger
from src.database import get_db
from src.models.database import Provider, ProviderAPIKey
from src.models.database import Provider, ProviderAPIKey, Usage
from src.services.provider.pool import redis_ops as pool_redis
from src.services.provider.pool.config import parse_pool_config
from src.services.provider.pool.scheduling_dimensions import (
PoolSchedulingSnapshot,
evaluate_pool_scheduling_dimensions,
summarize_pool_scheduling_dimensions,
)
from .schemas import (
BatchActionRequest,
@@ -38,6 +45,8 @@ from .schemas import (
PoolKeysPageResponse,
PoolOverviewItem,
PoolOverviewResponse,
PoolSchedulingDimension,
PoolSchedulingReason,
)
router = APIRouter(prefix="/api/admin/pool", tags=["pool-management"])
@@ -108,6 +117,33 @@ async def batch_import_keys(
ALLOWED_ACTIONS = {"enable", "disable", "delete", "clear_cooldown", "reset_cost"}
_COOLDOWN_REASON_LABELS: dict[str, str] = {
"rate_limited_429": "429 限流",
"forbidden_403": "403 禁止",
"overloaded_529": "529 过载",
"auth_failed_401": "401 认证失败",
"payment_required_402": "402 欠费",
"server_error_500": "500 错误",
}
_ACCOUNT_BLOCK_REASON_KEYWORDS: tuple[str, ...] = (
"account_block",
"account blocked",
"account has been disabled",
"account disabled",
"organization has been disabled",
"organization_disabled",
"validation_required",
"verify your account",
"forbidden",
"suspended",
"封禁",
"封号",
"被封",
"访问被禁止",
"账号异常",
)
def _to_float(value: Any) -> float | None:
if isinstance(value, bool):
@@ -125,6 +161,67 @@ def _to_float(value: Any) -> float | None:
return None
def _is_truthy_flag(value: Any) -> bool:
if isinstance(value, bool):
return value
if isinstance(value, (int, float)):
return value != 0
if isinstance(value, str):
normalized = value.strip().lower()
return normalized in {"1", "true", "yes", "y"}
return False
def _is_known_banned_reason(reason: str | None) -> bool:
if not reason:
return False
text = str(reason).strip()
if not text:
return False
lowered = text.lower()
# 结构化账号级别封禁标记（如 [ACCOUNT_BLOCK] ...）
try:
from src.services.provider.oauth_token import is_account_level_block
if is_account_level_block(text):
return True
except Exception:
pass
return any(keyword in lowered for keyword in _ACCOUNT_BLOCK_REASON_KEYWORDS)
def _is_known_banned_key(key: ProviderAPIKey, provider_type: str) -> bool:
upstream_metadata = getattr(key, "upstream_metadata", None)
normalized_provider = provider_type.strip().lower()
provider_bucket: dict[str, Any] | None = None
if isinstance(upstream_metadata, dict):
maybe_bucket = upstream_metadata.get(normalized_provider)
if isinstance(maybe_bucket, dict):
provider_bucket = maybe_bucket
if normalized_provider == "kiro" and provider_bucket:
if _is_truthy_flag(provider_bucket.get("is_banned")):
return True
if normalized_provider == "antigravity" and provider_bucket:
if _is_truthy_flag(provider_bucket.get("is_forbidden")):
return True
for source in (provider_bucket, upstream_metadata):
if not isinstance(source, dict):
continue
if _is_truthy_flag(source.get("is_banned")):
return True
if _is_truthy_flag(source.get("is_forbidden")):
return True
if _is_truthy_flag(source.get("account_disabled")):
return True
return _is_known_banned_reason(getattr(key, "oauth_invalid_reason", None))
def _format_percent(value: float) -> str:
clamped = max(0.0, min(value, 100.0))
return f"{clamped:.1f}%"
@@ -177,11 +274,7 @@ def _build_kiro_account_quota(upstream_metadata: dict[str, Any]) -> str | None:
remaining = 100.0 - usage_percentage
current_usage = _to_float(kiro.get("current_usage"))
usage_limit = _to_float(kiro.get("usage_limit"))
if (
current_usage is not None
and usage_limit is not None
and usage_limit > 0
):
if current_usage is not None and usage_limit is not None and usage_limit > 0:
return (
f"剩余 {_format_percent(remaining)} "
f"({_format_quota_value(current_usage)}/{_format_quota_value(usage_limit)})"
@@ -247,6 +340,215 @@ def _build_account_quota(provider_type: str, upstream_metadata: Any) -> str | No
return None
def _extract_quota_updated_at(provider_type: str, upstream_metadata: Any) -> int | None:
if not isinstance(upstream_metadata, dict):
return None
normalized_type = provider_type.strip().lower()
if normalized_type == "codex":
source = upstream_metadata.get("codex")
elif normalized_type == "antigravity":
source = upstream_metadata.get("antigravity")
elif normalized_type == "kiro":
source = upstream_metadata.get("kiro")
else:
return None
if not isinstance(source, dict):
return None
updated_at = _to_float(source.get("updated_at"))
if updated_at is None or updated_at <= 0:
return None
# 部分上游可能返回毫秒时间戳，统一转换为秒
if updated_at > 1_000_000_000_000:
updated_at /= 1000
return int(updated_at)
def _normalize_oauth_plan_type(plan_type: Any, provider_type: str) -> str | None:
if not isinstance(plan_type, str):
return None
text = plan_type.strip()
if not text:
return None
ptype = provider_type.strip().lower()
if ptype and text.lower().startswith(ptype):
trimmed = text[len(ptype) :].strip(" :-_")
if trimmed:
text = trimmed
return text or None
def _derive_oauth_plan_type(key: ProviderAPIKey, provider_type: str) -> str | None:
# Prefer persisted normalized field
persisted = _normalize_oauth_plan_type(getattr(key, "oauth_plan_type", None), provider_type)
if persisted:
return persisted
if str(getattr(key, "auth_type", "") or "").strip().lower() != "oauth":
return None
# Fallback 1: encrypted auth_config (common for Codex/Antigravity)
auth_config_raw = getattr(key, "auth_config", None)
if auth_config_raw:
try:
decrypted = crypto_service.decrypt(auth_config_raw)
auth_config = json.loads(decrypted)
if isinstance(auth_config, dict):
for plan_key in ("plan_type", "tier", "plan", "subscription_plan"):
normalized = _normalize_oauth_plan_type(
auth_config.get(plan_key), provider_type
)
if normalized:
return normalized
except Exception:
pass
# Fallback 2: upstream_metadata
upstream_metadata = getattr(key, "upstream_metadata", None)
if not isinstance(upstream_metadata, dict):
return None
provider_bucket = upstream_metadata.get(provider_type.strip().lower())
candidates: list[dict[str, Any]] = []
if isinstance(provider_bucket, dict):
candidates.append(provider_bucket)
candidates.append(upstream_metadata)
for source in candidates:
for plan_key in ("plan_type", "tier", "subscription_title", "subscription_plan"):
normalized = _normalize_oauth_plan_type(source.get(plan_key), provider_type)
if normalized:
return normalized
return None
def _compute_health_aggregate(
health_by_format: Any, circuit_breaker_by_format: Any
) -> tuple[float, bool]:
"""从按格式健康数据聚合出列表展示字段。"""
health_map = health_by_format if isinstance(health_by_format, dict) else {}
circuit_map = circuit_breaker_by_format if isinstance(circuit_breaker_by_format, dict) else {}
if health_map:
scores = [
float(item.get("health_score") or 1.0)
for item in health_map.values()
if isinstance(item, dict)
]
health_score = min(scores) if scores else 1.0
else:
health_score = 1.0
any_circuit_open = any(
bool(item.get("open", False)) for item in circuit_map.values() if isinstance(item, dict)
)
return health_score, any_circuit_open
def _format_cooldown_detail(raw: str | None) -> str | None:
if not raw:
return None
return _COOLDOWN_REASON_LABELS.get(raw, raw)
def _build_pool_scheduling_state(
*,
is_active: bool,
cooldown_reason: str | None,
cooldown_ttl_seconds: int | None,
circuit_breaker_open: bool,
cost_window_usage: int,
cost_limit: int | None,
cost_soft_threshold_percent: int,
health_score: float,
) -> tuple[
str,
str,
str,
list[PoolSchedulingReason],
float,
bool,
int,
int,
list[PoolSchedulingDimension],
]:
"""Build unified scheduling state for frontend display."""
snapshot = PoolSchedulingSnapshot(
is_active=is_active,
cooldown_reason=cooldown_reason,
cooldown_ttl_seconds=cooldown_ttl_seconds,
circuit_breaker_open=circuit_breaker_open,
cost_window_usage=cost_window_usage,
cost_limit=cost_limit,
cost_soft_threshold_percent=cost_soft_threshold_percent,
health_score=health_score,
)
dimensions_raw = evaluate_pool_scheduling_dimensions(snapshot)
summary = summarize_pool_scheduling_dimensions(dimensions_raw)
scheduling_dimensions: list[PoolSchedulingDimension] = []
scheduling_reasons: list[PoolSchedulingReason] = []
for item in dimensions_raw:
detail = item.detail
if item.code == "cooldown":
detail = _format_cooldown_detail(detail)
model = PoolSchedulingDimension(
code=item.code,
label=item.label,
status=item.status,
blocking=bool(item.blocking or item.status == "blocked"),
source=item.source,
weight=item.weight,
score=item.score,
ttl_seconds=item.ttl_seconds,
detail=detail,
)
scheduling_dimensions.append(model)
if item.status != "ok":
scheduling_reasons.append(
PoolSchedulingReason(
code=item.code,
label=item.label,
blocking=bool(item.blocking or item.status == "blocked"),
source=item.source,
ttl_seconds=item.ttl_seconds,
detail=detail,
)
)
return (
summary.status,
summary.reason,
summary.label,
scheduling_reasons,
summary.score,
summary.candidate_eligible,
summary.blocked_count,
summary.degraded_count,
scheduling_dimensions,
)
def _mask_proxy_password(proxy_config: Any) -> dict[str, Any] | None:
if not isinstance(proxy_config, dict):
return None
masked = dict(proxy_config)
password = masked.get("password")
if isinstance(password, str) and password:
masked["password"] = "******"
return masked
@router.post("/{provider_id}/keys/batch-action", response_model=BatchActionResponse)
async def batch_action_keys(
provider_id: str,
@@ -259,6 +561,17 @@ async def batch_action_keys(
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
@router.post("/{provider_id}/keys/cleanup-banned", response_model=BatchActionResponse)
async def cleanup_banned_keys(
provider_id: str,
request: Request,
db: Session = Depends(get_db),
) -> BatchActionResponse:
"""Delete known banned/suspended accounts for the provider."""
adapter = AdminCleanupBannedKeysAdapter(provider_id=provider_id)
return await pipeline.run(adapter=adapter, http_request=request, db=db, mode=adapter.mode)
# ---------------------------------------------------------------------------
# Adapters
# ---------------------------------------------------------------------------
@@ -380,29 +693,135 @@ class AdminListPoolKeysAdapter(AdminApiAdapter):
if pcfg
else asyncio.sleep(0, result={})
)
cooldowns, cooldown_ttls, lru_scores, cost_totals = await asyncio.gather(
cooldowns, cooldown_ttls, lru_scores, cost_totals, sticky_counts = await asyncio.gather(
pool_redis.batch_get_cooldowns(pid, key_ids),
pool_redis.batch_get_cooldown_ttls(pid, key_ids),
_lru_coro,
_cost_coro,
pool_redis.batch_get_key_sticky_counts(pid, key_ids),
)
else:
cooldowns, cooldown_ttls, lru_scores, cost_totals = {}, {}, {}, {}
# Sticky session count per key is expensive (SCAN+MGET per key).
# Only compute when the page is small enough to avoid timeout.
sticky_counts: dict[str, int] = {}
if key_ids and len(key_ids) <= 30:
counts = await asyncio.gather(
*(pool_redis.get_key_sticky_count(pid, kid) for kid in key_ids)
cooldowns, cooldown_ttls, lru_scores, cost_totals, sticky_counts = (
{},
{},
{},
{},
{},
)
sticky_counts = dict(zip(key_ids, counts))
usage_stats_by_key: dict[str, dict[str, Any]] = {}
if key_ids:
usage_rows = (
db.query(
Usage.provider_api_key_id.label("key_id"),
func.count(Usage.id).label("request_count"),
func.coalesce(func.sum(Usage.total_tokens), 0).label("total_tokens"),
func.coalesce(func.sum(Usage.total_cost_usd), 0.0).label("total_cost_usd"),
func.max(Usage.created_at).label("last_used_at"),
)
.filter(
Usage.provider_id == pid,
Usage.provider_api_key_id.in_(key_ids),
Usage.status.notin_(["pending", "streaming"]),
)
.group_by(Usage.provider_api_key_id)
.all()
)
usage_stats_by_key = {
str(row.key_id): {
"request_count": int(row.request_count or 0),
"total_tokens": int(row.total_tokens or 0),
"total_cost_usd": float(row.total_cost_usd or 0.0),
"last_used_at": getattr(row, "last_used_at", None),
}
for row in usage_rows
if getattr(row, "key_id", None)
}
key_details: list[PoolKeyDetail] = []
for k in keys:
kid = str(k.id)
cd_reason = cooldowns.get(kid)
cd_ttl = cooldown_ttls.get(kid) if cd_reason else None
health_score, any_circuit_open = _compute_health_aggregate(
getattr(k, "health_by_format", None),
getattr(k, "circuit_breaker_by_format", None),
)
cost_usage = int(cost_totals.get(kid, 0) or 0)
cost_limit = pcfg.cost_limit_per_key_tokens if pcfg else None
(
scheduling_status,
scheduling_reason,
scheduling_label,
scheduling_reasons,
scheduling_score,
candidate_eligible,
scheduling_blocked_count,
scheduling_degraded_count,
scheduling_dimensions,
) = _build_pool_scheduling_state(
is_active=bool(k.is_active),
cooldown_reason=cd_reason,
cooldown_ttl_seconds=cd_ttl,
circuit_breaker_open=any_circuit_open,
cost_window_usage=cost_usage,
cost_limit=cost_limit,
cost_soft_threshold_percent=(pcfg.cost_soft_threshold_percent if pcfg else 80),
health_score=health_score,
)
raw_allowed_models = getattr(k, "allowed_models", None)
allowed_models = (
[str(item) for item in raw_allowed_models]
if isinstance(raw_allowed_models, list)
else None
)
raw_locked_models = getattr(k, "locked_models", None)
locked_models = (
[str(item) for item in raw_locked_models]
if isinstance(raw_locked_models, list)
else None
)
raw_include_patterns = getattr(k, "model_include_patterns", None)
include_patterns = (
[str(item) for item in raw_include_patterns]
if isinstance(raw_include_patterns, list)
else None
)
raw_exclude_patterns = getattr(k, "model_exclude_patterns", None)
exclude_patterns = (
[str(item) for item in raw_exclude_patterns]
if isinstance(raw_exclude_patterns, list)
else None
)
capabilities = (
{str(name): bool(enabled) for name, enabled in k.capabilities.items()}
if isinstance(getattr(k, "capabilities", None), dict)
else None
)
rate_multipliers: dict[str, float] | None = None
if isinstance(getattr(k, "rate_multipliers", None), dict):
converted: dict[str, float] = {}
for fmt, raw_val in k.rate_multipliers.items():
num_val = _to_float(raw_val)
if num_val is None:
continue
converted[str(fmt)] = num_val
rate_multipliers = converted or None
api_formats = (
[str(fmt) for fmt in getattr(k, "api_formats", []) if isinstance(fmt, str)]
if isinstance(getattr(k, "api_formats", None), list)
else []
)
key_usage_stats = usage_stats_by_key.get(kid, {})
key_request_count = int(
key_usage_stats.get("request_count") or getattr(k, "request_count", 0) or 0
)
key_total_tokens = int(key_usage_stats.get("total_tokens") or 0)
key_total_cost_usd = float(key_usage_stats.get("total_cost_usd") or 0.0)
key_last_used_at = getattr(k, "last_used_at", None) or key_usage_stats.get(
"last_used_at"
)
key_details.append(
PoolKeyDetail(
@@ -410,22 +829,66 @@ class AdminListPoolKeysAdapter(AdminApiAdapter):
key_name=k.name or "",
is_active=bool(k.is_active),
auth_type=str(getattr(k, "auth_type", "api_key") or "api_key"),
oauth_expires_at=(
int(k.oauth_expires_at.timestamp())
if getattr(k, "oauth_expires_at", None)
else None
),
oauth_invalid_at=(
int(k.oauth_invalid_at.timestamp())
if getattr(k, "oauth_invalid_at", None)
else None
),
oauth_invalid_reason=getattr(k, "oauth_invalid_reason", None),
oauth_plan_type=_derive_oauth_plan_type(k, provider_type),
quota_updated_at=_extract_quota_updated_at(
provider_type,
getattr(k, "upstream_metadata", None),
),
health_score=health_score,
circuit_breaker_open=any_circuit_open,
api_formats=api_formats,
rate_multipliers=rate_multipliers,
internal_priority=int(getattr(k, "internal_priority", 50) or 50),
rpm_limit=getattr(k, "rpm_limit", None),
cache_ttl_minutes=int(getattr(k, "cache_ttl_minutes", 5) or 5),
max_probe_interval_minutes=int(
getattr(k, "max_probe_interval_minutes", 32) or 32
),
note=getattr(k, "note", None),
allowed_models=allowed_models,
capabilities=capabilities,
auto_fetch_models=bool(getattr(k, "auto_fetch_models", False)),
locked_models=locked_models,
model_include_patterns=include_patterns,
model_exclude_patterns=exclude_patterns,
proxy=_mask_proxy_password(getattr(k, "proxy", None)),
account_quota=_build_account_quota(
provider_type,
getattr(k, "upstream_metadata", None),
),
cooldown_reason=cd_reason,
cooldown_ttl_seconds=cd_ttl,
cost_window_usage=cost_totals.get(kid, 0),
cost_limit=pcfg.cost_limit_per_key_tokens if pcfg else None,
cost_window_usage=cost_usage,
cost_limit=cost_limit,
request_count=key_request_count,
total_tokens=key_total_tokens,
total_cost_usd=key_total_cost_usd,
sticky_sessions=sticky_counts.get(kid, 0),
lru_score=lru_scores.get(kid),
created_at=(
k.created_at.isoformat() if getattr(k, "created_at", None) else None
),
last_used_at=(
k.last_used_at.isoformat() if getattr(k, "last_used_at", None) else None
),
last_used_at=(key_last_used_at.isoformat() if key_last_used_at else None),
scheduling_status=scheduling_status,
scheduling_reason=scheduling_reason,
scheduling_label=scheduling_label,
scheduling_reasons=scheduling_reasons,
scheduling_score=scheduling_score,
candidate_eligible=candidate_eligible,
scheduling_blocked_count=scheduling_blocked_count,
scheduling_degraded_count=scheduling_degraded_count,
scheduling_dimensions=scheduling_dimensions,
)
)
@@ -447,6 +910,9 @@ class AdminBatchImportKeysAdapter(AdminApiAdapter):
provider = db.query(Provider).filter(Provider.id == self.provider_id).first()
if not provider:
raise NotFoundException("Provider not found", "provider")
key_proxy: dict[str, Any] | None = None
if self.body.proxy_node_id and self.body.proxy_node_id.strip():
key_proxy = {"node_id": self.body.proxy_node_id.strip(), "enabled": True}
imported = 0
skipped = 0
@@ -466,6 +932,7 @@ class AdminBatchImportKeysAdapter(AdminApiAdapter):
name=item.name or f"imported-{idx}",
api_key=encrypted_key,
auth_type=item.auth_type or "api_key",
proxy=key_proxy,
is_active=True,
created_at=now,
updated_at=now,
@@ -591,3 +1058,56 @@ class AdminBatchActionKeysAdapter(AdminApiAdapter):
affected=affected,
message=f"{affected} keys {action_labels.get(self.body.action, self.body.action)}",
)
@dataclass
class AdminCleanupBannedKeysAdapter(AdminApiAdapter):
provider_id: str = ""
async def handle(self, context: ApiRequestContext) -> Any: # type: ignore[override]
db = context.db
provider = db.query(Provider).filter(Provider.id == self.provider_id).first()
if not provider:
raise NotFoundException("Provider not found", "provider")
pid = str(provider.id)
provider_type = str(getattr(provider, "provider_type", "") or "").strip().lower()
keys = db.query(ProviderAPIKey).filter(ProviderAPIKey.provider_id == pid).all()
banned_keys = [key for key in keys if _is_known_banned_key(key, provider_type)]
if not banned_keys:
return BatchActionResponse(affected=0, message="未发现已知封号账号")
banned_key_ids = [str(key.id) for key in banned_keys]
for key in banned_keys:
db.delete(key)
try:
db.commit()
except Exception as exc:
db.rollback()
logger.error("cleanup banned keys commit failed: {}", exc)
return BatchActionResponse(affected=0, message=f"commit failed: {exc}")
# 清理 Redis 中可能残留的状态，避免删除后仍有旧状态占用资源。
cleanup_coros = []
for kid in banned_key_ids:
cleanup_coros.append(pool_redis.clear_cooldown(pid, kid))
cleanup_coros.append(pool_redis.clear_cost(pid, kid))
if cleanup_coros:
await asyncio.gather(*cleanup_coros, return_exceptions=True)
admin_name = context.user.username if context.user else "admin"
logger.warning(
"Pool cleanup banned by {}: provider={}, affected={}, key_ids={}",
admin_name,
self.provider_id[:8],
len(banned_key_ids),
[kid[:8] for kid in banned_key_ids],
)
return BatchActionResponse(
affected=len(banned_key_ids),
message=f"已清理 {len(banned_key_ids)} 个已知封号账号",
)

66
src/api/admin/pool/schemas.py
View File

@@ -2,6 +2,8 @@
from __future__ import annotations
from typing import Any
from pydantic import BaseModel, ConfigDict, Field
# ---------------------------------------------------------------------------
@@ -32,6 +34,31 @@ class PoolOverviewResponse(BaseModel):
# ---------------------------------------------------------------------------
class PoolSchedulingReason(BaseModel):
"""Structured scheduling reason for a key."""
code: str
label: str
blocking: bool = False
source: str = "pool" # manual / pool / health / policy
ttl_seconds: int | None = None
detail: str | None = None
class PoolSchedulingDimension(BaseModel):
"""Detailed scheduling dimension status."""
code: str
label: str
status: str = "ok" # ok / degraded / blocked
blocking: bool = False
source: str = "pool"
weight: int = 1
score: float = 1.0 # normalized 0~1
ttl_seconds: int | None = None
detail: str | None = None
class PoolKeyDetail(BaseModel):
"""Detailed status of a single pool key."""
@@ -39,15 +66,50 @@ class PoolKeyDetail(BaseModel):
key_name: str
is_active: bool
auth_type: str = "api_key"
oauth_expires_at: int | None = None
oauth_invalid_at: int | None = None
oauth_invalid_reason: str | None = None
oauth_plan_type: str | None = None
quota_updated_at: int | None = None
# 健康度聚合字段（与 Provider Key 列表口径一致）
health_score: float = 1.0
circuit_breaker_open: bool = False
# 编辑/权限/代理所需字段
api_formats: list[str] = Field(default_factory=list)
rate_multipliers: dict[str, float] | None = None
internal_priority: int = 50
rpm_limit: int | None = None
cache_ttl_minutes: int = 5
max_probe_interval_minutes: int = 32
note: str | None = None
allowed_models: list[str] | None = None
capabilities: dict[str, bool] | None = None
auto_fetch_models: bool = False
locked_models: list[str] | None = None
model_include_patterns: list[str] | None = None
model_exclude_patterns: list[str] | None = None
proxy: dict[str, Any] | None = None
account_quota: str | None = None
cooldown_reason: str | None = None
cooldown_ttl_seconds: int | None = None
cost_window_usage: int = 0
cost_limit: int | None = None
request_count: int = 0
total_tokens: int = 0
total_cost_usd: float = 0.0
sticky_sessions: int = 0
lru_score: float | None = None
created_at: str | None = None
last_used_at: str | None = None
scheduling_status: str = "available" # available / degraded / blocked
scheduling_reason: str = "available"
scheduling_label: str = "可用"
scheduling_reasons: list[PoolSchedulingReason] = Field(default_factory=list)
scheduling_score: float = 100.0
candidate_eligible: bool = True
scheduling_blocked_count: int = 0
scheduling_degraded_count: int = 0
scheduling_dimensions: list[PoolSchedulingDimension] = Field(default_factory=list)
model_config = ConfigDict(from_attributes=True)
@@ -76,6 +138,10 @@ class PoolKeyImportItem(BaseModel):
class BatchImportRequest(BaseModel):
keys: list[PoolKeyImportItem] = Field(..., max_length=500)
proxy_node_id: str | None = Field(
default=None,
description="导入时绑定到账号的代理节点 ID（可选）",
)
class BatchImportError(BaseModel):

70
src/api/admin/provider_oauth.py
View File

@@ -1368,6 +1368,61 @@ def _parse_kiro_import_input(raw_input: str) -> list[dict[str, Any]]:
返回: 凭据字典列表
"""
def _normalize_item(item: Any) -> dict[str, Any] | None:
"""规范化单条 Kiro 导入项，兼容导出结构。"""
if isinstance(item, str) and item.strip():
return {"refreshToken": item.strip()}
if not isinstance(item, dict):
return None
nested = item.get("auth_config") or item.get("authConfig")
if isinstance(nested, dict):
# 优先使用 auth_config，兼容导出对象形态：
# {"name": "...", "auth_config": {...}, ...}
merged = dict(nested)
# 若顶层也包含关键字段，允许覆盖 nested（便于手工修正）
for key in (
"provider_type",
"providerType",
"auth_method",
"authMethod",
"auth_type",
"authType",
"refresh_token",
"refreshToken",
"expires_at",
"expiresAt",
"profile_arn",
"profileArn",
"region",
"auth_region",
"authRegion",
"api_region",
"apiRegion",
"client_id",
"clientId",
"client_secret",
"clientSecret",
"machine_id",
"machineId",
"kiro_version",
"kiroVersion",
"system_version",
"systemVersion",
"node_version",
"nodeVersion",
"email",
"access_token",
"accessToken",
):
value = item.get(key)
if value is not None and value != "":
merged[key] = value
return merged
return item
raw = raw_input.strip()
if not raw:
return []
@@ -1380,18 +1435,15 @@ def _parse_kiro_import_input(raw_input: str) -> list[dict[str, Any]]:
if isinstance(parsed, list):
result: list[dict[str, Any]] = []
for item in parsed:
if isinstance(item, dict):
result.append(item)
elif isinstance(item, str) and item.strip():
result.append({"refreshToken": item.strip()})
normalized = _normalize_item(item)
if normalized:
result.append(normalized)
return result
if isinstance(parsed, dict):
# 兼容嵌套格式: {"auth_config": {...}} / {"authConfig": {...}}
nested = parsed.get("auth_config") or parsed.get("authConfig")
if isinstance(nested, dict):
return [nested]
return [parsed]
normalized = _normalize_item(parsed)
if normalized:
return [normalized]
except json.JSONDecodeError:
pass

278
src/api/handlers/base/base_handler.py
View File

@@ -145,6 +145,284 @@ class BaseMessageHandler:
return None
return {"perf": self.perf_metrics}
@staticmethod
def _normalize_candidate_status(candidate: dict[str, Any]) -> str:
status = candidate.get("status")
if isinstance(status, str) and status.strip():
return status.strip().lower()
attempt_status = candidate.get("attempt_status")
if isinstance(attempt_status, str) and attempt_status.strip():
return attempt_status.strip().lower()
if candidate.get("skipped"):
return "skipped"
return ""
@staticmethod
def _to_int(value: Any, default: int = 0) -> int:
try:
return int(value)
except Exception:
return default
def _load_request_candidate_keys(self) -> list[Any]:
if not self.request_id:
return []
try:
from src.services.candidate.recorder import CandidateRecorder
return CandidateRecorder(self.db).get_candidate_keys(self.request_id)
except Exception:
return []
def _compact_candidate_key_snapshot(self, item: Any) -> dict[str, Any] | None:
raw: dict[str, Any] | None = None
if isinstance(item, dict):
raw = dict(item)
elif hasattr(item, "to_dict"):
try:
converted = item.to_dict()
if isinstance(converted, dict):
raw = dict(converted)
except Exception:
raw = None
if raw is None:
return None
status = self._normalize_candidate_status(raw)
candidate_index = raw.get("candidate_index", raw.get("index", 0))
retry_index = raw.get("retry_index", 0)
snapshot: dict[str, Any] = {
"candidate_index": self._to_int(candidate_index, 0),
"retry_index": self._to_int(retry_index, 0),
}
passthrough_fields = (
"provider_id",
"provider_name",
"endpoint_id",
"key_id",
"key_name",
"auth_type",
"priority",
"is_cached",
"skip_reason",
"error_type",
"status_code",
"latency_ms",
)
for field in passthrough_fields:
value = raw.get(field)
if value is not None and value != "":
snapshot[field] = value
if status:
snapshot["status"] = status
if raw.get("skipped"):
snapshot["skipped"] = True
snapshot.setdefault("status", "skipped")
if "selected" in raw:
snapshot["selected"] = bool(raw.get("selected"))
error_message = raw.get("error_message")
if isinstance(error_message, str) and error_message:
snapshot["error_message"] = error_message[:240]
return snapshot
def _collect_candidate_snapshots(
self,
*,
candidate_keys: list[Any] | None = None,
fallback_from_request: bool = False,
) -> list[dict[str, Any]]:
source = candidate_keys
if (not source) and fallback_from_request:
source = self._load_request_candidate_keys()
snapshots: list[dict[str, Any]] = []
for item in source or []:
snapshot = self._compact_candidate_key_snapshot(item)
if snapshot:
snapshots.append(snapshot)
snapshots.sort(
key=lambda it: (
self._to_int(it.get("candidate_index"), 0),
self._to_int(it.get("retry_index"), 0),
)
)
return snapshots[:64]
def _build_scheduling_audit(
self,
snapshots: list[dict[str, Any]],
*,
selected_key_id: str | None = None,
) -> dict[str, Any] | None:
if not snapshots:
return None
# "unused" means the candidate was pre-created for audit but never actually attempted.
executed_status_exclude = {"", "available", "pending", "skipped", "unused"}
executed_count = 0
attempts: list[dict[str, Any]] = []
account_map: dict[str, dict[str, Any]] = {}
candidate_indices: set[int] = set()
key_ids: set[str] = set()
for snapshot in snapshots:
status = str(snapshot.get("status", "") or "").lower()
if status in executed_status_exclude:
continue
executed_count += 1
candidate_index = self._to_int(snapshot.get("candidate_index"), 0)
retry_index = self._to_int(snapshot.get("retry_index"), 0)
key_id = snapshot.get("key_id")
key_name = snapshot.get("key_name")
provider_id = snapshot.get("provider_id")
provider_name = snapshot.get("provider_name")
candidate_indices.add(candidate_index)
if isinstance(key_id, str) and key_id:
key_ids.add(key_id)
if len(attempts) < 24:
attempts.append(
{
"candidate_index": candidate_index,
"retry_index": retry_index,
"provider_id": provider_id,
"provider_name": provider_name,
"key_id": key_id,
"key_name": key_name,
"status": status,
"status_code": snapshot.get("status_code"),
"error_type": snapshot.get("error_type"),
}
)
if not isinstance(key_id, str) or not key_id:
continue
account = account_map.get(key_id)
if account is None:
account = {
"key_id": key_id,
"key_name": key_name,
"provider_id": provider_id,
"provider_name": provider_name,
"attempts": 0,
"successes": 0,
"last_status": status,
}
account_map[key_id] = account
account["attempts"] = self._to_int(account.get("attempts"), 0) + 1
if status in {"success", "streaming"}:
account["successes"] = self._to_int(account.get("successes"), 0) + 1
account["last_status"] = status
if executed_count == 0:
return {
"mode": "internal",
"attempted_count": 0,
"account_count": 0,
"retry_occurred": False,
"failover_occurred": False,
"accounts": [],
"attempts": [],
}
selected_key_id_norm = str(selected_key_id) if selected_key_id else None
accounts = list(account_map.values())[:12]
selected_account: dict[str, Any] | None = None
if selected_key_id_norm and selected_key_id_norm in account_map:
selected_account = dict(account_map[selected_key_id_norm])
else:
for account in account_map.values():
if self._to_int(account.get("successes"), 0) > 0:
selected_account = dict(account)
selected_key_id_norm = str(account.get("key_id", ""))
break
if selected_account is not None:
for account in accounts:
if account.get("key_id") == selected_account.get("key_id"):
account["selected"] = True
failover_occurred = executed_count > 1 and (len(candidate_indices) > 1 or len(key_ids) > 1)
return {
"mode": "internal",
"attempted_count": executed_count,
"account_count": len(account_map),
"retry_occurred": executed_count > 1,
"failover_occurred": bool(failover_occurred),
"selected_key_id": selected_key_id_norm,
"selected_account": selected_account,
"accounts": accounts,
"attempts": attempts,
}
def _build_scheduling_metadata(
self,
*,
candidate_keys: list[Any] | None = None,
selected_key_id: str | None = None,
pool_summary: dict[str, Any] | None = None,
fallback_from_request: bool = False,
) -> dict[str, Any]:
snapshots = self._collect_candidate_snapshots(
candidate_keys=candidate_keys,
fallback_from_request=fallback_from_request,
)
metadata: dict[str, Any] = {}
if pool_summary:
metadata["pool_summary"] = pool_summary
if snapshots:
metadata["candidate_keys"] = snapshots
scheduling_audit = self._build_scheduling_audit(
snapshots,
selected_key_id=selected_key_id,
)
if scheduling_audit:
metadata["scheduling_audit"] = scheduling_audit
return metadata
def _merge_scheduling_metadata(
self,
request_metadata: dict[str, Any] | None,
*,
exec_result: Any | None = None,
selected_key_id: str | None = None,
candidate_keys: list[Any] | None = None,
pool_summary: dict[str, Any] | None = None,
fallback_from_request: bool = True,
) -> dict[str, Any] | None:
merged = dict(request_metadata or {})
resolved_candidate_keys = (
candidate_keys
if candidate_keys is not None
else getattr(exec_result, "candidate_keys", None)
)
resolved_key_id = selected_key_id or getattr(exec_result, "key_id", None)
resolved_pool_summary = (
pool_summary if pool_summary is not None else getattr(exec_result, "pool_summary", None)
)
merged.update(
self._build_scheduling_metadata(
candidate_keys=resolved_candidate_keys,
selected_key_id=resolved_key_id,
pool_summary=resolved_pool_summary,
fallback_from_request=fallback_from_request,
)
)
return merged or None
def _resolve_capability_requirements(
self,
model_name: str,

15
src/api/handlers/base/chat_handler_base.py
View File

@@ -589,6 +589,21 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
ctx.key_id = key_id
if getattr(exec_result, "pool_summary", None):
ctx.pool_summary = exec_result.pool_summary
scheduling_metadata = (
self._merge_scheduling_metadata(
{},
exec_result=exec_result,
selected_key_id=key_id,
fallback_from_request=False,
)
or {}
)
candidate_keys = scheduling_metadata.get("candidate_keys")
if isinstance(candidate_keys, list):
ctx.candidate_keys = candidate_keys
scheduling_audit = scheduling_metadata.get("scheduling_audit")
if isinstance(scheduling_audit, dict):
ctx.scheduling_audit = scheduling_audit
# 同步整流状态（如果请求体被整流过）
ctx.rectified = request_body_ref.get("_rectified", False)

35
src/api/handlers/base/chat_sync_executor.py
View File

@@ -216,8 +216,11 @@ class ChatSyncExecutor:
request_metadata = handler._build_request_metadata() or {}
if ctx.sync_proxy_info:
request_metadata["proxy"] = ctx.sync_proxy_info
if getattr(exec_result, "pool_summary", None):
request_metadata["pool_summary"] = exec_result.pool_summary
request_metadata = handler._merge_scheduling_metadata(
request_metadata,
exec_result=exec_result,
selected_key_id=ctx.key_id,
)
total_cost = await handler.telemetry.record_success( # noqa: F841
provider=ctx.provider_name,
model=model,
@@ -251,7 +254,7 @@ class ChatSyncExecutor:
provider_api_key_id=ctx.key_id,
# 模型映射信息
target_model=ctx.mapped_model_result,
request_metadata=request_metadata or None,
request_metadata=request_metadata,
)
logger.debug(f"{handler.FORMAT_ID} 非流式响应完成")
@@ -273,6 +276,12 @@ class ChatSyncExecutor:
request_metadata = handler._build_request_metadata() or {}
if ctx.sync_proxy_info:
request_metadata["proxy"] = ctx.sync_proxy_info
request_metadata = handler._merge_scheduling_metadata(
request_metadata,
selected_key_id=ctx.key_id,
pool_summary=ctx.pool_summary,
fallback_from_request=True,
)
await handler.telemetry.record_failure(
provider=ctx.provider_name or "unknown",
model=model,
@@ -286,7 +295,7 @@ class ChatSyncExecutor:
provider_id=ctx.provider_id,
provider_endpoint_id=ctx.endpoint_id,
provider_api_key_id=ctx.key_id,
request_metadata=request_metadata or None,
request_metadata=request_metadata,
)
client_format = (ctx.client_api_format_for_error or "").upper()
provider_format = (ctx.provider_api_format_for_error or client_format).upper()
@@ -308,6 +317,12 @@ class ChatSyncExecutor:
request_metadata = handler._build_request_metadata() or {}
if ctx.sync_proxy_info:
request_metadata["proxy"] = ctx.sync_proxy_info
request_metadata = handler._merge_scheduling_metadata(
request_metadata,
selected_key_id=ctx.key_id,
pool_summary=ctx.pool_summary,
fallback_from_request=True,
)
client_format = (ctx.client_api_format_for_error or "").upper()
provider_format = (ctx.provider_api_format_for_error or client_format).upper()
payload = _build_error_json_payload(
@@ -372,6 +387,12 @@ class ChatSyncExecutor:
request_metadata = handler._build_request_metadata() or {}
if ctx.sync_proxy_info:
request_metadata["proxy"] = ctx.sync_proxy_info
request_metadata = handler._merge_scheduling_metadata(
request_metadata,
selected_key_id=ctx.key_id,
pool_summary=ctx.pool_summary,
fallback_from_request=True,
)
await handler.telemetry.record_failure(
provider=ctx.provider_name or "unknown",
model=model,
@@ -750,6 +771,12 @@ class ChatSyncExecutor:
stream_fail_metadata: dict[str, Any] | None = None
if ctx.proxy_info:
stream_fail_metadata = {"proxy": ctx.proxy_info}
stream_fail_metadata = handler._merge_scheduling_metadata(
stream_fail_metadata,
selected_key_id=ctx.key_id,
pool_summary=ctx.pool_summary,
fallback_from_request=True,
)
await handler.telemetry.record_failure(
provider=ctx.provider_name or "unknown",

32
src/api/handlers/base/cli_monitor_mixin.py
View File

@@ -299,7 +299,13 @@ class CliMonitorMixin:
if ctx.is_client_disconnected():
# 客户端取消：记录为 cancelled（不算系统失败）
request_metadata = {"perf": ctx.perf_metrics} if ctx.perf_metrics else None
request_metadata = self._merge_scheduling_metadata(
{"perf": ctx.perf_metrics} if ctx.perf_metrics else None,
selected_key_id=ctx.key_id,
candidate_keys=ctx.candidate_keys,
pool_summary=ctx.pool_summary,
fallback_from_request=True,
)
await bg_telemetry.record_cancelled(
provider=ctx.provider_name or "unknown",
model=ctx.model,
@@ -334,7 +340,13 @@ class CliMonitorMixin:
)
else:
# 服务端/上游异常：记录为失败
request_metadata = {"perf": ctx.perf_metrics} if ctx.perf_metrics else None
request_metadata = self._merge_scheduling_metadata(
{"perf": ctx.perf_metrics} if ctx.perf_metrics else None,
selected_key_id=ctx.key_id,
candidate_keys=ctx.candidate_keys,
pool_summary=ctx.pool_summary,
fallback_from_request=True,
)
await bg_telemetry.record_failure(
provider=ctx.provider_name or "unknown",
model=ctx.model,
@@ -412,7 +424,13 @@ class CliMonitorMixin:
f"provider={ctx.provider_name}, model={ctx.model}, "
f"in={ctx.input_tokens}, out={ctx.output_tokens}"
)
request_metadata = {"perf": ctx.perf_metrics} if ctx.perf_metrics else None
request_metadata = self._merge_scheduling_metadata(
{"perf": ctx.perf_metrics} if ctx.perf_metrics else None,
selected_key_id=ctx.key_id,
candidate_keys=ctx.candidate_keys,
pool_summary=ctx.pool_summary,
fallback_from_request=True,
)
total_cost = await bg_telemetry.record_success(
provider=ctx.provider_name,
model=ctx.model,
@@ -570,7 +588,13 @@ class CliMonitorMixin:
# 失败时返回给客户端的是 JSON 错误响应
client_response_headers = {"content-type": "application/json"}
request_metadata = {"perf": ctx.perf_metrics} if ctx.perf_metrics else None
request_metadata = self._merge_scheduling_metadata(
{"perf": ctx.perf_metrics} if ctx.perf_metrics else None,
selected_key_id=ctx.key_id,
candidate_keys=ctx.candidate_keys,
pool_summary=ctx.pool_summary,
fallback_from_request=True,
)
await self.telemetry.record_failure(
provider=ctx.provider_name or "unknown",
model=ctx.model,

11
src/api/handlers/base/cli_protocol.py
View File

@@ -87,6 +87,17 @@ class CliHandlerProtocol(Protocol):
http_request: Any | None = ...,
) -> dict[str, Any] | None: ...
def _merge_scheduling_metadata(
self,
request_metadata: dict[str, Any] | None,
*,
exec_result: Any | None = ...,
selected_key_id: str | None = ...,
candidate_keys: list[Any] | None = ...,
pool_summary: dict[str, Any] | None = ...,
fallback_from_request: bool = ...,
) -> dict[str, Any] | None: ...
def _resolve_capability_requirements(
self,
model_name: str,

15
src/api/handlers/base/cli_stream_mixin.py
View File

@@ -198,6 +198,21 @@ class CliStreamMixin:
ctx.key_id = key_id
if getattr(exec_result, "pool_summary", None):
ctx.pool_summary = exec_result.pool_summary
scheduling_metadata = (
self._merge_scheduling_metadata(
{},
exec_result=exec_result,
selected_key_id=key_id,
fallback_from_request=False,
)
or {}
)
candidate_keys = scheduling_metadata.get("candidate_keys")
if isinstance(candidate_keys, list):
ctx.candidate_keys = candidate_keys
scheduling_audit = scheduling_metadata.get("scheduling_audit")
if isinstance(scheduling_audit, dict):
ctx.scheduling_audit = scheduling_audit
# 同步整流状态（如果请求体被整流过）
ctx.rectified = request_body_ref.get("_rectified", False)

26
src/api/handlers/base/cli_sync_mixin.py
View File

@@ -101,6 +101,7 @@ class CliSyncMixin:
provider_id = None # Provider ID（用于失败记录）
endpoint_id = None # Endpoint ID（用于失败记录）
key_id = None # Key ID（用于失败记录）
exec_result = None
mapped_model_result = None # 映射后的目标模型名（用于 Usage 记录）
response_metadata_result: dict[str, Any] = {} # Provider 响应元数据
needs_conversion = False # 是否需要格式转换（由 candidate 决定）
@@ -558,8 +559,11 @@ class CliSyncMixin:
request_metadata = self._build_request_metadata() or {}
if sync_proxy_info:
request_metadata["proxy"] = sync_proxy_info
if getattr(exec_result, "pool_summary", None):
request_metadata["pool_summary"] = exec_result.pool_summary
request_metadata = self._merge_scheduling_metadata(
request_metadata,
exec_result=exec_result,
selected_key_id=key_id,
)
total_cost = await self.telemetry.record_success(
provider=provider_name,
model=model,
@@ -592,7 +596,7 @@ class CliSyncMixin:
target_model=mapped_model_result,
# Provider 响应元数据（如 Gemini 的 modelVersion）
response_metadata=response_metadata_result if response_metadata_result else None,
request_metadata=request_metadata or None,
request_metadata=request_metadata,
)
logger.info("{} 非流式响应处理完成", self.FORMAT_ID)
@@ -607,6 +611,12 @@ class CliSyncMixin:
request_metadata = self._build_request_metadata() or {}
if sync_proxy_info:
request_metadata["proxy"] = sync_proxy_info
request_metadata = self._merge_scheduling_metadata(
request_metadata,
selected_key_id=key_id,
pool_summary=getattr(exec_result, "pool_summary", None),
fallback_from_request=True,
)
await self.telemetry.record_failure(
provider=provider_name or "unknown",
model=model,
@@ -620,7 +630,7 @@ class CliSyncMixin:
api_format=api_format,
api_family=self.api_family,
endpoint_kind=self.endpoint_kind,
request_metadata=request_metadata or None,
request_metadata=request_metadata,
)
raise
@@ -645,6 +655,12 @@ class CliSyncMixin:
request_metadata = self._build_request_metadata() or {}
if sync_proxy_info:
request_metadata["proxy"] = sync_proxy_info
request_metadata = self._merge_scheduling_metadata(
request_metadata,
selected_key_id=key_id,
pool_summary=getattr(exec_result, "pool_summary", None),
fallback_from_request=True,
)
await self.telemetry.record_failure(
provider=provider_name or "unknown",
model=model,
@@ -667,7 +683,7 @@ class CliSyncMixin:
has_format_conversion=is_format_converted(provider_api_format, str(api_format)),
# 模型映射信息
target_model=mapped_model_result,
request_metadata=request_metadata or None,
request_metadata=request_metadata,
)
raise

7
src/api/handlers/base/stream_context.py
View File

@@ -139,6 +139,10 @@ class StreamContext:
# 号池调度摘要（来自 ExecutionResult.pool_summary）
pool_summary: dict[str, Any] | None = None
# 候选轨迹（来自 ExecutionResult.candidate_keys，写入 usage metadata）
candidate_keys: list[dict[str, Any]] = field(default_factory=list)
# 内部调度审计摘要（重试/故障转移/账号使用轨迹）
scheduling_audit: dict[str, Any] | None = None
# 流式格式转换状态（跨 chunk 追踪）
stream_conversion_state: StreamState | None = None
@@ -175,6 +179,9 @@ class StreamContext:
self.final_usage = None
self.final_response = None
self.proxy_info = None
self.pool_summary = None
self.candidate_keys = []
self.scheduling_audit = None
self.stream_conversion_state = None
self.stream_conversion_event_count = 0
self.needs_conversion = False

16
src/api/handlers/base/stream_telemetry.py
View File

@@ -210,6 +210,10 @@ class StreamTelemetryRecorder:
metadata["proxy"] = ctx.proxy_info
if ctx.pool_summary:
metadata["pool_summary"] = ctx.pool_summary
if ctx.candidate_keys:
metadata["candidate_keys"] = ctx.candidate_keys
if ctx.scheduling_audit:
metadata["scheduling_audit"] = ctx.scheduling_audit
await writer.record_success(
provider=ctx.provider_name or "unknown",
@@ -268,6 +272,12 @@ class StreamTelemetryRecorder:
metadata["perf"] = ctx.perf_metrics
if ctx.proxy_info:
metadata["proxy"] = ctx.proxy_info
if ctx.pool_summary:
metadata["pool_summary"] = ctx.pool_summary
if ctx.candidate_keys:
metadata["candidate_keys"] = ctx.candidate_keys
if ctx.scheduling_audit:
metadata["scheduling_audit"] = ctx.scheduling_audit
await writer.record_failure(
provider=ctx.provider_name or "unknown",
@@ -327,6 +337,12 @@ class StreamTelemetryRecorder:
metadata["perf"] = ctx.perf_metrics
if ctx.proxy_info:
metadata["proxy"] = ctx.proxy_info
if ctx.pool_summary:
metadata["pool_summary"] = ctx.pool_summary
if ctx.candidate_keys:
metadata["candidate_keys"] = ctx.candidate_keys
if ctx.scheduling_audit:
metadata["scheduling_audit"] = ctx.scheduling_audit
await writer.record_cancelled(
provider=ctx.provider_name or "unknown",

1
src/models/database.py
View File

@@ -306,6 +306,7 @@ class Usage(Base):
Index("idx_usage_provider_model_created", "provider_name", "model", "created_at"),
Index("idx_usage_provider_created", "provider_name", "created_at"),
Index("idx_usage_model_created", "model", "created_at"),
Index("idx_usage_provider_key", "provider_id", "provider_api_key_id"),
)
id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()), index=True)

7
src/services/orchestration/candidate_resolver.py
View File

@@ -214,6 +214,11 @@ class CandidateResolver:
provider = candidate.provider
endpoint = candidate.endpoint
key = candidate.key
pool_extra = (
getattr(candidate, "_pool_extra_data", None)
if isinstance(getattr(candidate, "_pool_extra_data", None), dict)
else {}
)
if candidate.is_skipped:
record_id = str(uuid.uuid4())
@@ -235,6 +240,7 @@ class CandidateResolver:
"needs_conversion": candidate.needs_conversion,
"provider_api_format": candidate.provider_api_format or None,
"mapping_matched_model": candidate.mapping_matched_model or None,
**pool_extra,
},
"required_capabilities": active_capabilities,
"created_at": datetime.now(timezone.utc),
@@ -269,6 +275,7 @@ class CandidateResolver:
"needs_conversion": candidate.needs_conversion,
"provider_api_format": candidate.provider_api_format or None,
"mapping_matched_model": candidate.mapping_matched_model or None,
**pool_extra,
},
"required_capabilities": active_capabilities,
"created_at": datetime.now(timezone.utc),

48
src/services/provider/adapters/kiro/models/credentials.py
View File

@@ -36,6 +36,28 @@ def _nonempty(s: str | None) -> str | None:
return None
def _normalize_auth_method(value: str | None) -> str:
method = (value or "").strip().lower()
if not method:
return "social"
# 历史/别名兼容：统一映射到 idc
if method in {
"idc",
"builder-id",
"builder_id",
"builderid",
"identity-center",
"identity_center",
"identitycenter",
"iam",
"device",
"device_authorization",
"device-auth",
}:
return "idc"
return method
def _parse_iso_to_epoch_seconds(value: object) -> int | None:
if not isinstance(value, str) or not value.strip():
return None
@@ -111,6 +133,11 @@ class KiroAuthConfig:
- 包含 clientId + clientSecret -> IdC
- 仅含 refreshToken -> Social
"""
explicit_method = _get_str(raw, "auth_method", "authMethod", "auth_type", "authType")
normalized_explicit = _normalize_auth_method(explicit_method)
if normalized_explicit != "social":
return normalized_explicit
client_id = raw.get("client_id") or raw.get("clientId")
client_secret = raw.get("client_secret") or raw.get("clientSecret")
@@ -136,7 +163,12 @@ class KiroAuthConfig:
return False, "refreshToken 不完整（含有 ...），请导出完整的 Token"
# IdC 类型需要 clientId 和 clientSecret
auth_method = KiroAuthConfig.infer_auth_method(raw)
explicit_method = _get_str(raw, "auth_method", "authMethod", "auth_type", "authType")
auth_method = (
_normalize_auth_method(explicit_method)
if explicit_method
else KiroAuthConfig.infer_auth_method(raw)
)
if auth_method == "idc":
client_id = raw.get("client_id") or raw.get("clientId")
client_secret = raw.get("client_secret") or raw.get("clientSecret")
@@ -155,8 +187,12 @@ class KiroAuthConfig:
provider_type = _get_str(raw, "provider_type", "providerType") or "kiro"
# 自动推断 auth_method（如果未显式指定）
explicit_method = _get_str(raw, "auth_method", "authMethod")
auth_method = explicit_method.lower() if explicit_method else cls.infer_auth_method(raw)
explicit_method = _get_str(raw, "auth_method", "authMethod", "auth_type", "authType")
auth_method = (
_normalize_auth_method(explicit_method)
if explicit_method
else cls.infer_auth_method(raw)
)
refresh_token = (_get_str(raw, "refresh_token", "refreshToken") or "").strip()
@@ -168,7 +204,7 @@ class KiroAuthConfig:
cfg = cls(
provider_type=provider_type,
auth_method=(auth_method or "social").lower(),
auth_method=_normalize_auth_method(auth_method),
refresh_token=refresh_token,
expires_at=int(expires_at),
profile_arn=_get_str(raw, "profile_arn", "profileArn"),
@@ -185,10 +221,6 @@ class KiroAuthConfig:
access_token=_get_str(raw, "access_token", "accessToken"),
)
# Normalize auth_method aliases.
if cfg.auth_method in {"builder-id", "builder_id", "iam"}:
cfg.auth_method = "idc"
return cfg
def to_dict(self) -> dict[str, Any]:

44
src/services/provider/pool/redis_ops.py
View File

@@ -440,6 +440,50 @@ async def get_key_sticky_count(provider_id: str, key_id: str) -> int:
return 0
async def batch_get_key_sticky_counts(
provider_id: str,
key_ids: list[str],
) -> dict[str, int]:
"""Count sticky sessions for multiple keys in a single scan (admin only)."""
if not key_ids:
return {}
redis = await _get_redis()
if redis is None:
return {kid: 0 for kid in key_ids}
target_ids = set(key_ids)
counts: dict[str, int] = {kid: 0 for kid in key_ids}
try:
pattern = f"{PREFIX}:{provider_id}:sticky:*"
batch: list[bytes | str] = []
async for key in redis.scan_iter(match=pattern, count=200):
batch.append(key)
if len(batch) >= 200:
vals = await redis.mget(batch)
for val in vals:
if not val:
continue
bound_id = val.decode() if isinstance(val, bytes) else str(val)
if bound_id in target_ids:
counts[bound_id] = counts.get(bound_id, 0) + 1
batch.clear()
if batch:
vals = await redis.mget(batch)
for val in vals:
if not val:
continue
bound_id = val.decode() if isinstance(val, bytes) else str(val)
if bound_id in target_ids:
counts[bound_id] = counts.get(bound_id, 0) + 1
return counts
except Exception:
return {kid: 0 for kid in key_ids}
async def get_cooldown_ttl(provider_id: str, key_id: str) -> int | None:
"""Get remaining cooldown TTL in seconds. None = no cooldown."""
redis = await _get_redis()

373
src/services/provider/pool/scheduling_dimensions.py Normal file
View File

@@ -0,0 +1,373 @@
"""Pool scheduling dimension registry and evaluation helpers.
This module keeps pool scheduling scoring isolated from API layer code.
Callers build a :class:`PoolSchedulingSnapshot` and evaluate it against
registered dimensions to obtain a normalized summary.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Protocol
PoolDimensionStatus = str # ok / degraded / blocked
@dataclass(frozen=True, slots=True)
class PoolSchedulingSnapshot:
"""Point-in-time scheduling inputs for one key."""
is_active: bool
cooldown_reason: str | None
cooldown_ttl_seconds: int | None
circuit_breaker_open: bool
cost_window_usage: int
cost_limit: int | None
cost_soft_threshold_percent: int = 80
health_score: float = 1.0
@dataclass(frozen=True, slots=True)
class PoolSchedulingDimensionResult:
"""Evaluation output for one scheduling dimension."""
code: str
label: str
status: PoolDimensionStatus = "ok"
blocking: bool = False
source: str = "pool"
weight: int = 1
score: float = 1.0
detail: str | None = None
ttl_seconds: int | None = None
@dataclass(frozen=True, slots=True)
class PoolSchedulingSummary:
"""Merged scheduling state across all dimensions."""
status: str # available / degraded / blocked
reason: str
label: str
score: float
candidate_eligible: bool
blocked_count: int
degraded_count: int
class PoolSchedulingDimension(Protocol):
"""Dimension evaluator protocol."""
code: str
label: str
source: str
weight: int
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
"""Evaluate one dimension from snapshot."""
@dataclass(frozen=True, slots=True)
class _ManualEnableDimension:
code: str = "manual_disabled"
label: str = "已禁用"
source: str = "manual"
weight: int = 8
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
if snapshot.is_active:
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="ok",
score=1.0,
)
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="blocked",
blocking=True,
score=0.0,
detail="账号被手动禁用",
)
@dataclass(frozen=True, slots=True)
class _CooldownDimension:
code: str = "cooldown"
label: str = "冷却中"
source: str = "pool"
weight: int = 7
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
if not snapshot.cooldown_reason:
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="ok",
score=1.0,
)
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="blocked",
blocking=True,
score=0.0,
detail=snapshot.cooldown_reason,
ttl_seconds=snapshot.cooldown_ttl_seconds,
)
@dataclass(frozen=True, slots=True)
class _CircuitBreakerDimension:
code: str = "circuit_open"
label: str = "熔断中"
source: str = "health"
weight: int = 6
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
if not snapshot.circuit_breaker_open:
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="ok",
score=1.0,
)
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="blocked",
blocking=True,
score=0.0,
)
@dataclass(frozen=True, slots=True)
class _CostDimension:
code: str = "cost"
label: str = "成本"
source: str = "pool"
weight: int = 5
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
limit = snapshot.cost_limit
usage = max(snapshot.cost_window_usage, 0)
if limit is None or limit <= 0:
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="ok",
score=1.0,
detail=f"{usage}/-",
)
ratio = usage / limit
detail = f"{usage}/{limit}"
if ratio >= 1.0:
return PoolSchedulingDimensionResult(
code="cost_exhausted",
label="成本超限",
source=self.source,
weight=self.weight,
status="blocked",
blocking=True,
score=0.0,
detail=detail,
)
soft_threshold = max(1, min(snapshot.cost_soft_threshold_percent, 100))
if ratio * 100 >= soft_threshold:
return PoolSchedulingDimensionResult(
code="cost_soft",
label="成本接近上限",
source=self.source,
weight=self.weight,
status="degraded",
score=0.45,
detail=detail,
)
if ratio >= 0.6:
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="degraded",
score=0.72,
detail=detail,
)
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="ok",
score=1.0,
detail=detail,
)
@dataclass(frozen=True, slots=True)
class _HealthDimension:
code: str = "health"
label: str = "健康度"
source: str = "health"
weight: int = 4
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
score = max(0.0, min(snapshot.health_score, 1.0))
detail = f"{score:.2f}"
if score < 0.5:
return PoolSchedulingDimensionResult(
code="health_low",
label="健康度过低",
source=self.source,
weight=self.weight,
status="degraded",
score=0.3,
detail=detail,
)
if score < 0.8:
return PoolSchedulingDimensionResult(
code="health_degraded",
label="健康度下降",
source=self.source,
weight=self.weight,
status="degraded",
score=0.65,
detail=detail,
)
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="ok",
score=1.0,
detail=detail,
)
_POOL_DIMENSION_REGISTRY: dict[str, PoolSchedulingDimension] = {}
_POOL_DIMENSION_ORDER: list[str] = []
def register_pool_scheduling_dimension(name: str, dimension: PoolSchedulingDimension) -> None:
"""Register a dimension evaluator by name."""
normalized = name.strip()
if not normalized:
return
if normalized not in _POOL_DIMENSION_ORDER:
_POOL_DIMENSION_ORDER.append(normalized)
_POOL_DIMENSION_REGISTRY[normalized] = dimension
def get_pool_scheduling_dimension(name: str) -> PoolSchedulingDimension | None:
"""Fetch a registered dimension evaluator."""
return _POOL_DIMENSION_REGISTRY.get(name.strip())
def list_pool_scheduling_dimensions() -> tuple[str, ...]:
"""List registered dimension names in evaluation order."""
return tuple(_POOL_DIMENSION_ORDER)
def evaluate_pool_scheduling_dimensions(
snapshot: PoolSchedulingSnapshot,
*,
dimension_names: tuple[str, ...] | None = None,
) -> list[PoolSchedulingDimensionResult]:
"""Evaluate snapshot across all registered dimensions."""
names = dimension_names or list_pool_scheduling_dimensions()
results: list[PoolSchedulingDimensionResult] = []
for name in names:
dimension = get_pool_scheduling_dimension(name)
if dimension is None:
continue
results.append(dimension.evaluate(snapshot))
return results
def summarize_pool_scheduling_dimensions(
dimensions: list[PoolSchedulingDimensionResult],
) -> PoolSchedulingSummary:
"""Summarize dimension outputs into a unified scheduling state."""
if not dimensions:
return PoolSchedulingSummary(
status="available",
reason="available",
label="可用",
score=100.0,
candidate_eligible=True,
blocked_count=0,
degraded_count=0,
)
blocked = [item for item in dimensions if item.status == "blocked" or item.blocking]
degraded = [item for item in dimensions if item.status == "degraded"]
total_weight = sum(max(item.weight, 1) for item in dimensions)
weighted_score = sum(
max(item.weight, 1) * max(min(item.score, 1.0), 0.0) for item in dimensions
) / max(total_weight, 1)
if blocked:
primary = blocked[0]
return PoolSchedulingSummary(
status="blocked",
reason=primary.code,
label=primary.label,
score=round(weighted_score * 100, 1),
candidate_eligible=False,
blocked_count=len(blocked),
degraded_count=len(degraded),
)
if degraded:
primary = degraded[0]
return PoolSchedulingSummary(
status="degraded",
reason=primary.code,
label=primary.label,
score=round(weighted_score * 100, 1),
candidate_eligible=True,
blocked_count=0,
degraded_count=len(degraded),
)
return PoolSchedulingSummary(
status="available",
reason="available",
label="可用",
score=round(weighted_score * 100, 1),
candidate_eligible=True,
blocked_count=0,
degraded_count=0,
)
def _register_default_dimensions() -> None:
register_pool_scheduling_dimension("manual", _ManualEnableDimension())
register_pool_scheduling_dimension("cooldown", _CooldownDimension())
register_pool_scheduling_dimension("circuit", _CircuitBreakerDimension())
register_pool_scheduling_dimension("cost", _CostDimension())
register_pool_scheduling_dimension("health", _HealthDimension())
_register_default_dimensions()

20
src/services/provider/pool/trace.py
View File

@@ -63,7 +63,12 @@ class PoolSchedulingTrace:
session_uuid: str | None = None
candidate_traces: dict[str, PoolCandidateTrace] = field(default_factory=dict)
def build_summary(self, success_key_id: str | None = None) -> dict[str, Any]:
def build_summary(
self,
success_key_id: str | None = None,
*,
attempted_key_ids: set[str] | None = None,
) -> dict[str, Any]:
"""Build compact dict for ``Usage.request_metadata["pool_summary"]``."""
skipped_cooldown = 0
skipped_cost = 0
@@ -74,8 +79,17 @@ class PoolSchedulingTrace:
skipped_cooldown += 1
elif t.skip_type == "cost_exhausted":
skipped_cost += 1
else:
attempted += 1
if attempted_key_ids is None:
# Backward-compatible behavior: count all schedulable keys.
attempted = sum(1 for t in self.candidate_traces.values() if not t.skipped)
else:
# Preferred behavior: count only keys that were actually executed.
attempted = sum(
1
for kid in attempted_key_ids
if kid in self.candidate_traces and not self.candidate_traces[kid].skipped
)
success_reason: str | None = None
if success_key_id and success_key_id in self.candidate_traces:

35
src/services/provider_keys/quota_cooldown.py Normal file
View File

@@ -0,0 +1,35 @@
"""配额冷却判定工具。"""
from __future__ import annotations
from typing import Any
from src.core.logger import logger
from src.services.scheduling.quota_skipper import is_key_quota_exhausted
def resolve_effective_cooldown_reason(
*,
provider_type: str | None,
key: Any,
redis_reason: str | None,
) -> str | None:
"""返回 Key 的有效冷却原因。
规则：
- Redis 冷却存在时，优先返回 Redis 原因（429/403/quota_exhausted 等）。
- Redis 冷却不存在时，回退到 upstream_metadata 配额判断：
若账号级配额耗尽（Codex/Kiro），返回 ``quota_exhausted``。
"""
if redis_reason:
return redis_reason
try:
exhausted, _ = is_key_quota_exhausted(provider_type, key, model_name="")
except Exception:
logger.opt(exception=True).debug(
"quota_cooldown: is_key_quota_exhausted failed for key={}",
getattr(key, "id", "?"),
)
return None
return "quota_exhausted" if exhausted else None

56
src/services/scheduling/candidate_builder.py
View File

@@ -576,32 +576,21 @@ class CandidateBuilder:
if not active_keys:
continue
# --- Pool branch: select a single key internally ------
# Pool provider should still expose all key candidates here.
# Runtime pool scheduling/failover is handled later by TaskService._apply_pool_reorder.
use_random = all((key.cache_ttl_minutes or 0) == 0 for key in active_keys)
if pool_cfg is not None:
selected_key = await self._pool_select_key(
db, provider, pool_cfg, active_keys, request_body
)
if selected_key is None:
logger.debug(
"Pool[{}]: no schedulable key for endpoint {}",
str(provider.id)[:8],
endpoint_format_str,
)
continue
keys_to_check: list[ProviderAPIKey] = [selected_key]
else:
# --- Normal branch: check all keys ----
use_random = all((key.cache_ttl_minutes or 0) == 0 for key in active_keys)
if use_random and len(active_keys) > 1:
logger.debug(
" Provider {} 启用 Key 轮换模式 (endpoint_format={}, {} keys)",
provider.name,
endpoint_format_str,
len(active_keys),
)
keys_to_check = self._sorter.shuffle_keys_by_internal_priority(
active_keys, affinity_key, use_random
use_random = False
elif use_random and len(active_keys) > 1:
logger.debug(
" Provider {} 启用 Key 轮换模式 (endpoint_format={}, {} keys)",
provider.name,
endpoint_format_str,
len(active_keys),
)
keys_to_check = self._sorter.shuffle_keys_by_internal_priority(
active_keys, affinity_key, use_random
)
for key in keys_to_check:
# Key 级别检查（健康度/熔断按 provider_format bucket）
@@ -660,22 +649,3 @@ class CandidateBuilder:
candidates = candidates[:max_candidates]
return candidates
async def _pool_select_key(
self,
db: Session,
provider: Provider,
pool_cfg: "PoolConfig",
active_keys: list[ProviderAPIKey],
request_body: dict | None,
) -> ProviderAPIKey | None:
"""Select a single key via pool scheduling (sticky -> cooldown/cost -> LRU)."""
from src.services.provider.pool.hooks import get_pool_hook
from src.services.provider.pool.manager import PoolManager
provider_type = str(getattr(provider, "provider_type", "") or "")
hook = get_pool_hook(provider_type)
session_uuid = hook.extract_session_uuid(request_body) if hook and request_body else None
mgr = PoolManager(str(provider.id), pool_cfg)
release_db_connection_before_await(db)
return await mgr.select_key(session_uuid, active_keys)

25
src/services/task/service.py
View File

@@ -119,6 +119,7 @@ class TaskService:
allow_format_conversion=allow_format_conversion,
capability_requirements=capability_requirements,
max_candidates=max_candidates,
request_body=request_body,
)
candidate_keys = []
@@ -598,8 +599,22 @@ class TaskService:
# Build pool scheduling summary from traces collected during reorder.
if pool_traces and result.key_id:
try:
attempted_key_ids: set[str] = set()
for ck in result.candidate_keys or []:
status = str(getattr(ck, "status", "") or "").strip().lower()
if status in {"", "available", "pending", "skipped", "unused"}:
continue
kid = getattr(ck, "key_id", None)
if isinstance(kid, str) and kid:
attempted_key_ids.add(kid)
if not attempted_key_ids:
attempted_key_ids.add(str(result.key_id))
for pt in pool_traces:
summary = pt.build_summary(result.key_id)
summary = pt.build_summary(
result.key_id,
attempted_key_ids=attempted_key_ids,
)
if summary:
result.pool_summary = summary
break
@@ -1204,6 +1219,7 @@ class TaskService:
allow_format_conversion: bool = False,
capability_requirements: dict[str, bool] | None = None,
max_candidates: int | None = None,
request_body: dict[str, Any] | None = None,
) -> Any:
"""
Unified ASYNC submit entrypoint (Phase 3.2).
@@ -1300,6 +1316,7 @@ class TaskService:
request_id=request_id,
is_stream=False,
capability_requirements=capability_requirements,
request_body=request_body,
)
if not candidates:
@@ -1309,6 +1326,12 @@ class TaskService:
last_status_code=None,
)
# Account Pool: keep internal key failover order/skip behavior
# consistent with the SYNC path.
candidates, _pool_traces = await self._apply_pool_reorder(
candidates, request_body=request_body
)
if max_candidates is not None and max_candidates > 0:
candidates = candidates[:max_candidates]

1
src/services/usage/_recording_helpers.py
View File

@@ -39,6 +39,7 @@ METADATA_KEEP_KEYS: frozenset[str] = frozenset(
"billing_updated_at",
"perf",
"pool_summary",
"scheduling_audit",
"_metadata_truncated",
}
)