feat(test,quota,failover): 模型并发测试、统一配额读取器与故障转移取消支持

- 新增 QuotaReader 抽象层,统一 Codex/Kiro/Antigravity 配额解析逻辑,
  替换 pool/routes.py 中分散的配额构建函数
- 模型测试支持并发执行多候选,前端新增 useModelTest composable 统一
  ModelsTab 和 ModelMappingTab 的测试逻辑
- ModelTestDialog 增加结果概览摘要、超长结果折叠、端点列和新状态支持,
  删除已合并的 TestResultDialog
- FailoverEngine 新增客户端断开检测,支持取消剩余候选并标记记录
- 刷新配额改为分批执行,直连测试候选按可用性排序
- 修复 error 判断从 "error" in dict 改为 dict.get("error") 避免误判
This commit is contained in:
fawney19
2026-03-07 02:16:03 +08:00
parent 1f3693d3a2
commit fb1aeb789a
22 changed files with 2145 additions and 987 deletions

View File

@@ -37,6 +37,7 @@ from src.services.provider.pool.scheduling_dimensions import (
evaluate_pool_scheduling_dimensions,
summarize_pool_scheduling_dimensions,
)
from src.services.provider_keys.quota_reader import get_quota_reader
from .schemas import (
BatchActionRequest,
@@ -197,180 +198,12 @@ def _is_known_banned_key(key: ProviderAPIKey, provider_type: str) -> bool:
return state.blocked
def _format_percent(value: float) -> str:
clamped = max(0.0, min(value, 100.0))
return f"{clamped:.1f}%"
def _format_quota_value(value: float) -> str:
rounded = round(value)
if abs(value - rounded) < 1e-6:
return str(rounded)
return f"{value:.1f}"
def _format_reset_after(seconds_raw: Any) -> str | None:
seconds = _to_float(seconds_raw)
if seconds is None:
return None
total_seconds = int(seconds)
if total_seconds <= 0:
return "已重置"
days = total_seconds // 86400
hours = (total_seconds % 86400) // 3600
minutes = (total_seconds % 3600) // 60
if days > 0:
return f"{days}{hours}小时后重置"
if hours > 0:
return f"{hours}小时{minutes}分钟后重置"
if minutes > 0:
return f"{minutes}分钟后重置"
return "即将重置"
def _build_codex_account_quota(upstream_metadata: dict[str, Any]) -> str | None:
codex = upstream_metadata.get("codex")
if not isinstance(codex, dict):
return None
parts: list[str] = []
primary_used = _to_float(codex.get("primary_used_percent"))
if primary_used is not None:
part = f"周剩余 {_format_percent(100.0 - primary_used)}"
reset_text = _format_reset_after(codex.get("primary_reset_seconds"))
if reset_text:
part = f"{part} ({reset_text})"
parts.append(part)
secondary_used = _to_float(codex.get("secondary_used_percent"))
if secondary_used is not None:
part = f"5H剩余 {_format_percent(100.0 - secondary_used)}"
reset_text = _format_reset_after(codex.get("secondary_reset_seconds"))
if reset_text:
part = f"{part} ({reset_text})"
parts.append(part)
if parts:
return " | ".join(parts)
has_credits = codex.get("has_credits")
credits_balance = _to_float(codex.get("credits_balance"))
if has_credits is True and credits_balance is not None:
return f"积分 {credits_balance:.2f}"
if has_credits is True:
return "有积分"
return None
def _build_kiro_account_quota(upstream_metadata: dict[str, Any]) -> str | None:
kiro = upstream_metadata.get("kiro")
if not isinstance(kiro, dict):
return None
if kiro.get("is_banned") is True:
return "账号已封禁"
usage_percentage = _to_float(kiro.get("usage_percentage"))
if usage_percentage is not 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:
return (
f"剩余 {_format_percent(remaining)} "
f"({_format_quota_value(current_usage)}/{_format_quota_value(usage_limit)})"
)
return f"剩余 {_format_percent(remaining)}"
remaining = _to_float(kiro.get("remaining"))
usage_limit = _to_float(kiro.get("usage_limit"))
if remaining is not None and usage_limit is not None and usage_limit > 0:
return f"剩余 {_format_quota_value(remaining)}/{_format_quota_value(usage_limit)}"
return None
def _build_antigravity_account_quota(upstream_metadata: dict[str, Any]) -> str | None:
antigravity = upstream_metadata.get("antigravity")
if not isinstance(antigravity, dict):
return None
if antigravity.get("is_forbidden") is True:
return "访问受限"
quota_by_model = antigravity.get("quota_by_model")
if not isinstance(quota_by_model, dict) or not quota_by_model:
return None
remaining_list: list[float] = []
for raw_info in quota_by_model.values():
if not isinstance(raw_info, dict):
continue
used_percent = _to_float(raw_info.get("used_percent"))
if used_percent is None:
remaining_fraction = _to_float(raw_info.get("remaining_fraction"))
if remaining_fraction is not None:
used_percent = (1.0 - remaining_fraction) * 100.0
if used_percent is None:
continue
remaining = max(0.0, min(100.0 - used_percent, 100.0))
remaining_list.append(remaining)
if not remaining_list:
return None
min_remaining = min(remaining_list)
if len(remaining_list) == 1:
return f"剩余 {_format_percent(min_remaining)}"
return f"最低剩余 {_format_percent(min_remaining)} ({len(remaining_list)} 模型)"
def _build_account_quota(provider_type: str, upstream_metadata: Any) -> str | None:
if not isinstance(upstream_metadata, dict):
return None
normalized_type = provider_type.strip().lower()
if normalized_type == "codex":
return _build_codex_account_quota(upstream_metadata)
if normalized_type == "kiro":
return _build_kiro_account_quota(upstream_metadata)
if normalized_type == "antigravity":
return _build_antigravity_account_quota(upstream_metadata)
return None
return get_quota_reader(provider_type, upstream_metadata).display_summary()
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)
return get_quota_reader(provider_type, upstream_metadata).updated_at()
def _normalize_oauth_plan_type(plan_type: Any, provider_type: str) -> str | None:

View File

@@ -7,13 +7,17 @@ from __future__ import annotations
import asyncio
import json
import time
from collections.abc import Awaitable, Callable
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any
from uuid import uuid4
import httpx
from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel
from sqlalchemy.orm import Session, joinedload
from fastapi import APIRouter, Depends, HTTPException, Request
from pydantic import BaseModel, Field
from sqlalchemy import update
from sqlalchemy.orm import Session, joinedload, make_transient
from src.config.constants import TimeoutDefaults
from src.core.api_format import get_extra_headers_from_endpoint
@@ -21,8 +25,9 @@ from src.core.cache_service import CacheService
from src.core.crypto import crypto_service
from src.core.logger import logger
from src.core.provider_types import ProviderType
from src.database import create_session
from src.database.database import get_db
from src.models.database import Provider, ProviderEndpoint, User
from src.models.database import Provider, ProviderAPIKey, ProviderEndpoint, RequestCandidate, User
from src.services.model.fetch_scheduler import (
MODEL_FETCH_HTTP_TIMEOUT,
UPSTREAM_MODELS_CACHE_TTL_SECONDS,
@@ -39,6 +44,7 @@ from src.services.model.upstream_fetcher import (
)
from src.services.provider.oauth_token import resolve_oauth_access_token
from src.services.proxy_node.resolver import resolve_effective_proxy
from src.services.request.candidate import RequestCandidateService
from src.utils.auth_utils import get_current_user
if TYPE_CHECKING:
@@ -209,6 +215,7 @@ class TestModelFailoverRequest(BaseModel):
endpoint_id: str | None = None # 指定仅使用该端点测试
message: str | None = "Hello"
request_id: str | None = None
concurrency: int = Field(default=1, ge=1, le=20)
class TestAttemptDetail(BaseModel):
@@ -893,7 +900,7 @@ async def test_model(
def _response_has_error(resp: dict) -> bool:
"""快速判断响应是否包含错误"""
if "error" in resp:
if resp.get("error"):
return True
if resp.get("status_code", 0) != 200:
return True
@@ -905,7 +912,7 @@ async def test_model(
parsed = json.loads(resp_body)
except (json.JSONDecodeError, ValueError):
pass
if isinstance(parsed, dict) and "error" in parsed:
if isinstance(parsed, dict) and parsed.get("error"):
return True
return False
@@ -987,7 +994,7 @@ async def test_model(
except json.JSONDecodeError:
pass
if isinstance(parsed_body, dict) and "error" in parsed_body:
if isinstance(parsed_body, dict) and parsed_body.get("error"):
error_obj = parsed_body["error"]
# 兼容 error 可能是字典或字符串的情况
if isinstance(error_obj, dict):
@@ -1055,7 +1062,7 @@ async def test_model(
status_code=upstream_status,
detail=str(error_obj)[:500] if error_obj else "Provider error",
)
elif "error" in response:
elif response.get("error"):
logger.debug(f"[test-model] Error: {response['error']}")
upstream_status = int(response.get("status_code", 0) or 500)
if not (400 <= upstream_status <= 599):
@@ -1134,6 +1141,7 @@ def _build_direct_test_candidates(
为直接测试模式构建候选列表。
遍历 Provider 的活跃 Endpoint 和 Key不经过 GlobalModel 解析。
按可用性排序:熔断器关闭 > 健康度高 > 连续失败少 > Key 优先级。
"""
from src.services.scheduling.schemas import ProviderCandidate
@@ -1165,9 +1173,50 @@ def _build_direct_test_candidates(
provider_api_format=ep_format,
)
)
candidates.sort(key=lambda c: _direct_candidate_sort_key(c))
return candidates
def _direct_candidate_sort_key(candidate: ProviderCandidate) -> tuple[int, float, int, int]:
"""
按可用性排序候选:
1. 熔断器状态:关闭(0) > 打开(2)
2. 健康度评分:越高越好(取负值以升序排列)
3. 连续失败次数:越少越好
4. Key 优先级:数字越小越优先
"""
key = candidate.key
ep_format = candidate.provider_api_format
# 熔断器状态
circuit_breaker_order = 0
cb_data = getattr(key, "circuit_breaker_by_format", None) or {}
cb_entry = cb_data.get(ep_format, {}) if isinstance(cb_data, dict) else {}
if isinstance(cb_entry, dict) and cb_entry.get("open"):
circuit_breaker_order = 2
# 健康度评分(默认 1.0 表示完全健康)
health_score = 1.0
consecutive_failures = 0
health_data = getattr(key, "health_by_format", None) or {}
health_entry = health_data.get(ep_format, {}) if isinstance(health_data, dict) else {}
if isinstance(health_entry, dict):
health_score = health_entry.get("health_score", 1.0)
consecutive_failures = health_entry.get("consecutive_failures", 0)
# Key 优先级
internal_priority_raw = getattr(key, "internal_priority", None)
try:
internal_priority = (
int(internal_priority_raw) if internal_priority_raw is not None else 999999
)
except (TypeError, ValueError):
internal_priority = 999999
return (circuit_breaker_order, -health_score, consecutive_failures, internal_priority)
def _filter_test_candidates_by_endpoint(
candidates: list[ProviderCandidate],
endpoint_id: str | None,
@@ -1278,6 +1327,457 @@ def _build_test_candidate_meta(
return by_pair, by_candidate
def _flatten_test_candidates_for_concurrency(
candidates: list[ProviderCandidate],
) -> list[ProviderCandidate]:
from src.services.scheduling.schemas import (
PoolCandidate,
)
from src.services.scheduling.schemas import ProviderCandidate as SchedulerCandidate
flattened: list[ProviderCandidate] = []
for candidate in candidates:
if not isinstance(candidate, PoolCandidate):
flattened.append(candidate)
continue
for pool_key in candidate.pool_keys or []:
key_skipped = candidate.is_skipped or bool(getattr(pool_key, "_pool_skipped", False))
key_skip_reason_raw = (
getattr(pool_key, "_pool_skip_reason", None) if key_skipped else None
)
key_skip_reason = (
str(key_skip_reason_raw)
if key_skip_reason_raw
else (str(candidate.skip_reason) if candidate.skip_reason else None)
)
flattened.append(
SchedulerCandidate(
provider=candidate.provider,
endpoint=candidate.endpoint,
key=pool_key,
is_cached=bool(getattr(candidate, "is_cached", False)),
is_skipped=key_skipped,
skip_reason=key_skip_reason,
mapping_matched_model=(
getattr(pool_key, "_pool_mapping_matched_model", None)
or getattr(candidate, "mapping_matched_model", None)
),
needs_conversion=bool(getattr(candidate, "needs_conversion", False)),
provider_api_format=(
getattr(candidate, "provider_api_format", "")
or str(getattr(candidate.endpoint, "api_format", "") or "")
),
output_limit=getattr(candidate, "output_limit", None),
capability_miss_count=int(getattr(candidate, "capability_miss_count", 0) or 0),
)
)
return flattened
def _build_test_candidate_extra_data(candidate: ProviderCandidate) -> dict[str, Any]:
extra_data: dict[str, Any] = {
"needs_conversion": bool(getattr(candidate, "needs_conversion", False)),
"provider_api_format": (
getattr(candidate, "provider_api_format", None)
or getattr(getattr(candidate, "endpoint", None), "api_format", None)
),
"mapping_matched_model": (
getattr(candidate, "mapping_matched_model", None)
or getattr(getattr(candidate, "key", None), "_pool_mapping_matched_model", None)
),
}
key_extra = getattr(getattr(candidate, "key", None), "_pool_extra_data", None)
if isinstance(key_extra, dict):
extra_data.update(key_extra)
return extra_data
def _precreate_concurrent_test_records(
*,
db: Session,
request_id: str,
candidates: list[ProviderCandidate],
user: User | None,
) -> dict[int, str]:
record_map: dict[int, str] = {}
rows: list[dict[str, Any]] = []
user_id = str(getattr(user, "id", "") or "") or None
now = datetime.now(timezone.utc)
for candidate_index, candidate in enumerate(candidates):
record_id = str(uuid4())
record_map[candidate_index] = record_id
rows.append(
{
"id": record_id,
"request_id": request_id,
"candidate_index": candidate_index,
"retry_index": 0,
"user_id": user_id,
"api_key_id": None,
"provider_id": str(getattr(candidate.provider, "id", "") or "") or None,
"endpoint_id": str(getattr(candidate.endpoint, "id", "") or "") or None,
"key_id": str(getattr(candidate.key, "id", "") or "") or None,
"status": (
"skipped" if bool(getattr(candidate, "is_skipped", False)) else "available"
),
"skip_reason": getattr(candidate, "skip_reason", None),
"is_cached": bool(getattr(candidate, "is_cached", False)),
"extra_data": _build_test_candidate_extra_data(candidate),
"required_capabilities": None,
"created_at": now,
}
)
if rows:
db.bulk_insert_mappings(RequestCandidate, rows) # type: ignore[arg-type]
db.commit()
return record_map
def _mark_concurrent_test_record_cancelled(record_id: str) -> None:
if not record_id:
return
with create_session() as local_db:
RequestCandidateService.mark_candidate_cancelled(
db=local_db,
candidate_id=record_id,
status_code=499,
)
def _cancel_remaining_concurrent_test_records(request_id: str) -> None:
if not request_id:
return
with create_session() as local_db:
local_db.execute(
update(RequestCandidate)
.where(RequestCandidate.request_id == request_id)
.where(RequestCandidate.status.in_(["available", "pending"]))
.values(
status="cancelled",
status_code=499,
finished_at=datetime.now(timezone.utc),
)
)
local_db.commit()
async def _execute_test_check(
*,
provider_obj: Any,
endpoint: Any,
key: Any,
effective_model: str,
request_payload: dict[str, Any],
request_timeout: float,
provider_type: str,
user: User | None,
db: Session | None,
) -> tuple[dict[str, Any], str]:
effective_proxy = resolve_effective_proxy(
getattr(provider_obj, "proxy", None), getattr(key, "proxy", None)
)
try:
api_key_value, auth_config = await _resolve_key_auth(
key,
provider_obj,
provider_proxy_config=effective_proxy,
)
except _KeyAuthError as e:
raise RuntimeError(e.message) from e
auth_type = str(getattr(key, "auth_type", "api_key") or "api_key").lower()
extra_headers = get_extra_headers_from_endpoint(endpoint) or {}
if auth_type == "oauth":
account_id = (auth_config or {}).get("account_id")
if account_id:
extra_headers["chatgpt-account-id"] = str(account_id)
adapter_class = get_adapter_for_format(endpoint.api_format)
if not adapter_class:
raise ValueError(f"Unknown API format: {endpoint.api_format}")
response = await adapter_class.check_endpoint(
None,
endpoint.base_url,
api_key_value,
{
**request_payload,
"model": effective_model,
},
extra_headers if extra_headers else None,
body_rules=getattr(endpoint, "body_rules", None),
header_rules=getattr(endpoint, "header_rules", None),
db=db,
user=user,
provider_name=provider_obj.name,
provider_id=str(provider_obj.id),
api_key_id=str(key.id),
model_name=effective_model,
auth_type=auth_type,
provider_type=provider_type if provider_type else None,
decrypted_auth_config=auth_config if auth_config else None,
provider_endpoint=endpoint,
provider_api_key=key,
proxy_config=effective_proxy,
timeout_seconds=request_timeout,
)
return response, auth_type
async def _run_concurrent_test(
*,
candidates: list[ProviderCandidate],
concurrency: int,
is_cancelled: Callable[[], Awaitable[bool]],
request_id: str,
request_payload: dict[str, Any],
effective_model_by_candidate_index: dict[int, str],
request_timeout: float,
provider_type: str,
user: User | None,
db: Session,
) -> dict[str, Any]:
from src.core.exceptions import EmbeddedErrorException
from src.services.candidate.recorder import CandidateRecorder
from src.services.task.service import pool_on_error
semaphore = asyncio.Semaphore(max(1, concurrency))
record_map = _precreate_concurrent_test_records(
db=db,
request_id=request_id,
candidates=candidates,
user=user,
)
# 预加载所有候选的 provider/endpoint/key避免每个 worker 重复查询
_preloaded: dict[int, tuple[Provider, ProviderEndpoint, ProviderAPIKey]] = {}
with create_session() as preload_db:
provider_ids = {str(getattr(c.provider, "id", "") or "") for c in candidates}
endpoint_ids = {str(getattr(c.endpoint, "id", "") or "") for c in candidates}
key_ids = {str(getattr(c.key, "id", "") or "") for c in candidates}
providers_by_id = {
str(p.id): p
for p in preload_db.query(Provider).filter(Provider.id.in_(provider_ids)).all()
}
endpoints_by_id = {
str(e.id): e
for e in preload_db.query(ProviderEndpoint)
.filter(ProviderEndpoint.id.in_(endpoint_ids))
.all()
}
keys_by_id = {
str(k.id): k
for k in preload_db.query(ProviderAPIKey).filter(ProviderAPIKey.id.in_(key_ids)).all()
}
_already_detached: set[int] = set()
for idx, cand in enumerate(candidates):
p = providers_by_id.get(str(getattr(cand.provider, "id", "") or ""))
e = endpoints_by_id.get(str(getattr(cand.endpoint, "id", "") or ""))
k = keys_by_id.get(str(getattr(cand.key, "id", "") or ""))
if p is not None and e is not None and k is not None:
# make_transient 将对象脱离 session 并保留已加载属性,
# 避免 expired 状态导致跨协程访问时触发 lazy load 报错。
# 同一个对象(多个 candidate 可能共享同一 provider/endpoint
# 只需处理一次。
for obj in (p, e, k):
obj_id = id(obj)
if obj_id not in _already_detached:
make_transient(obj)
_already_detached.add(obj_id)
_preloaded[idx] = (p, e, k)
success_payload: dict[str, Any] = {}
success_event = asyncio.Event()
candidate_recorder = CandidateRecorder(db)
last_error: Exception | None = None
async def _worker(candidate_index: int) -> dict[str, Any]:
nonlocal last_error
record_id = record_map[candidate_index]
started = False
started_at = 0.0
try:
preloaded = _preloaded.get(candidate_index)
if preloaded is None:
raise RuntimeError("并发测试目标不存在或已被删除")
local_provider, local_endpoint, local_key = preloaded
if success_event.is_set() or await is_cancelled():
_mark_concurrent_test_record_cancelled(record_id)
return {"status": "cancelled"}
async with semaphore:
if success_event.is_set() or await is_cancelled():
_mark_concurrent_test_record_cancelled(record_id)
return {"status": "cancelled"}
with create_session() as update_db:
RequestCandidateService.mark_candidate_started(update_db, record_id)
started = True
started_at = time.perf_counter()
response, auth_type = await _execute_test_check(
provider_obj=local_provider,
endpoint=local_endpoint,
key=local_key,
effective_model=effective_model_by_candidate_index.get(
candidate_index,
str(request_payload.get("model", "") or ""),
),
request_payload=request_payload,
request_timeout=request_timeout,
provider_type=provider_type,
user=user,
db=None,
)
elapsed_ms = max(0, int((time.perf_counter() - started_at) * 1000))
with create_session() as parse_db:
parse_key = (
parse_db.query(ProviderAPIKey)
.filter(ProviderAPIKey.id == str(getattr(local_key, "id", "") or ""))
.first()
)
parsed = _extract_test_response_or_raise(
response=response,
endpoint=local_endpoint,
provider_name=str(local_provider.name),
auth_type=auth_type,
api_key=parse_key or local_key,
db=parse_db,
)
with create_session() as update_db:
RequestCandidateService.mark_candidate_success(
db=update_db,
candidate_id=record_id,
status_code=200,
latency_ms=elapsed_ms,
)
if not success_event.is_set():
success_payload.update(
{
"response": parsed,
"candidate_index": candidate_index,
"key_id": str(getattr(local_key, "id", "") or "") or None,
}
)
success_event.set()
return {"status": "success"}
except asyncio.CancelledError:
if started or not success_event.is_set():
_mark_concurrent_test_record_cancelled(record_id)
return {"status": "cancelled"}
except Exception as exc:
last_error = exc
elapsed_ms = max(0, int((time.perf_counter() - started_at) * 1000)) if started else None
status_code = None
if isinstance(exc, httpx.HTTPStatusError):
status_code = int(exc.response.status_code)
elif isinstance(exc, httpx.TimeoutException):
status_code = 408
elif isinstance(exc, EmbeddedErrorException):
status_code = int(exc.error_code or 200)
loaded = _preloaded.get(candidate_index)
if loaded is not None and status_code is not None:
await pool_on_error(loaded[0], loaded[2], status_code, exc)
with create_session() as update_db:
RequestCandidateService.mark_candidate_failed(
db=update_db,
candidate_id=record_id,
error_type=type(exc).__name__,
error_message=str(
getattr(exc, "error_message", None)
or getattr(exc, "upstream_response", None)
or exc
),
status_code=status_code,
latency_ms=elapsed_ms,
)
return {"status": "failed", "error": exc}
async def _watch_disconnect() -> bool:
while not success_event.is_set():
if await is_cancelled():
return True
await asyncio.sleep(0.1)
return False
tasks = [
asyncio.create_task(_worker(candidate_index))
for candidate_index, candidate in enumerate(candidates)
if not bool(getattr(candidate, "is_skipped", False))
]
disconnect_task = asyncio.create_task(_watch_disconnect())
pending: set[asyncio.Task[Any]] = set(tasks)
pending.add(disconnect_task)
try:
while pending:
if pending == {disconnect_task}:
disconnect_task.cancel()
pending.clear()
break
done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
if disconnect_task in done and disconnect_task.result() is True:
for task in pending:
task.cancel()
break
for finished in done:
if finished is disconnect_task:
continue
result = finished.result()
if result.get("status") == "success":
for task in pending:
task.cancel()
pending.discard(disconnect_task)
disconnect_task.cancel()
break
if success_event.is_set():
break
finally:
await asyncio.gather(*pending, return_exceptions=True)
if not disconnect_task.done():
disconnect_task.cancel()
await asyncio.gather(disconnect_task, return_exceptions=True)
if success_event.is_set():
_cancel_remaining_concurrent_test_records(request_id)
elif await is_cancelled():
_cancel_remaining_concurrent_test_records(request_id)
try:
db.expire_all()
candidate_keys = candidate_recorder.get_candidate_keys(request_id)
except Exception:
candidate_keys = []
attempt_count = sum(
1
for item in candidate_keys
if str(getattr(item, "status", "") or "")
not in {"skipped", "cancelled", "available", "unused"}
)
return {
"success": success_event.is_set(),
"candidate_keys": candidate_keys,
"attempt_count": attempt_count,
"run_error": last_error,
"response": success_payload.get("response"),
}
def _maybe_mark_test_oauth_key_invalid(
*,
db: Session,
@@ -1326,7 +1826,7 @@ def _extract_test_response_or_raise(
if isinstance(parsed_payload, dict) and "response_body" in parsed_payload:
parsed_payload = _parse_jsonish(parsed_payload.get("response_body"))
if isinstance(parsed_payload, dict) and "error" in parsed_payload:
if isinstance(parsed_payload, dict) and parsed_payload.get("error"):
_maybe_mark_test_oauth_key_invalid(
db=db,
key=api_key,
@@ -1444,6 +1944,7 @@ def _build_test_attempts_from_candidate_keys(
@router.post("/test-model-failover")
async def test_model_failover(
request: TestModelFailoverRequest,
http_request: Request,
db: Session = Depends(get_db),
current_user: User = Depends(get_current_user),
) -> Any:
@@ -1571,25 +2072,6 @@ async def test_model_failover(
provider_type = str(getattr(provider, "provider_type", "") or "").lower()
async def _request_func(provider_obj: Any, endpoint: Any, key: Any, candidate: Any) -> Any:
effective_proxy = resolve_effective_proxy(
getattr(provider_obj, "proxy", None), getattr(key, "proxy", None)
)
try:
api_key_value, auth_config = await _resolve_key_auth(
key,
provider_obj,
provider_proxy_config=effective_proxy,
)
except _KeyAuthError as e:
raise RuntimeError(e.message) from e
auth_type = str(getattr(key, "auth_type", "api_key") or "api_key").lower()
extra_headers = get_extra_headers_from_endpoint(endpoint) or {}
if auth_type == "oauth":
account_id = (auth_config or {}).get("account_id")
if account_id:
extra_headers["chatgpt-account-id"] = str(account_id)
effective_model = _resolve_test_effective_model(
provider=provider,
candidate=candidate,
@@ -1597,34 +2079,16 @@ async def test_model_failover(
gm_obj=gm_obj,
key=key,
)
adapter_class = get_adapter_for_format(endpoint.api_format)
if not adapter_class:
raise ValueError(f"Unknown API format: {endpoint.api_format}")
response = await adapter_class.check_endpoint(
None,
endpoint.base_url,
api_key_value,
{
**request_payload,
"model": effective_model,
},
extra_headers if extra_headers else None,
body_rules=getattr(endpoint, "body_rules", None),
header_rules=getattr(endpoint, "header_rules", None),
db=db,
response, auth_type = await _execute_test_check(
provider_obj=provider_obj,
endpoint=endpoint,
key=key,
effective_model=effective_model,
request_payload=request_payload,
request_timeout=request_timeout,
provider_type=provider_type,
user=current_user,
provider_name=provider_obj.name,
provider_id=str(provider_obj.id),
api_key_id=str(key.id),
model_name=effective_model,
auth_type=auth_type,
provider_type=provider_type if provider_type else None,
decrypted_auth_config=auth_config if auth_config else None,
provider_endpoint=endpoint,
provider_api_key=key,
proxy_config=effective_proxy,
timeout_seconds=request_timeout,
db=db,
)
return _extract_test_response_or_raise(
response=response,
@@ -1639,46 +2103,97 @@ async def test_model_failover(
task_service = TaskService(db)
exec_result = None
run_error: Exception | None = None
concurrent_result: dict[str, Any] | None = None
result_candidates = candidates
if request.concurrency > 1:
result_candidates = _flatten_test_candidates_for_concurrency(candidates)
candidate_meta_by_pair, candidate_meta_by_index = _build_test_candidate_meta(
candidates=result_candidates,
provider=provider,
request=request,
gm_obj=gm_obj,
)
effective_model_by_candidate_index = {
index: str(meta.get("effective_model") or request.model_name)
for index, meta in candidate_meta_by_index.items()
}
try:
exec_result = await task_service.execute_sync_candidates(
api_format=client_format or "openai:chat",
model_name=request.model_name,
candidates=candidates,
request_func=_request_func,
request_id=request_id,
current_user=current_user,
user_api_key=None,
is_stream=False,
capability_requirements=None,
request_body_ref={"body": dict(request_payload)},
request_headers=None,
request_body=dict(request_payload),
affinity_key=f"provider-test:{provider.id}",
create_pending_usage=False,
enable_cache_affinity=False,
)
if request.concurrency > 1:
concurrent_result = await _run_concurrent_test(
candidates=result_candidates,
concurrency=request.concurrency,
is_cancelled=http_request.is_disconnected,
request_id=request_id,
request_payload=dict(request_payload),
effective_model_by_candidate_index=effective_model_by_candidate_index,
request_timeout=request_timeout,
provider_type=provider_type,
user=current_user,
db=db,
)
else:
exec_result = await task_service.execute_sync_candidates(
api_format=client_format or "openai:chat",
model_name=request.model_name,
candidates=result_candidates,
request_func=_request_func,
request_id=request_id,
current_user=current_user,
user_api_key=None,
is_stream=False,
capability_requirements=None,
request_body_ref={"body": dict(request_payload)},
request_headers=None,
request_body=dict(request_payload),
affinity_key=f"provider-test:{provider.id}",
create_pending_usage=False,
enable_cache_affinity=False,
is_cancelled=http_request.is_disconnected,
)
except Exception as exc:
run_error = exc
logger.error("[test-model-failover] Error: {}", exc)
try:
candidate_keys = candidate_recorder.get_candidate_keys(request_id)
candidate_keys = (
list(concurrent_result.get("candidate_keys", []))
if concurrent_result is not None
else candidate_recorder.get_candidate_keys(request_id)
)
except Exception:
candidate_keys = list(exec_result.candidate_keys) if exec_result else []
candidate_meta_by_pair, candidate_meta_by_index = _build_test_candidate_meta(
candidates=candidates,
provider=provider,
request=request,
gm_obj=gm_obj,
)
attempts = _build_test_attempts_from_candidate_keys(
candidate_keys=candidate_keys,
candidate_meta_by_pair=candidate_meta_by_pair,
candidate_meta_by_index=candidate_meta_by_index,
)
total_attempts = sum(1 for attempt in attempts if attempt.status != "skipped")
total_attempts = (
int(exec_result.attempt_count)
if exec_result is not None
else (
int(concurrent_result.get("attempt_count", 0))
if concurrent_result is not None
else sum(1 for attempt in attempts if attempt.status not in {"skipped", "cancelled"})
)
)
if concurrent_result is not None and concurrent_result.get("success"):
return TestModelFailoverResponse(
success=True,
model=request.model_name,
provider={"id": str(provider.id), "name": provider.name},
attempts=attempts,
total_candidates=len(result_candidates),
total_attempts=total_attempts,
data={
"stream": True,
"response": concurrent_result.get("response"),
},
error=None,
).model_dump()
if exec_result and exec_result.success:
return TestModelFailoverResponse(
@@ -1686,7 +2201,7 @@ async def test_model_failover(
model=request.model_name,
provider={"id": str(provider.id), "name": provider.name},
attempts=attempts,
total_candidates=len(candidates),
total_candidates=len(result_candidates),
total_attempts=exec_result.attempt_count,
data={
"stream": True,
@@ -1703,6 +2218,10 @@ async def test_model_failover(
error_message = str(run_error.upstream_response)[:500]
if not error_message:
error_message = str(run_error)
if not error_message and concurrent_result is not None and concurrent_result.get("run_error"):
error_message = str(concurrent_result.get("run_error"))
if not error_message and exec_result is not None and exec_result.error_message:
error_message = str(exec_result.error_message)
if not error_message:
failed_attempt = next(
(attempt for attempt in reversed(attempts) if attempt.error_message),
@@ -1717,7 +2236,7 @@ async def test_model_failover(
model=request.model_name,
provider={"id": str(provider.id), "name": provider.name},
attempts=attempts,
total_candidates=len(candidates),
total_candidates=len(result_candidates),
total_attempts=total_attempts,
error=str(error_message)[:500],
).model_dump()

View File

@@ -181,6 +181,147 @@ class FailoverEngine:
)
await asyncio.sleep(backoff_seconds)
async def _check_cancellation(
self,
is_cancelled: Callable[[], Awaitable[bool]] | None,
) -> bool:
if is_cancelled is None:
return False
try:
return bool(await is_cancelled())
except Exception:
return False
def _mark_remaining_cancelled(
self,
*,
candidate_record_map: dict[tuple[int, int], str] | None,
candidates: list[ProviderCandidate],
from_candidate_idx: int,
from_retry_idx: int,
retry_policy: RetryPolicy,
) -> None:
if not candidate_record_map:
return
now = datetime.now(timezone.utc)
updated = False
for candidate_idx, cand in enumerate(candidates):
if candidate_idx < from_candidate_idx:
continue
max_retries = self._get_max_retries(cand, retry_policy)
for retry_idx in range(max_retries):
if candidate_idx == from_candidate_idx and retry_idx < from_retry_idx:
continue
record_id = candidate_record_map.get((candidate_idx, retry_idx))
if not record_id:
continue
self.db.execute(
update(RequestCandidate)
.where(RequestCandidate.id == record_id)
.where(RequestCandidate.status.in_(["available", "pending"]))
.values(
status="cancelled",
status_code=499,
error_message="cancelled_by_client",
finished_at=now,
)
)
updated = True
if updated:
self.db.commit()
def _append_cancelled_fallback_candidate_keys(
self,
*,
fallback: list[CandidateKey],
candidates: list[ProviderCandidate],
from_candidate_idx: int,
from_retry_idx: int,
retry_policy: RetryPolicy,
) -> None:
existing = {(item.candidate_index, item.retry_index) for item in fallback}
for candidate_idx, cand in enumerate(candidates):
if candidate_idx < from_candidate_idx:
continue
max_retries = self._get_max_retries(cand, retry_policy)
for retry_idx in range(max_retries):
if candidate_idx == from_candidate_idx and retry_idx < from_retry_idx:
continue
key = (candidate_idx, retry_idx)
if key in existing:
continue
original_key = getattr(cand, "key", None)
original_pool_key_index = getattr(cand, "_pool_key_index", 0)
if isinstance(cand, PoolCandidate) and cand.pool_keys:
retry_slots_per_key = self._get_pool_key_max_retries(cand, retry_policy)
pool_key_index = min(retry_idx // retry_slots_per_key, len(cand.pool_keys) - 1)
cand.key = cand.pool_keys[pool_key_index]
cand._pool_key_index = pool_key_index
fallback.append(
self._make_candidate_key(
candidate=cand,
candidate_index=candidate_idx,
retry_index=retry_idx,
status="cancelled",
error_message="cancelled_by_client",
status_code=499,
)
)
if isinstance(cand, PoolCandidate):
cand.key = original_key
cand._pool_key_index = original_pool_key_index
existing.add(key)
async def _maybe_cancel_execution(
self,
*,
is_cancelled: Callable[[], Awaitable[bool]] | None,
candidate_record_map: dict[tuple[int, int], str] | None,
candidate_keys_fallback: list[CandidateKey],
candidates: list[ProviderCandidate],
from_candidate_idx: int,
from_retry_idx: int,
retry_policy: RetryPolicy,
request_id: str | None,
attempt_count: int,
) -> ExecutionResult | None:
if not await self._check_cancellation(is_cancelled):
return None
logger.info(
"[FailoverEngine] Request cancelled by client at candidate_index={}, retry_index={}",
from_candidate_idx,
from_retry_idx,
)
self._mark_remaining_cancelled(
candidate_record_map=candidate_record_map,
candidates=candidates,
from_candidate_idx=from_candidate_idx,
from_retry_idx=from_retry_idx,
retry_policy=retry_policy,
)
self._append_cancelled_fallback_candidate_keys(
fallback=candidate_keys_fallback,
candidates=candidates,
from_candidate_idx=from_candidate_idx,
from_retry_idx=from_retry_idx,
retry_policy=retry_policy,
)
return ExecutionResult(
success=False,
error_type="cancelled",
error_message="cancelled_by_client",
last_status_code=499,
candidate_keys=self._get_candidate_keys(
request_id=request_id,
fallback=candidate_keys_fallback,
candidates=candidates,
),
attempt_count=attempt_count,
)
async def execute(
self,
*,
@@ -201,6 +342,7 @@ class FailoverEngine:
]
| None
) = None,
is_cancelled: Callable[[], Awaitable[bool]] | None = None,
) -> ExecutionResult:
"""
Execute candidate traversal + retry + failover.
@@ -229,6 +371,20 @@ class FailoverEngine:
max_attempts = computed
for candidate_index, candidate in enumerate(candidates):
cancelled_result = await self._maybe_cancel_execution(
is_cancelled=is_cancelled,
candidate_record_map=candidate_record_map,
candidate_keys_fallback=candidate_keys_fallback,
candidates=candidates,
from_candidate_idx=candidate_index,
from_retry_idx=0,
retry_policy=retry_policy,
request_id=request_id,
attempt_count=attempt_count,
)
if cancelled_result is not None:
return cancelled_result
should_skip, skip_reason = self._should_skip(candidate, skip_policy)
if should_skip:
# PRE_EXPAND: mark all retry slots skipped.
@@ -279,6 +435,7 @@ class FailoverEngine:
max_attempts=max_attempts,
execution_error_handler=execution_error_handler,
consecutive_failures=consecutive_failures,
is_cancelled=is_cancelled,
)
)
if pool_result is not None:
@@ -288,6 +445,20 @@ class FailoverEngine:
max_retries = self._get_max_retries(candidate, retry_policy)
retry_index = 0
while retry_index < max_retries:
cancelled_result = await self._maybe_cancel_execution(
is_cancelled=is_cancelled,
candidate_record_map=candidate_record_map,
candidate_keys_fallback=candidate_keys_fallback,
candidates=candidates,
from_candidate_idx=candidate_index,
from_retry_idx=retry_index,
retry_policy=retry_policy,
request_id=request_id,
attempt_count=attempt_count,
)
if cancelled_result is not None:
return cancelled_result
attempt_count += 1
# Resolve/create record_id
@@ -456,6 +627,7 @@ class FailoverEngine:
consecutive_failures: int,
max_attempts: int | None,
execution_error_handler: Any,
is_cancelled: Callable[[], Awaitable[bool]] | None,
) -> tuple[ExecutionResult | None, int, int, int | None]:
"""Execute a PoolCandidate with in-pool key failover."""
last_status_code: int | None = None
@@ -463,6 +635,20 @@ class FailoverEngine:
for key_index, pool_key in enumerate(candidate.pool_keys or []):
base_retry_index = key_index * retry_slots_per_key
cancelled_result = await self._maybe_cancel_execution(
is_cancelled=is_cancelled,
candidate_record_map=candidate_record_map,
candidate_keys_fallback=candidate_keys_fallback,
candidates=candidates,
from_candidate_idx=candidate_index,
from_retry_idx=base_retry_index,
retry_policy=retry_policy,
request_id=request_id,
attempt_count=attempt_count,
)
if cancelled_result is not None:
return cancelled_result, attempt_count, consecutive_failures, last_status_code
candidate.key = pool_key
candidate._pool_key_index = key_index
candidate.mapping_matched_model = getattr(pool_key, "_pool_mapping_matched_model", None)
@@ -507,8 +693,22 @@ class FailoverEngine:
max_retries_for_key = retry_slots_per_key
retry_index = 0
while retry_index < max_retries_for_key:
attempt_count += 1
composite_retry_index = base_retry_index + retry_index
cancelled_result = await self._maybe_cancel_execution(
is_cancelled=is_cancelled,
candidate_record_map=candidate_record_map,
candidate_keys_fallback=candidate_keys_fallback,
candidates=candidates,
from_candidate_idx=candidate_index,
from_retry_idx=composite_retry_index,
retry_policy=retry_policy,
request_id=request_id,
attempt_count=attempt_count,
)
if cancelled_result is not None:
return cancelled_result, attempt_count, consecutive_failures, last_status_code
attempt_count += 1
record_id = None
if candidate_record_map:

View File

@@ -9,6 +9,8 @@ from __future__ import annotations
from dataclasses import dataclass
from typing import Any
from src.services.provider_keys.quota_reader import get_quota_reader
OAUTH_ACCOUNT_BLOCK_PREFIX = "[ACCOUNT_BLOCK] "
OAUTH_REFRESH_FAILED_PREFIX = "[REFRESH_FAILED] "
OAUTH_EXPIRED_PREFIX = "[OAUTH_EXPIRED] "
@@ -69,7 +71,7 @@ def _classify_block_reason(text: str) -> tuple[str, str]:
return "oauth_expired", "Token 失效"
if any(kw in lowered for kw in _KEYWORDS_VERIFICATION):
return "account_verification", "需要验证"
if 'deactivated_workspace' in lowered:
if "deactivated_workspace" in lowered:
return "workspace_deactivated", "工作区停用"
if any(kw in lowered for kw in _KEYWORDS_DISABLED):
return "account_disabled", "账号停用"
@@ -130,30 +132,13 @@ def _resolve_from_metadata(
if isinstance(maybe_bucket, dict):
provider_bucket = maybe_bucket
if (
normalized_provider == "kiro"
and provider_bucket
and _is_truthy_flag(provider_bucket.get("is_banned"))
):
reason = _extract_reason(provider_bucket, "ban_reason", "reason", "message")
quota_block = get_quota_reader(normalized_provider, upstream_metadata).account_block()
if quota_block.blocked:
return PoolAccountState(
blocked=True,
code="account_banned",
label="账号封禁",
reason=reason or "Kiro 账号已封禁",
)
if (
normalized_provider == "antigravity"
and provider_bucket
and _is_truthy_flag(provider_bucket.get("is_forbidden"))
):
reason = _extract_reason(provider_bucket, "forbidden_reason", "reason", "message")
return PoolAccountState(
blocked=True,
code="account_forbidden",
label="访问受限",
reason=reason or "Antigravity 账户访问受限",
code=quota_block.code,
label=quota_block.label,
reason=quota_block.reason,
)
for source in (provider_bucket, upstream_metadata):

View File

@@ -3,9 +3,11 @@
from __future__ import annotations
import math
import time
from typing import Any
from src.core.provider_types import ProviderType
from src.services.provider_keys.quota_reader import get_quota_reader
def safe_float(value: Any) -> float | None:
try:
@@ -98,26 +100,10 @@ def extract_plan_type(key_obj: Any) -> str | None:
return direct
metadata = safe_metadata(key_obj)
codex = metadata.get("codex")
if isinstance(codex, dict):
codex_plan = normalize_plan(codex.get("plan_type"))
if codex_plan:
return codex_plan
kiro = metadata.get("kiro")
if isinstance(kiro, dict):
subscription_title = normalize_plan(kiro.get("subscription_title"))
if subscription_title:
# Normalize common Kiro labels into free/team buckets used by free_team_first.
if "team" in subscription_title:
return "team"
if "free" in subscription_title:
return "free"
if "pro" in subscription_title:
return "pro"
if "plus" in subscription_title:
return "plus"
return subscription_title
for provider_type in (ProviderType.CODEX, ProviderType.KIRO, ProviderType.ANTIGRAVITY):
plan_type = get_quota_reader(provider_type, metadata).plan_type()
if plan_type:
return plan_type
return None
@@ -126,19 +112,11 @@ def extract_reset_seconds(key_obj: Any) -> float | None:
metadata = safe_metadata(key_obj)
candidates: list[float] = []
codex = metadata.get("codex")
if isinstance(codex, dict):
for field in ("secondary_reset_seconds", "primary_reset_seconds"):
parsed = safe_float(codex.get(field))
if parsed is None or parsed < 0:
continue
candidates.append(parsed)
kiro = metadata.get("kiro")
if isinstance(kiro, dict):
next_reset_at = safe_float(kiro.get("next_reset_at"))
if next_reset_at is not None and next_reset_at > 0:
candidates.append(max(0.0, next_reset_at - time.time()))
for provider_type in (ProviderType.CODEX, ProviderType.KIRO, ProviderType.ANTIGRAVITY):
reset_seconds = get_quota_reader(provider_type, metadata).reset_seconds()
if reset_seconds is None:
continue
candidates.append(reset_seconds)
if not candidates:
return None
@@ -148,41 +126,10 @@ def extract_reset_seconds(key_obj: Any) -> float | None:
def extract_usage_ratio(key_obj: Any) -> float | None:
metadata = safe_metadata(key_obj)
codex = metadata.get("codex")
if isinstance(codex, dict):
codex_values: list[float] = []
for field in ("primary_used_percent", "secondary_used_percent"):
parsed = safe_float(codex.get(field))
if parsed is None:
continue
codex_values.append(max(0.0, min(parsed, 100.0)) / 100.0)
if codex_values:
return sum(codex_values) / len(codex_values)
kiro = metadata.get("kiro")
if isinstance(kiro, dict):
parsed = safe_float(kiro.get("usage_percentage"))
if parsed is not None:
return max(0.0, min(parsed, 100.0)) / 100.0
antigravity = metadata.get("antigravity")
if isinstance(antigravity, dict):
quota_by_model = antigravity.get("quota_by_model")
if isinstance(quota_by_model, dict):
usage_values: list[float] = []
for model_info in quota_by_model.values():
if not isinstance(model_info, dict):
continue
used_percent = safe_float(model_info.get("used_percent"))
if used_percent is None:
remaining_fraction = safe_float(model_info.get("remaining_fraction"))
if remaining_fraction is not None:
used_percent = (1.0 - remaining_fraction) * 100.0
if used_percent is None:
continue
usage_values.append(max(0.0, min(used_percent, 100.0)) / 100.0)
if usage_values:
return sum(usage_values) / len(usage_values)
for provider_type in (ProviderType.CODEX, ProviderType.KIRO, ProviderType.ANTIGRAVITY):
usage_ratio = get_quota_reader(provider_type, metadata).usage_ratio()
if usage_ratio is not None:
return usage_ratio
return None

View File

@@ -26,6 +26,13 @@ from src.services.provider_keys.quota_refresh import (
QuotaRefreshHandler = Callable[..., Awaitable[dict]]
_QUOTA_REFRESH_HANDLERS: dict[str, QuotaRefreshHandler] = {
ProviderType.CODEX: refresh_codex_key_quota,
ProviderType.ANTIGRAVITY: refresh_antigravity_key_quota,
ProviderType.KIRO: refresh_kiro_key_quota,
}
def _normalize_api_format(api_format: Any) -> str:
"""规范化 api_format兼容大小写和首尾空白。"""
if not isinstance(api_format, str):
@@ -55,12 +62,9 @@ def _select_refresh_endpoint(provider: Provider, provider_type: str) -> Provider
def _resolve_quota_refresh_handler(provider_type: str) -> QuotaRefreshHandler:
"""按 provider 类型返回刷新策略。"""
if provider_type == ProviderType.CODEX:
return refresh_codex_key_quota
if provider_type == ProviderType.ANTIGRAVITY:
return refresh_antigravity_key_quota
if provider_type == ProviderType.KIRO:
return refresh_kiro_key_quota
handler = _QUOTA_REFRESH_HANDLERS.get(provider_type)
if handler is not None:
return handler
raise InvalidRequestException("仅支持 Codex / Antigravity / Kiro 类型的 Provider 刷新限额")

View File

@@ -0,0 +1,439 @@
"""Unified quota readers for provider key upstream metadata."""
from __future__ import annotations
import math
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Any
from src.core.provider_types import ProviderType, normalize_provider_type
def _to_float(value: Any) -> float | None:
try:
parsed = float(value)
except (TypeError, ValueError):
return None
if math.isnan(parsed) or math.isinf(parsed):
return None
return parsed
def _normalize_plan(value: Any) -> str | None:
if not isinstance(value, str):
return None
normalized = value.strip().lower()
return normalized or 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 _extract_reason(source: dict[str, Any], *fields: str) -> str | None:
for field in fields:
value = source.get(field)
if not isinstance(value, str):
continue
text = value.strip()
if text:
return text
return None
def _pct_is_exhausted(value: Any) -> bool:
pct = _to_float(value)
if pct is None:
return False
return pct >= 100.0 - 1e-6
def _format_percent(value: float) -> str:
clamped = max(0.0, min(value, 100.0))
return f"{clamped:.1f}%"
def _format_quota_value(value: float) -> str:
rounded = round(value)
if abs(value - rounded) < 1e-6:
return str(rounded)
return f"{value:.1f}"
def _format_reset_after(seconds_raw: Any) -> str | None:
seconds = _to_float(seconds_raw)
if seconds is None:
return None
total_seconds = int(seconds)
if total_seconds <= 0:
return "已重置"
days = total_seconds // 86400
hours = (total_seconds % 86400) // 3600
minutes = (total_seconds % 3600) // 60
if days > 0:
return f"{days}{hours}小时后重置"
if hours > 0:
return f"{hours}小时{minutes}分钟后重置"
if minutes > 0:
return f"{minutes}分钟后重置"
return "即将重置"
@dataclass(frozen=True, slots=True)
class QuotaExhaustedResult:
exhausted: bool
reason: str | None = None
@dataclass(frozen=True, slots=True)
class AccountBlockResult:
blocked: bool
code: str | None = None
label: str | None = None
reason: str | None = None
class PoolQuotaReader(ABC):
"""Read-only view over one provider namespace in upstream_metadata."""
namespace: str | None = None
def __init__(self, data: dict[str, Any] | None) -> None:
self._data: dict[str, Any] = data if isinstance(data, dict) else {}
@abstractmethod
def is_exhausted(self, model_name: str | None = None) -> QuotaExhaustedResult:
"""Return whether this key/model should be skipped for quota exhaustion."""
@abstractmethod
def usage_ratio(self) -> float | None:
"""Return usage ratio within [0, 1], when available."""
@abstractmethod
def plan_type(self) -> str | None:
"""Return normalized plan type, when available."""
@abstractmethod
def reset_seconds(self) -> float | None:
"""Return seconds until next reset, when available."""
@abstractmethod
def account_block(self) -> AccountBlockResult:
"""Return account-level block state derived from metadata."""
@abstractmethod
def display_summary(self) -> str | None:
"""Return admin-facing quota summary string."""
def updated_at(self) -> int | None:
updated_at = _to_float(self._data.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)
class NullQuotaReader(PoolQuotaReader):
def is_exhausted(self, model_name: str | None = None) -> QuotaExhaustedResult:
_ = model_name
return QuotaExhaustedResult(exhausted=False)
def usage_ratio(self) -> float | None:
return None
def plan_type(self) -> str | None:
return None
def reset_seconds(self) -> float | None:
return None
def account_block(self) -> AccountBlockResult:
return AccountBlockResult(blocked=False)
def display_summary(self) -> str | None:
return None
class CodexQuotaReader(PoolQuotaReader):
namespace = "codex"
def is_exhausted(self, model_name: str | None = None) -> QuotaExhaustedResult:
_ = model_name
exhausted_parts: list[str] = []
if _pct_is_exhausted(self._data.get("primary_used_percent")):
exhausted_parts.append("周限额剩余 0%")
if _pct_is_exhausted(self._data.get("secondary_used_percent")):
exhausted_parts.append("5H 限额剩余 0%")
if exhausted_parts:
return QuotaExhaustedResult(True, "Codex " + "".join(exhausted_parts))
return QuotaExhaustedResult(False)
def usage_ratio(self) -> float | None:
values: list[float] = []
for field in ("primary_used_percent", "secondary_used_percent"):
parsed = _to_float(self._data.get(field))
if parsed is None:
continue
values.append(max(0.0, min(parsed, 100.0)) / 100.0)
if not values:
return None
return sum(values) / len(values)
def plan_type(self) -> str | None:
return _normalize_plan(self._data.get("plan_type"))
def reset_seconds(self) -> float | None:
candidates: list[float] = []
for field in ("secondary_reset_seconds", "primary_reset_seconds"):
parsed = _to_float(self._data.get(field))
if parsed is None or parsed < 0:
continue
candidates.append(parsed)
if not candidates:
return None
return min(candidates)
def account_block(self) -> AccountBlockResult:
if not _is_truthy_flag(self._data.get("account_disabled")):
return AccountBlockResult(blocked=False)
reason = _extract_reason(self._data, "forbidden_reason", "ban_reason", "reason", "message")
return AccountBlockResult(
blocked=True,
code="account_forbidden",
label="访问受限",
reason=reason or "账号访问受限",
)
def display_summary(self) -> str | None:
parts: list[str] = []
primary_used = _to_float(self._data.get("primary_used_percent"))
if primary_used is not None:
part = f"周剩余 {_format_percent(100.0 - primary_used)}"
reset_text = _format_reset_after(self._data.get("primary_reset_seconds"))
if reset_text:
part = f"{part} ({reset_text})"
parts.append(part)
secondary_used = _to_float(self._data.get("secondary_used_percent"))
if secondary_used is not None:
part = f"5H剩余 {_format_percent(100.0 - secondary_used)}"
reset_text = _format_reset_after(self._data.get("secondary_reset_seconds"))
if reset_text:
part = f"{part} ({reset_text})"
parts.append(part)
if parts:
return " | ".join(parts)
has_credits = self._data.get("has_credits")
credits_balance = _to_float(self._data.get("credits_balance"))
if has_credits is True and credits_balance is not None:
return f"积分 {credits_balance:.2f}"
if has_credits is True:
return "有积分"
return None
class KiroQuotaReader(PoolQuotaReader):
namespace = "kiro"
def is_exhausted(self, model_name: str | None = None) -> QuotaExhaustedResult:
_ = model_name
remaining = _to_float(self._data.get("remaining"))
if remaining is not None and remaining <= 0.0:
return QuotaExhaustedResult(True, "Kiro 账号配额剩余 0")
return QuotaExhaustedResult(False)
def usage_ratio(self) -> float | None:
parsed = _to_float(self._data.get("usage_percentage"))
if parsed is None:
return None
return max(0.0, min(parsed, 100.0)) / 100.0
def plan_type(self) -> str | None:
subscription_title = _normalize_plan(self._data.get("subscription_title"))
if not subscription_title:
return None
if "team" in subscription_title:
return "team"
if "free" in subscription_title:
return "free"
if "pro" in subscription_title:
return "pro"
if "plus" in subscription_title:
return "plus"
return subscription_title
def reset_seconds(self) -> float | None:
next_reset_at = _to_float(self._data.get("next_reset_at"))
if next_reset_at is None or next_reset_at <= 0:
return None
return max(0.0, next_reset_at - time.time())
def account_block(self) -> AccountBlockResult:
if not _is_truthy_flag(self._data.get("is_banned")):
return AccountBlockResult(blocked=False)
reason = _extract_reason(self._data, "ban_reason", "reason", "message")
return AccountBlockResult(
blocked=True,
code="account_banned",
label="账号封禁",
reason=reason or "Kiro 账号已封禁",
)
def display_summary(self) -> str | None:
if self._data.get("is_banned") is True:
return "账号已封禁"
usage_percentage = _to_float(self._data.get("usage_percentage"))
if usage_percentage is not None:
remaining = 100.0 - usage_percentage
current_usage = _to_float(self._data.get("current_usage"))
usage_limit = _to_float(self._data.get("usage_limit"))
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)})"
)
return f"剩余 {_format_percent(remaining)}"
remaining = _to_float(self._data.get("remaining"))
usage_limit = _to_float(self._data.get("usage_limit"))
if remaining is not None and usage_limit is not None and usage_limit > 0:
return f"剩余 {_format_quota_value(remaining)}/{_format_quota_value(usage_limit)}"
return None
class AntigravityQuotaReader(PoolQuotaReader):
namespace = "antigravity"
def _quota_by_model(self) -> dict[str, Any]:
quota_by_model = self._data.get("quota_by_model")
if not isinstance(quota_by_model, dict):
return {}
return quota_by_model
def _used_percent(self, model_info: dict[str, Any]) -> float | None:
used_percent = _to_float(model_info.get("used_percent"))
if used_percent is not None:
return max(0.0, min(used_percent, 100.0))
remaining_fraction = _to_float(model_info.get("remaining_fraction"))
if remaining_fraction is None:
return None
return max(0.0, min((1.0 - remaining_fraction) * 100.0, 100.0))
def is_exhausted(self, model_name: str | None = None) -> QuotaExhaustedResult:
if not model_name:
return QuotaExhaustedResult(False)
model_quota = self._quota_by_model().get(model_name)
if not isinstance(model_quota, dict):
return QuotaExhaustedResult(False)
remaining_fraction = _to_float(model_quota.get("remaining_fraction"))
if remaining_fraction is not None and remaining_fraction <= 0.0:
return QuotaExhaustedResult(True, f"Antigravity 模型 {model_name} 配额剩余 0%")
if _pct_is_exhausted(model_quota.get("used_percent")):
return QuotaExhaustedResult(True, f"Antigravity 模型 {model_name} 配额剩余 0%")
return QuotaExhaustedResult(False)
def usage_ratio(self) -> float | None:
usage_values: list[float] = []
for model_info in self._quota_by_model().values():
if not isinstance(model_info, dict):
continue
used_percent = self._used_percent(model_info)
if used_percent is None:
continue
usage_values.append(used_percent / 100.0)
if not usage_values:
return None
return sum(usage_values) / len(usage_values)
def plan_type(self) -> str | None:
return None
def reset_seconds(self) -> float | None:
return None
def account_block(self) -> AccountBlockResult:
if not _is_truthy_flag(self._data.get("is_forbidden")):
return AccountBlockResult(blocked=False)
reason = _extract_reason(self._data, "forbidden_reason", "reason", "message")
return AccountBlockResult(
blocked=True,
code="account_forbidden",
label="访问受限",
reason=reason or "Antigravity 账户访问受限",
)
def display_summary(self) -> str | None:
if self._data.get("is_forbidden") is True:
return "访问受限"
remaining_list: list[float] = []
for raw_info in self._quota_by_model().values():
if not isinstance(raw_info, dict):
continue
used_percent = self._used_percent(raw_info)
if used_percent is None:
continue
remaining_list.append(max(0.0, min(100.0 - used_percent, 100.0)))
if not remaining_list:
return None
min_remaining = min(remaining_list)
if len(remaining_list) == 1:
return f"剩余 {_format_percent(min_remaining)}"
return f"最低剩余 {_format_percent(min_remaining)} ({len(remaining_list)} 模型)"
_READER_CLASSES: dict[str, type[PoolQuotaReader]] = {
ProviderType.CODEX: CodexQuotaReader,
ProviderType.KIRO: KiroQuotaReader,
ProviderType.ANTIGRAVITY: AntigravityQuotaReader,
}
def get_quota_reader(provider_type: str | None, upstream_metadata: Any) -> PoolQuotaReader:
"""Return a quota reader for one provider namespace in upstream_metadata."""
normalized_type = normalize_provider_type(provider_type)
reader_cls = _READER_CLASSES.get(normalized_type)
if reader_cls is None or not isinstance(upstream_metadata, dict):
return NullQuotaReader(None)
namespace = reader_cls.namespace
if not namespace:
return NullQuotaReader(None)
data = upstream_metadata.get(namespace)
if not isinstance(data, dict):
return NullQuotaReader(None)
return reader_cls(data)
__all__ = [
"AccountBlockResult",
"AntigravityQuotaReader",
"CodexQuotaReader",
"KiroQuotaReader",
"NullQuotaReader",
"PoolQuotaReader",
"QuotaExhaustedResult",
"get_quota_reader",
]

View File

@@ -1,26 +1,7 @@
from __future__ import annotations
from src.core.provider_types import ProviderType, normalize_provider_type
from src.models.database import ProviderAPIKey
def _pct_is_exhausted(value: object) -> bool:
"""Return True when used_percent indicates 0% remaining."""
try:
pct = float(value) # type: ignore[arg-type]
except (TypeError, ValueError):
return False
# Some upstreams may return values slightly above 100 due to rounding.
return pct >= 100.0 - 1e-6
def _float_or_none(value: object) -> float | None:
try:
if value is None:
return None
return float(value) # type: ignore[arg-type]
except (TypeError, ValueError):
return None
from src.services.provider_keys.quota_reader import get_quota_reader
def is_key_quota_exhausted(
@@ -29,72 +10,8 @@ def is_key_quota_exhausted(
*,
model_name: str,
) -> tuple[bool, str | None]:
"""Check ProviderAPIKey.upstream_metadata quota and decide whether to skip.
"""Check ProviderAPIKey.upstream_metadata quota and decide whether to skip."""
Requirements:
- Kiro: account-level quota. When remaining == 0, skip this key; allow again when remaining > 0.
- Codex: only consider weekly quota + 5H quota.
If either remaining is 0%, skip this key.
- Antigravity: quota is per-model; do not disable the account.
When the requested model's quota is 0%, skip this key.
"""
pt = normalize_provider_type(provider_type)
upstream = getattr(key, "upstream_metadata", None) or {}
if not isinstance(upstream, dict):
return False, None
if pt == ProviderType.KIRO:
kiro_meta = upstream.get("kiro")
if not isinstance(kiro_meta, dict):
return False, None
remaining = _float_or_none(kiro_meta.get("remaining"))
if remaining is not None and remaining <= 0.0:
return True, "Kiro 账号配额剩余 0"
return False, None
if pt == ProviderType.CODEX:
codex_meta = upstream.get("codex")
if not isinstance(codex_meta, dict):
return False, None
weekly_used = codex_meta.get("primary_used_percent")
five_hour_used = codex_meta.get("secondary_used_percent")
exhausted_parts: list[str] = []
if _pct_is_exhausted(weekly_used):
exhausted_parts.append("周限额剩余 0%")
if _pct_is_exhausted(five_hour_used):
exhausted_parts.append("5H 限额剩余 0%")
if exhausted_parts:
return True, "Codex " + "".join(exhausted_parts)
return False, None
if pt == ProviderType.ANTIGRAVITY:
ag_meta = upstream.get("antigravity")
if not isinstance(ag_meta, dict):
return False, None
quota_by_model = ag_meta.get("quota_by_model")
if not isinstance(quota_by_model, dict):
return False, None
model_quota = quota_by_model.get(model_name)
if not isinstance(model_quota, dict):
return False, None
remaining_fraction = _float_or_none(model_quota.get("remaining_fraction"))
if remaining_fraction is not None and remaining_fraction <= 0.0:
return True, f"Antigravity 模型 {model_name} 配额剩余 0%"
if _pct_is_exhausted(model_quota.get("used_percent")):
return True, f"Antigravity 模型 {model_name} 配额剩余 0%"
return False, None
return False, None
reader = get_quota_reader(provider_type, getattr(key, "upstream_metadata", None))
result = reader.is_exhausted(model_name)
return result.exhausted, result.reason

View File

@@ -1,7 +1,7 @@
from __future__ import annotations
import re
from collections.abc import Callable
from collections.abc import Awaitable, Callable
from typing import Any
from uuid import uuid4
@@ -56,6 +56,47 @@ _SENSITIVE_PATTERN = re.compile(
)
async def pool_on_error(
provider: Any,
key: Any,
status_code: int,
cause: Any,
) -> None:
"""Notify the pool manager about an upstream error (health policy)."""
try:
from src.services.provider.pool.config import parse_pool_config
from src.services.provider.pool.health_policy import apply_health_policy
pool_cfg = parse_pool_config(getattr(provider, "config", None))
if pool_cfg is None:
return
error_text = ""
resp_headers: dict[str, str] = {}
if getattr(cause, "response", None) is not None:
try:
error_text = (cause.response.text or "")[:4000]
except Exception:
pass
try:
resp_headers = dict(cause.response.headers)
except Exception:
pass
elif isinstance(getattr(cause, "error_message", None), str):
error_text = str(getattr(cause, "error_message", "") or "")[:4000]
await apply_health_policy(
provider_id=str(provider.id),
key_id=str(key.id),
status_code=status_code,
error_body=error_text,
response_headers=resp_headers,
config=pool_cfg,
)
except Exception:
pass
class TaskService:
"""
Unified task service facade (Phase 3).
@@ -205,6 +246,7 @@ class TaskService:
affinity_key: str | None = None,
create_pending_usage: bool = False,
enable_cache_affinity: bool = False,
is_cancelled: Callable[[], Awaitable[bool]] | None = None,
) -> ExecutionResult:
"""Execute a pre-built candidate set through the unified SYNC runtime."""
from src.services.rate_limit.adaptive_rpm import get_adaptive_rpm_manager
@@ -478,6 +520,7 @@ class TaskService:
candidate_record_map=candidate_record_map,
max_attempts=max_attempts,
execution_error_handler=_handle_exec_err,
is_cancelled=is_cancelled,
)
if result.success:
@@ -687,44 +730,7 @@ class TaskService:
except Exception:
logger.opt(exception=True).debug("Pool on_request_success failed (non-blocking)")
@staticmethod
async def _pool_on_error(
provider: Any,
key: Any,
status_code: int,
cause: Any,
) -> None:
"""Notify the pool manager about an upstream error (health policy)."""
try:
from src.services.provider.pool.config import parse_pool_config
from src.services.provider.pool.health_policy import apply_health_policy
pool_cfg = parse_pool_config(getattr(provider, "config", None))
if pool_cfg is None:
return
error_text = ""
resp_headers: dict[str, str] = {}
if getattr(cause, "response", None) is not None:
try:
error_text = (cause.response.text or "")[:4000]
except Exception:
pass
try:
resp_headers = dict(cause.response.headers)
except Exception:
pass
await apply_health_policy(
provider_id=str(provider.id),
key_id=str(key.id),
status_code=status_code,
error_body=error_text,
response_headers=resp_headers,
config=pool_cfg,
)
except Exception:
pass
_pool_on_error = staticmethod(pool_on_error)
async def _execute_sync_unified(
self,
@@ -1472,11 +1478,13 @@ class TaskService:
if isinstance(cause, EmbeddedErrorException):
error_message = cause.error_message or ""
embedded_status = cause.error_code or 200
embedded_detail = error_message[:200] or cause.error_status or f"code={embedded_status}"
if error_classifier.is_client_error(error_message):
logger.warning(
" [{}] 嵌入式客户端错误继续转移: {}",
" [{}] 嵌入式客户端错误 (HTTP 200, status={}), 继续转移: {}",
request_id,
error_message[:200],
cause.error_status or embedded_status,
embedded_detail,
)
RequestCandidateService.mark_candidate_failed(
db=self.db,
@@ -1488,12 +1496,14 @@ class TaskService:
concurrent_requests=captured_key_concurrent,
extra_data=_proxy_extra,
)
await self._pool_on_error(provider, key, embedded_status, cause)
return "break"
logger.warning(
" [{}] 嵌入式服务端错误尝试重试: {}",
" [{}] 嵌入式服务端错误 (HTTP 200, status={}), 尝试重试: {}",
request_id,
error_message[:200],
cause.error_status or embedded_status,
embedded_detail,
)
RequestCandidateService.mark_candidate_failed(
db=self.db,
@@ -1505,6 +1515,7 @@ class TaskService:
concurrent_requests=captured_key_concurrent,
extra_data=_proxy_extra,
)
await self._pool_on_error(provider, key, embedded_status, cause)
return "continue" if has_retry_left else "break"
if isinstance(cause, httpx.HTTPStatusError):