refactor: 移除 Python 后端源码,全面迁移至 Rust gateway 架构

- 删除全部 Python 源码 (src/) 及 Alembic 迁移脚本,归档至 _deprecated_py_src/
- 重构 Rust gateway ai_pipeline: 拆分 planner/finalize 模块,新增 contracts/adaptation 层
- 重组 handlers 模块为 admin/public/proxy/internal/shared 子模块结构
- 新增 executor 模块,引入 Rust 原生数据库迁移 (aether-data/migrations)
- 简化 CI/Docker 构建流程,移除 base image 二级构建,统一为单一 app image
- 移除 Python 相关基础设施文件 (entrypoint.sh, gunicorn_conf.py, Dockerfile.base)
This commit is contained in:
fawney19
2026-04-03 16:26:16 +08:00
parent 8f26e1a31f
commit 1d9c77522a
868 changed files with 1735 additions and 2433 deletions

View File

@@ -0,0 +1,20 @@
"""Generic Account Pool management for any Provider type.
Re-exports the main public API for convenience.
"""
from src.services.provider.pool.config import (
PoolConfig,
ScoringWeights,
UnschedulableRule,
parse_pool_config,
)
from src.services.provider.pool.manager import PoolManager
__all__ = [
"PoolConfig",
"PoolManager",
"ScoringWeights",
"UnschedulableRule",
"parse_pool_config",
]

View File

@@ -0,0 +1,594 @@
"""Pool account state helpers.
Provides a shared way to classify account-level hard-block states
from upstream metadata and OAuth invalid reasons.
"""
from __future__ import annotations
import re
import time
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] "
OAUTH_REQUEST_FAILED_PREFIX = "[REQUEST_FAILED] "
# -- 按原因细分的关键词组 --
# 封禁类 (suspended / banned)
_KEYWORDS_SUSPENDED: tuple[str, ...] = (
"suspended",
"account_block",
"account blocked",
"封禁",
"封号",
"被封",
"账户已封禁",
"账号异常",
)
# 停用类 (disabled / deactivated)
_KEYWORDS_DISABLED: tuple[str, ...] = (
"account has been disabled",
"account disabled",
"account has been deactivated",
"account_deactivated",
"account deactivated",
"organization has been disabled",
"organization_disabled",
"deactivated_workspace",
"deactivated",
"访问被禁止",
"账户访问被禁止",
)
_TOKEN_INVALID_KEYWORDS: tuple[str, ...] = (
"authentication token has been invalidated",
"token has been invalidated",
"codex token 无效或已过期",
)
# 需要验证类
_KEYWORDS_VERIFICATION: tuple[str, ...] = (
"validation_required",
"verify your account",
"需要验证",
"验证账号",
"验证身份",
)
# 合并的完整列表(用于 is_account_level_block_reason 快速判断)
ACCOUNT_BLOCK_REASON_KEYWORDS: tuple[str, ...] = (
*_KEYWORDS_SUSPENDED,
*_KEYWORDS_DISABLED,
*_TOKEN_INVALID_KEYWORDS,
*_KEYWORDS_VERIFICATION,
)
AUTO_REMOVABLE_ACCOUNT_STATE_CODES: frozenset[str] = frozenset(
{
"account_banned",
"account_suspended",
"account_disabled",
"workspace_deactivated",
"account_forbidden",
}
)
def _classify_block_reason(text: str) -> tuple[str, str]:
"""Return (code, label) based on the oauth_invalid_reason text."""
lowered = text.lower()
if any(kw in lowered for kw in _TOKEN_INVALID_KEYWORDS):
return "oauth_expired", "Token 失效"
if any(kw in lowered for kw in _KEYWORDS_VERIFICATION):
return "account_verification", "需要验证"
if "deactivated_workspace" in lowered:
return "workspace_deactivated", "工作区停用"
if any(kw in lowered for kw in _KEYWORDS_DISABLED):
return "account_disabled", "账号停用"
if any(kw in lowered for kw in _KEYWORDS_SUSPENDED):
return "account_suspended", "账号封禁"
return "account_blocked", "账号异常"
@dataclass(frozen=True, slots=True)
class PoolAccountState:
"""Resolved account-level state for one key."""
blocked: bool
code: str | None = None # account_banned / account_forbidden / account_blocked
label: str | None = None
reason: str | None = None
source: str | None = None # metadata / oauth_invalid / oauth_refresh / oauth_request
recoverable: bool = False
@dataclass(frozen=True, slots=True)
class OAuthStatusSnapshot:
code: str = "none" # none / valid / expiring / expired / invalid / check_failed
label: str | None = None
reason: str | None = None
expires_at: int | None = None
invalid_at: int | None = None
source: str | None = None
requires_reauth: bool = False
expiring_soon: bool = False
@dataclass(frozen=True, slots=True)
class AccountStatusSnapshot:
code: str = "ok"
label: str | None = None
reason: str | None = None
blocked: bool = False
source: str | None = None
recoverable: bool = False
@dataclass(frozen=True, slots=True)
class QuotaStatusSnapshot:
code: str = "unknown" # unknown / ok / exhausted
label: str | None = None
reason: str | None = None
exhausted: bool = False
usage_ratio: float | None = None
updated_at: int | None = None
reset_seconds: float | None = None
plan_type: str | None = None
@dataclass(frozen=True, slots=True)
class ProviderKeyStatusSnapshot:
oauth: OAuthStatusSnapshot
account: AccountStatusSnapshot
quota: QuotaStatusSnapshot
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 _clean_text(value: Any) -> str | None:
if not isinstance(value, str):
return None
text = value.strip()
return text or None
def _extract_reason(source: dict[str, Any] | None, *fields: str) -> str | None:
if not isinstance(source, dict):
return None
for field in fields:
text = _clean_text(source.get(field))
if text:
return text
return None
def _is_workspace_deactivated_reason(reason: str | None) -> bool:
text = _clean_text(reason)
return bool(text and "deactivated_workspace" in text.lower())
_TAGGED_REASON_PATTERN = re.compile(
r"(?:^|\n)\[(?P<tag>[A-Z_]+)\]\s*(?P<detail>.*?)(?=\n\[[A-Z_]+\]|\Z)",
re.S,
)
def _extract_tagged_reason_sections(reason: str | None) -> dict[str, str]:
text = _clean_text(reason)
if not text:
return {}
sections: dict[str, str] = {}
for match in _TAGGED_REASON_PATTERN.finditer(text):
tag = str(match.group("tag") or "").strip().upper()
if not tag or tag in sections:
continue
detail = str(match.group("detail") or "").strip()
sections[tag] = detail
return sections
def _resolve_from_metadata(
provider_type: str | None,
upstream_metadata: Any,
) -> PoolAccountState | None:
if not isinstance(upstream_metadata, dict):
return None
normalized_provider = str(provider_type or "").strip().lower()
provider_bucket: dict[str, Any] | None = None
if normalized_provider:
maybe_bucket = upstream_metadata.get(normalized_provider)
if isinstance(maybe_bucket, dict):
provider_bucket = maybe_bucket
quota_block = get_quota_reader(normalized_provider, upstream_metadata).account_block()
if quota_block.blocked:
return PoolAccountState(
blocked=True,
code=quota_block.code,
label=quota_block.label,
reason=quota_block.reason,
source="metadata",
)
for source in (provider_bucket, upstream_metadata):
if not isinstance(source, dict):
continue
if _is_truthy_flag(source.get("is_banned")):
reason = _extract_reason(source, "ban_reason", "forbidden_reason", "reason", "message")
return PoolAccountState(
blocked=True,
code="account_banned",
label="账号封禁",
reason=reason or "账号已封禁",
source="metadata",
)
if _is_truthy_flag(source.get("is_forbidden")) or _is_truthy_flag(
source.get("account_disabled")
):
reason = _extract_reason(source, "forbidden_reason", "ban_reason", "reason", "message")
if _is_workspace_deactivated_reason(reason):
return PoolAccountState(
blocked=True,
code="workspace_deactivated",
label="工作区停用",
reason=reason or "工作区已停用",
source="metadata",
)
return PoolAccountState(
blocked=True,
code="account_forbidden",
label="访问受限",
reason=reason or "账号访问受限",
source="metadata",
)
return None
def _resolve_from_oauth_invalid_reason(reason: str | None) -> PoolAccountState | None:
text = _clean_text(reason)
if not text:
return None
if text.startswith(OAUTH_ACCOUNT_BLOCK_PREFIX):
cleaned = text[len(OAUTH_ACCOUNT_BLOCK_PREFIX) :].strip()
code, label = (
_classify_block_reason(cleaned) if cleaned else ("account_blocked", "账号异常")
)
return PoolAccountState(
blocked=True,
code=code,
label=label,
reason=cleaned or "账号异常",
source="oauth_invalid",
)
if text.startswith(OAUTH_EXPIRED_PREFIX):
cleaned = text[len(OAUTH_EXPIRED_PREFIX) :].strip()
return PoolAccountState(
blocked=True,
code="oauth_expired",
label="Token 失效",
reason=cleaned or "OAuth Token 已过期且无法续期",
source="oauth_invalid",
recoverable=True,
)
if text.startswith(OAUTH_REFRESH_FAILED_PREFIX):
cleaned = text[len(OAUTH_REFRESH_FAILED_PREFIX) :].strip()
return PoolAccountState(
blocked=False,
code="oauth_refresh_failed",
label="续期失败",
reason=cleaned or "OAuth Token 续期失败",
source="oauth_refresh",
recoverable=True,
)
if text.startswith(OAUTH_REQUEST_FAILED_PREFIX):
cleaned = text[len(OAUTH_REQUEST_FAILED_PREFIX) :].strip()
return PoolAccountState(
blocked=False,
code="oauth_request_failed",
label="请求失败",
reason=cleaned or "账号状态检查失败",
source="oauth_request",
recoverable=True,
)
if text.startswith("["):
return None
lowered = text.lower()
if any(keyword in lowered for keyword in ACCOUNT_BLOCK_REASON_KEYWORDS):
code, label = _classify_block_reason(text)
return PoolAccountState(
blocked=True,
code=code,
label=label,
reason=text,
source="oauth_invalid",
)
return None
def resolve_account_status_snapshot(
*,
provider_type: str | None,
upstream_metadata: Any,
oauth_invalid_reason: str | None,
) -> AccountStatusSnapshot:
from_metadata = _resolve_from_metadata(provider_type, upstream_metadata)
if from_metadata is not None:
return AccountStatusSnapshot(
code=from_metadata.code or "ok",
label=from_metadata.label,
reason=from_metadata.reason,
blocked=from_metadata.blocked,
source=from_metadata.source,
recoverable=from_metadata.recoverable,
)
text = _clean_text(oauth_invalid_reason)
if not text:
return AccountStatusSnapshot()
tagged_sections = _extract_tagged_reason_sections(text)
if "ACCOUNT_BLOCK" in tagged_sections:
cleaned = tagged_sections["ACCOUNT_BLOCK"]
code, label = (
_classify_block_reason(cleaned) if cleaned else ("account_blocked", "账号异常")
)
return AccountStatusSnapshot(
code=code,
label=label,
reason=cleaned or "账号异常",
blocked=True,
source="oauth_invalid",
)
if text.startswith("["):
return AccountStatusSnapshot()
lowered = text.lower()
if any(keyword in lowered for keyword in ACCOUNT_BLOCK_REASON_KEYWORDS):
code, label = _classify_block_reason(text)
return AccountStatusSnapshot(
code=code,
label=label,
reason=text,
blocked=True,
source="oauth_invalid",
)
return AccountStatusSnapshot()
def resolve_oauth_status_snapshot(
*,
auth_type: str | None,
oauth_expires_at: int | None,
oauth_invalid_at: int | None,
oauth_invalid_reason: str | None,
now_ts: int | None = None,
) -> OAuthStatusSnapshot:
if str(auth_type or "").strip().lower() != "oauth":
return OAuthStatusSnapshot()
now = int(now_ts if now_ts is not None else time.time())
invalid_at = int(oauth_invalid_at) if isinstance(oauth_invalid_at, int) else None
tagged_sections = _extract_tagged_reason_sections(oauth_invalid_reason)
raw_reason = _clean_text(oauth_invalid_reason)
expired_reason = tagged_sections.get("OAUTH_EXPIRED")
if expired_reason:
return OAuthStatusSnapshot(
code="invalid",
label="已失效",
reason=expired_reason,
invalid_at=invalid_at,
expires_at=oauth_expires_at,
source="oauth_invalid",
requires_reauth=True,
)
refresh_failed_reason = tagged_sections.get("REFRESH_FAILED")
if refresh_failed_reason:
return OAuthStatusSnapshot(
code="invalid",
label="已失效",
reason=refresh_failed_reason,
invalid_at=invalid_at,
expires_at=oauth_expires_at,
source="oauth_refresh",
requires_reauth=True,
)
request_failed_reason = tagged_sections.get("REQUEST_FAILED")
if request_failed_reason:
return OAuthStatusSnapshot(
code="check_failed",
label="检查失败",
reason=request_failed_reason,
expires_at=oauth_expires_at,
source="oauth_request",
)
account_snapshot = resolve_account_status_snapshot(
provider_type=None,
upstream_metadata=None,
oauth_invalid_reason=raw_reason,
)
if account_snapshot.blocked:
if oauth_expires_at is None:
return OAuthStatusSnapshot()
elif raw_reason or invalid_at is not None:
return OAuthStatusSnapshot(
code="invalid",
label="已失效",
reason=raw_reason,
invalid_at=invalid_at,
expires_at=oauth_expires_at,
source="oauth_invalid",
requires_reauth=True,
)
expires_at = int(oauth_expires_at) if isinstance(oauth_expires_at, int) else None
if expires_at is None:
return OAuthStatusSnapshot()
if expires_at <= now:
return OAuthStatusSnapshot(
code="expired",
label="已过期",
reason="Token 已过期,请重新授权",
expires_at=expires_at,
source="expires_at",
requires_reauth=True,
)
expiring_soon = (expires_at - now) < 24 * 3600
return OAuthStatusSnapshot(
code="expiring" if expiring_soon else "valid",
label="即将过期" if expiring_soon else "有效",
expires_at=expires_at,
source="expires_at",
expiring_soon=expiring_soon,
)
def resolve_quota_status_snapshot(
*,
provider_type: str | None,
upstream_metadata: Any,
) -> QuotaStatusSnapshot:
normalized_provider = str(provider_type or "").strip().lower()
reader = get_quota_reader(normalized_provider, upstream_metadata)
quota_state = reader.is_exhausted()
usage_ratio = reader.usage_ratio()
updated_at = reader.updated_at()
reset_seconds = reader.reset_seconds()
plan_type = reader.plan_type()
if quota_state.exhausted:
return QuotaStatusSnapshot(
code="exhausted",
label="额度耗尽",
reason=quota_state.reason,
exhausted=True,
usage_ratio=usage_ratio,
updated_at=updated_at,
reset_seconds=reset_seconds,
plan_type=plan_type,
)
if any(value is not None for value in (usage_ratio, updated_at, reset_seconds, plan_type)):
return QuotaStatusSnapshot(
code="ok",
exhausted=False,
usage_ratio=usage_ratio,
updated_at=updated_at,
reset_seconds=reset_seconds,
plan_type=plan_type,
)
return QuotaStatusSnapshot()
def build_provider_key_status_snapshot(
*,
auth_type: str | None,
oauth_expires_at: int | None,
oauth_invalid_at: int | None,
oauth_invalid_reason: str | None,
provider_type: str | None,
upstream_metadata: Any,
now_ts: int | None = None,
) -> ProviderKeyStatusSnapshot:
account = resolve_account_status_snapshot(
provider_type=provider_type,
upstream_metadata=upstream_metadata,
oauth_invalid_reason=oauth_invalid_reason,
)
oauth = resolve_oauth_status_snapshot(
auth_type=auth_type,
oauth_expires_at=oauth_expires_at,
oauth_invalid_at=oauth_invalid_at,
oauth_invalid_reason=oauth_invalid_reason,
now_ts=now_ts,
)
quota = resolve_quota_status_snapshot(
provider_type=provider_type,
upstream_metadata=upstream_metadata,
)
return ProviderKeyStatusSnapshot(oauth=oauth, account=account, quota=quota)
def resolve_pool_account_state(
*,
provider_type: str | None,
upstream_metadata: Any,
oauth_invalid_reason: str | None,
) -> PoolAccountState:
"""Resolve account-level hard-block state for pool scheduling."""
from_metadata = _resolve_from_metadata(provider_type, upstream_metadata)
if from_metadata is not None:
return from_metadata
from_oauth = _resolve_from_oauth_invalid_reason(oauth_invalid_reason)
if from_oauth is not None:
return from_oauth
return PoolAccountState(blocked=False)
def should_auto_remove_account_state(state: PoolAccountState) -> bool:
"""Whether a resolved account state is safe to auto-remove.
Auto-removal is limited to hard, non-recoverable account abnormalities.
Pure token failures (`oauth_expired`, `oauth_refresh_failed`) and
softer/manual-recoverable states like `account_verification` are excluded.
"""
return bool(
state.blocked
and not state.recoverable
and str(state.code or "").strip().lower() in AUTO_REMOVABLE_ACCOUNT_STATE_CODES
)
__all__ = [
"ACCOUNT_BLOCK_REASON_KEYWORDS",
"AUTO_REMOVABLE_ACCOUNT_STATE_CODES",
"AccountStatusSnapshot",
"OAUTH_ACCOUNT_BLOCK_PREFIX",
"OAUTH_EXPIRED_PREFIX",
"OAUTH_REFRESH_FAILED_PREFIX",
"OAUTH_REQUEST_FAILED_PREFIX",
"OAuthStatusSnapshot",
"PoolAccountState",
"ProviderKeyStatusSnapshot",
"QuotaStatusSnapshot",
"build_provider_key_status_snapshot",
"resolve_account_status_snapshot",
"resolve_oauth_status_snapshot",
"resolve_pool_account_state",
"resolve_quota_status_snapshot",
"should_auto_remove_account_state",
]

View File

@@ -0,0 +1,341 @@
"""Account Pool configuration (provider-agnostic)."""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any
from src.core.logger import logger
from src.services.provider.pool.dimensions import get_preset_dimension, get_preset_names
@dataclass(frozen=True, slots=True)
class ScoringWeights:
"""Weights used by multi-score scheduling."""
lru: float = 0.3
latency: float = 0.25
health: float = 0.2
cost_remaining: float = 0.25
@dataclass(frozen=True, slots=True)
class SchedulingPreset:
"""Single scheduling preset item with enable/disable and optional sub-config."""
preset: str
enabled: bool = True
mode: str | None = None
@dataclass(frozen=True, slots=True)
class UnschedulableRule:
"""Keyword-based temporary unschedule rule."""
keyword: str
duration_minutes: int = 5
@dataclass(frozen=True, slots=True)
class PoolConfig:
"""Parsed pool configuration for any Provider.
All transient state lives in Redis; this dataclass only holds
the *configuration* that controls pool behaviour.
"""
# -- Sticky Session -------------------------------------------------------
sticky_session_ttl_seconds: int = 3600 # 1 hour
# Key 优先模式下号池整体优先级None 时回退 provider_priority
global_priority: int | None = None
# -- Load-Aware Selection -------------------------------------------------
load_threshold_percent: int = 80
# -- Scheduling (unified preset list) -------------------------------------
scheduling_presets: tuple[SchedulingPreset, ...] = (
SchedulingPreset(preset="cache_affinity", enabled=True),
)
# Derived from scheduling_presets at parse time (backward compat for consumers)
lru_enabled: bool = True
scheduling_mode: str = "lru" # lru | multi_score
scoring_weights: ScoringWeights = field(default_factory=ScoringWeights)
latency_window_seconds: int = 3600
latency_sample_limit: int = 50
# -- Rolling-Window Cost Tracking -----------------------------------------
cost_window_seconds: int = 18000 # 5 hours
cost_limit_per_key_tokens: int | None = None # None = unlimited
cost_soft_threshold_percent: int = 80
# -- Cooldown Defaults ----------------------------------------------------
rate_limit_cooldown_seconds: int = 300 # 429
overload_cooldown_seconds: int = 30 # 529
# -- OAuth Proactive Refresh ----------------------------------------------
proactive_refresh_seconds: int = 180 # 3 minutes before expiry
# -- Health Policy --------------------------------------------------------
health_policy_enabled: bool = True
# -- Temporary Unschedulable Rules ----------------------------------------
unschedulable_rules: list[UnschedulableRule] = field(default_factory=list)
# -- Batch Operations -----------------------------------------------------
batch_concurrency: int = 8
# -- Quota Probing --------------------------------------------------------
probing_enabled: bool = False
probing_interval_minutes: int = 10
auto_remove_banned_keys: bool = False
# -- Stream Timeout Auto-Pause --------------------------------------------
stream_timeout_threshold: int = 3 # N timeouts within window trigger cooldown
stream_timeout_window_seconds: int = 1800 # 30 min counting window
stream_timeout_cooldown_seconds: int = 300 # 5 min cooldown
# -- Pluggable Strategies -------------------------------------------------
strategies: tuple[str, ...] = ()
def parse_pool_config(provider_config: Any) -> PoolConfig | None:
"""Parse PoolConfig from ``Provider.config``.
Only looks for the explicit ``pool_advanced`` key. Returns ``None``
when the provider has no pool section configured, meaning the caller
should use the normal (non-pool) scheduling path.
"""
config_dict = provider_config if isinstance(provider_config, dict) else {}
raw_advanced = config_dict.get("pool_advanced")
if raw_advanced is None:
return None
if not isinstance(raw_advanced, dict):
# Could be a pre-validated Pydantic model; grab its dict.
try:
raw_advanced = raw_advanced.model_dump() # type: ignore[union-attr]
except Exception:
logger.warning(
"PoolConfig: advanced config type invalid ({}), falling back to defaults",
type(raw_advanced).__name__,
)
return PoolConfig()
rules: list[UnschedulableRule] = []
raw_rules = raw_advanced.get("unschedulable_rules")
if isinstance(raw_rules, list):
for r in raw_rules:
if isinstance(r, dict) and isinstance(r.get("keyword"), str):
rules.append(
UnschedulableRule(
keyword=r["keyword"],
duration_minutes=int(r.get("duration_minutes", 5)),
)
)
def _int_or(key: str, default: int) -> int:
v = raw_advanced.get(key)
if v is None:
return default
try:
return int(v)
except (TypeError, ValueError):
return default
def _bool_or(key: str, default: bool) -> bool:
v = raw_advanced.get(key)
if v is None:
return default
return bool(v)
def _opt_int(key: str) -> int | None:
v = raw_advanced.get(key)
if v is None:
return None
try:
return int(v)
except (TypeError, ValueError):
return None
scoring_weights = _parse_scoring_weights(raw_advanced.get("scoring_weights"))
# Parse scheduling presets (new object-list format or legacy string-list)
presets = _parse_scheduling_presets_v2(
raw_advanced.get("scheduling_presets"),
legacy_mode=raw_advanced.get("scheduling_mode"),
legacy_lru=raw_advanced.get("lru_enabled"),
)
# Derive scheduling_mode and lru_enabled from the presets list
enabled = [p for p in presets if p.enabled]
lru_enabled = any(p.preset == "lru" for p in enabled)
non_lru_enabled = [p for p in enabled if p.preset != "lru"]
scheduling_mode = "multi_score" if non_lru_enabled else "lru"
strategies = list(_parse_strategies(raw_advanced.get("strategies")))
if scheduling_mode == "multi_score" and "multi_score" not in strategies:
strategies.append("multi_score")
return PoolConfig(
sticky_session_ttl_seconds=_int_or("sticky_session_ttl_seconds", 3600),
global_priority=_opt_int("global_priority"),
load_threshold_percent=_int_or("load_threshold_percent", 80),
scheduling_presets=presets,
lru_enabled=lru_enabled,
scheduling_mode=scheduling_mode,
scoring_weights=scoring_weights,
latency_window_seconds=_int_or("latency_window_seconds", 3600),
latency_sample_limit=_int_or("latency_sample_limit", 50),
cost_window_seconds=_int_or("cost_window_seconds", 18000),
cost_limit_per_key_tokens=_opt_int("cost_limit_per_key_tokens"),
cost_soft_threshold_percent=_int_or("cost_soft_threshold_percent", 80),
rate_limit_cooldown_seconds=_int_or("rate_limit_cooldown_seconds", 300),
overload_cooldown_seconds=_int_or("overload_cooldown_seconds", 30),
proactive_refresh_seconds=_int_or("proactive_refresh_seconds", 180),
health_policy_enabled=_bool_or("health_policy_enabled", True),
unschedulable_rules=rules,
batch_concurrency=max(1, min(_int_or("batch_concurrency", 8), 32)),
probing_enabled=_bool_or("probing_enabled", False),
probing_interval_minutes=max(1, min(_int_or("probing_interval_minutes", 10), 1440)),
auto_remove_banned_keys=_bool_or("auto_remove_banned_keys", False),
stream_timeout_threshold=_int_or("stream_timeout_threshold", 3),
stream_timeout_window_seconds=_int_or("stream_timeout_window_seconds", 1800),
stream_timeout_cooldown_seconds=_int_or("stream_timeout_cooldown_seconds", 300),
strategies=tuple(strategies),
)
# ---------------------------------------------------------------------------
# Internal parsers
# ---------------------------------------------------------------------------
def _allowed_preset_names() -> set[str]:
return get_preset_names() | {"lru"}
def _get_preset_mode_meta(preset_name: str) -> tuple[tuple[str, ...], str | None]:
dim = get_preset_dimension(preset_name)
if dim is None or not dim.modes:
return (), None
modes = tuple(str(mode).strip().lower() for mode in dim.modes if str(mode).strip())
if not modes:
return (), None
raw_default = str(dim.default_mode or "").strip().lower()
default_mode = raw_default if raw_default in modes else modes[0]
return modes, default_mode
def _parse_strategies(raw: Any) -> tuple[str, ...]:
"""Parse strategy names from config (list[str] -> tuple[str, ...])."""
if not isinstance(raw, list):
return ()
return tuple(str(s) for s in raw if isinstance(s, str) and s)
def _parse_scoring_weights(raw: Any) -> ScoringWeights:
"""Parse scoring weights with graceful fallback."""
if not isinstance(raw, dict):
return ScoringWeights()
def _float_or(value: Any, default: float) -> float:
try:
parsed = float(value)
except (TypeError, ValueError):
return default
return max(0.0, min(parsed, 1.0))
return ScoringWeights(
lru=_float_or(raw.get("lru"), 0.3),
latency=_float_or(raw.get("latency"), 0.25),
health=_float_or(raw.get("health"), 0.2),
cost_remaining=_float_or(raw.get("cost_remaining"), 0.25),
)
def _parse_scheduling_presets_v2(
raw: Any,
*,
legacy_mode: Any = None,
legacy_lru: Any = None,
) -> tuple[SchedulingPreset, ...]:
"""Parse scheduling presets, supporting both new and legacy formats.
New format::
[{"preset": "lru", "enabled": true},
{"preset": "free_team_first", "enabled": true, "mode": "free_only"},
...]
Legacy format::
["free_team_first", "recent_refresh"] (with separate scheduling_mode / lru_enabled)
"""
if isinstance(raw, list) and raw:
first = raw[0]
if isinstance(first, dict):
return _parse_preset_object_list(raw)
if isinstance(first, str):
return _convert_legacy_string_list(raw, legacy_mode, legacy_lru)
# No presets at all: derive from legacy fields
return _build_from_legacy_fields(legacy_mode, legacy_lru)
def _parse_preset_object_list(raw: list[Any]) -> tuple[SchedulingPreset, ...]:
"""Parse new-format object list into SchedulingPreset tuple."""
allowed = _allowed_preset_names()
ordered: list[SchedulingPreset] = []
seen: set[str] = set()
for item in raw:
if not isinstance(item, dict):
continue
name = str(item.get("preset", "")).strip().lower()
if name not in allowed or name in seen:
continue
seen.add(name)
enabled = bool(item.get("enabled", True))
mode: str | None = None
modes, default_mode = _get_preset_mode_meta(name)
if modes:
raw_mode = str(item.get("mode", default_mode) or "").strip().lower()
mode = raw_mode if raw_mode in modes else default_mode
ordered.append(SchedulingPreset(preset=name, enabled=enabled, mode=mode))
return tuple(ordered) if ordered else (SchedulingPreset(preset="lru", enabled=True),)
def _convert_legacy_string_list(
raw: list[Any],
legacy_mode: Any,
legacy_lru: Any,
) -> tuple[SchedulingPreset, ...]:
"""Convert legacy string list + mode/lru fields to new format."""
lru_enabled = legacy_lru if isinstance(legacy_lru, bool) else True
allowed_non_lru = _allowed_preset_names() - {"lru"}
items: list[SchedulingPreset] = [SchedulingPreset(preset="lru", enabled=lru_enabled)]
seen: set[str] = {"lru"}
for p in raw:
if not isinstance(p, str):
continue
name = p.strip().lower()
if name not in allowed_non_lru or name in seen:
continue
seen.add(name)
items.append(SchedulingPreset(preset=name, enabled=True))
return tuple(items)
def _build_from_legacy_fields(legacy_mode: Any, legacy_lru: Any) -> tuple[SchedulingPreset, ...]:
"""Build presets from legacy scheduling_mode / lru_enabled only."""
# Explicit lru_enabled=True -> LRU; explicit lru_enabled=False -> cache_affinity.
# No legacy fields at all -> default to cache_affinity.
if isinstance(legacy_lru, bool):
if legacy_lru:
return (SchedulingPreset(preset="lru", enabled=True),)
return (SchedulingPreset(preset="cache_affinity", enabled=True),)
return (SchedulingPreset(preset="cache_affinity", enabled=True),)

View File

@@ -0,0 +1,66 @@
"""Rolling-window cost tracking for the Account Pool.
Each key has a configurable token budget per rolling window (e.g. 5 hours).
When the budget is exhausted the key is marked as unschedulable by the pool
manager. A "soft threshold" (default 80 %) causes the pool to *prefer*
other keys but still allows traffic if no alternatives exist.
All state is stored in Redis sorted sets via :mod:`redis_ops`.
"""
from __future__ import annotations
from typing import TYPE_CHECKING
from src.services.provider.pool import redis_ops
if TYPE_CHECKING:
from src.services.provider.pool.config import PoolConfig
async def record_usage(
provider_id: str,
key_id: str,
tokens: int,
config: PoolConfig,
) -> None:
"""Record *tokens* used by *key_id* in the rolling cost window."""
if tokens <= 0:
return
if config.cost_limit_per_key_tokens is None:
return # cost tracking disabled
await redis_ops.add_cost_entry(provider_id, key_id, tokens, config.cost_window_seconds)
async def get_window_usage(
provider_id: str,
key_id: str,
config: PoolConfig,
) -> int:
"""Return total tokens used by *key_id* within the current window."""
return await redis_ops.get_cost_window_total(provider_id, key_id, config.cost_window_seconds)
async def is_at_limit(
provider_id: str,
key_id: str,
config: PoolConfig,
) -> bool:
"""Return ``True`` if the key has exhausted its budget."""
if config.cost_limit_per_key_tokens is None:
return False
total = await get_window_usage(provider_id, key_id, config)
return total >= config.cost_limit_per_key_tokens
async def is_approaching_limit(
provider_id: str,
key_id: str,
config: PoolConfig,
) -> bool:
"""Return ``True`` if the key is above the soft threshold."""
if config.cost_limit_per_key_tokens is None:
return False
total = await get_window_usage(provider_id, key_id, config)
threshold = config.cost_limit_per_key_tokens * config.cost_soft_threshold_percent / 100
return total >= threshold

View File

@@ -0,0 +1,39 @@
"""Pool scheduling preset dimensions.
Importing this package registers all built-in preset dimensions.
"""
from __future__ import annotations
from . import cache_affinity # noqa: F401
from . import cost_first # noqa: F401
from . import free_first # noqa: F401
from . import free_team_first # noqa: F401
from . import health_first # noqa: F401
from . import latency_first # noqa: F401
from . import load_balance # noqa: F401
from . import plus_first # noqa: F401
from . import priority_first # noqa: F401
from . import quota_balanced # noqa: F401
from . import recent_refresh # noqa: F401
from . import single_account # noqa: F401
from . import team_first # noqa: F401
from .registry import (
PresetDimensionBase,
PresetDimensionMeta,
get_all_preset_dimensions,
get_preset_dimension,
get_preset_dimension_metas,
get_preset_names,
register_preset_dimension,
)
__all__ = [
"PresetDimensionBase",
"PresetDimensionMeta",
"get_all_preset_dimensions",
"get_preset_dimension",
"get_preset_dimension_metas",
"get_preset_names",
"register_preset_dimension",
]

View File

@@ -0,0 +1,284 @@
"""Shared helpers for pool preset dimensions."""
from __future__ import annotations
import math
import time
from typing import Any
from src.core.provider_types import ProviderType, normalize_provider_type
from src.services.provider_keys.quota_reader import get_quota_reader
def safe_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 safe_metadata(key_obj: Any) -> dict[str, Any]:
raw = getattr(key_obj, "upstream_metadata", None)
return raw if isinstance(raw, dict) else {}
def normalize_plan(value: Any) -> str | None:
if not isinstance(value, str):
return None
normalized = value.strip().lower()
return normalized or None
def rank_ascending(key_id: str, scores: dict[str, float], all_ids: list[str]) -> float:
"""Rank score within all IDs; lower value means better rank."""
if not all_ids:
return 0.0
valid_count = sum(1 for kid in all_ids if safe_float(scores.get(kid)) is not None)
if valid_count <= 0:
return 0.5
decorated: list[tuple[int, float, int, str]] = []
for idx, kid in enumerate(all_ids):
score_raw = safe_float(scores.get(kid))
if score_raw is None:
decorated.append((1, float("inf"), idx, kid))
else:
decorated.append((0, score_raw, idx, kid))
decorated.sort(key=lambda item: (item[0], item[1], item[2]))
rank_idx = 0
for idx, (_missing, _value, _order, kid) in enumerate(decorated):
if kid == key_id:
rank_idx = idx
break
n = len(all_ids)
if n <= 1:
return 0.0
return rank_idx / float(n - 1)
def rank_descending(key_id: str, scores: dict[str, float], all_ids: list[str]) -> float:
"""Rank score within all IDs; higher value means better rank."""
if not all_ids:
return 0.0
valid_count = sum(1 for kid in all_ids if safe_float(scores.get(kid)) is not None)
if valid_count <= 0:
return 0.5
decorated: list[tuple[int, float, int, str]] = []
for idx, kid in enumerate(all_ids):
score_raw = safe_float(scores.get(kid))
if score_raw is None:
decorated.append((1, float("inf"), idx, kid))
else:
# 排序时取负值使分值越大排名越靠前rank 越小)
decorated.append((0, -score_raw, idx, kid))
decorated.sort(key=lambda item: (item[0], item[1], item[2]))
rank_idx = 0
for idx, (_missing, _value, _order, kid) in enumerate(decorated):
if kid == key_id:
rank_idx = idx
break
n = len(all_ids)
if n <= 1:
return 0.0
return rank_idx / float(n - 1)
def extract_plan_type(key_obj: Any) -> str | None:
direct = normalize_plan(getattr(key_obj, "oauth_plan_type", None))
if direct:
return direct
metadata = safe_metadata(key_obj)
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
def _resolve_key_provider_type(key_obj: Any, provider_type: str | None = None) -> str | None:
explicit = normalize_provider_type(provider_type)
if explicit:
return explicit
direct = normalize_provider_type(getattr(key_obj, "provider_type", None))
if direct:
return direct
provider = getattr(key_obj, "provider", None)
related = normalize_provider_type(getattr(provider, "provider_type", None))
if related:
return related
metadata = safe_metadata(key_obj)
candidates = [
provider.value
for provider in (ProviderType.CODEX, ProviderType.KIRO, ProviderType.ANTIGRAVITY)
if isinstance(metadata.get(provider.value), dict)
]
if len(candidates) == 1:
return candidates[0]
return None
def _extract_codex_weekly_reset_seconds(metadata: dict[str, Any]) -> float | None:
codex = metadata.get(ProviderType.CODEX.value)
if not isinstance(codex, dict):
return None
weekly_used_percent = safe_float(codex.get("primary_used_percent"))
if weekly_used_percent is not None:
clamped_used = max(0.0, min(weekly_used_percent, 100.0))
if clamped_used <= 1e-6:
# 周额度仍为满额时,不启用周窗口重置倒计时。
return None
now = time.time()
# 优先绝对时间戳,避免 reset_seconds 快照随时间漂移。
reset_at = safe_float(codex.get("primary_reset_at"))
if reset_at is not None and reset_at > 0:
remaining = reset_at - now
return remaining if remaining > 0 else 0.0
reset_seconds = safe_float(codex.get("primary_reset_seconds"))
if reset_seconds is None or reset_seconds < 0:
return None
updated_at = safe_float(codex.get("updated_at"))
if updated_at is not None and updated_at > 0:
# 时钟偏移下 updated_at 可能晚于当前时间elapsed 需要下限钳制到 0。
elapsed = max(now - updated_at, 0.0)
return max(reset_seconds - elapsed, 0.0)
return reset_seconds
def extract_reset_seconds(key_obj: Any, provider_type: str | None = None) -> float | None:
metadata = safe_metadata(key_obj)
resolved_provider_type = _resolve_key_provider_type(key_obj, provider_type)
if resolved_provider_type == ProviderType.CODEX:
# Codex metadata 已统一约定primary_* 表示周限额secondary_* 表示 5H 限额。
return _extract_codex_weekly_reset_seconds(metadata)
if resolved_provider_type in (ProviderType.KIRO, ProviderType.ANTIGRAVITY):
return get_quota_reader(resolved_provider_type, metadata).reset_seconds()
candidates: list[float] = []
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
return min(candidates)
def extract_usage_ratio(key_obj: Any) -> float | None:
metadata = safe_metadata(key_obj)
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
def extract_internal_priority(key_obj: Any) -> int:
raw = getattr(key_obj, "internal_priority", None)
parsed = safe_float(raw)
if parsed is None:
return 999999
return max(0, int(parsed))
def extract_health_score(key_obj: Any) -> float | None:
direct = safe_float(getattr(key_obj, "health_score", None))
if direct is not None:
return max(0.0, min(direct, 1.0))
health_by_format = getattr(key_obj, "health_by_format", None)
if not isinstance(health_by_format, dict) or not health_by_format:
return None
scores: list[float] = []
for payload in health_by_format.values():
if not isinstance(payload, dict):
continue
score = safe_float(payload.get("health_score"))
if score is None:
continue
scores.append(max(0.0, min(score, 1.0)))
if not scores:
return None
return min(scores)
def plan_priority_score(plan_type: str | None, mode: str | None = None) -> float:
"""Score a key based on plan type and scheduling mode.
Lower score = higher priority.
"""
effective_mode = (mode or "both").strip().lower()
if effective_mode == "free_only":
if plan_type == "free":
return 0.0
if plan_type == "team":
return 0.5
elif effective_mode == "team_only":
if plan_type == "team":
return 0.0
if plan_type == "free":
return 0.5
elif effective_mode == "plus_only":
if plan_type in {"plus", "pro"}:
return 0.0
if plan_type in {"enterprise", "business"}:
return 0.3
else:
# "both" or unrecognized -> original behavior
if plan_type in {"free", "team"}:
return 0.0
if plan_type in {"enterprise", "business"}:
return 0.2
if plan_type in {"plus", "pro"}:
return 0.6
if plan_type:
return 0.7
return 0.8
__all__ = [
"extract_health_score",
"extract_internal_priority",
"extract_plan_type",
"extract_reset_seconds",
"extract_usage_ratio",
"normalize_plan",
"plan_priority_score",
"rank_ascending",
"rank_descending",
"safe_float",
"safe_metadata",
]

View File

@@ -0,0 +1,45 @@
"""cache_affinity preset dimension."""
from __future__ import annotations
from typing import Any
from ._helpers import rank_descending
from .registry import PresetDimensionBase, register_preset_dimension
class CacheAffinityDimension(PresetDimensionBase):
@property
def name(self) -> str:
return "cache_affinity"
@property
def label(self) -> str:
return "缓存亲和"
@property
def description(self) -> str:
return "优先复用最近使用过的 Key利用 Prompt Caching"
@property
def mutex_group(self) -> str | None:
return "distribution_mode"
@property
def evidence_hint(self) -> str | None:
return "依据 LRU 时间戳(最近使用优先,与 LRU 轮转相反)"
def compute_metric(
self,
*,
key_id: str,
all_key_ids: list[str],
keys_by_id: dict[str, Any],
lru_scores: dict[str, Any],
context: dict[str, Any],
mode: str | None,
) -> float:
return rank_descending(key_id, lru_scores, all_key_ids)
register_preset_dimension(CacheAffinityDimension())

View File

@@ -0,0 +1,62 @@
"""cost_first preset dimension."""
from __future__ import annotations
from typing import Any
from ._helpers import extract_usage_ratio, rank_ascending, safe_float
from .registry import PresetDimensionBase, register_preset_dimension
class CostFirstDimension(PresetDimensionBase):
@property
def name(self) -> str:
return "cost_first"
@property
def label(self) -> str:
return "成本优先"
@property
def description(self) -> str:
return "优先选择窗口消耗更低的账号"
@property
def evidence_hint(self) -> str | None:
return "依据窗口成本/Token 用量,缺失时回退配额使用率"
def compute_metric(
self,
*,
key_id: str,
all_key_ids: list[str],
keys_by_id: dict[str, Any],
lru_scores: dict[str, Any],
context: dict[str, Any],
mode: str | None,
) -> float:
cost_totals = context.get("cost_totals")
if not isinstance(cost_totals, dict):
cost_totals = {}
cost_limit = safe_float(context.get("cost_limit_per_key_tokens"))
cost_scores: dict[str, float] = {}
for kid in all_key_ids:
used = safe_float(cost_totals.get(kid))
if used is not None and used >= 0:
if cost_limit is not None and cost_limit > 0:
cost_scores[kid] = max(0.0, min(used / cost_limit, 1.0))
else:
cost_scores[kid] = min(1.0, used / (used + 10000.0))
continue
usage_ratio = extract_usage_ratio(keys_by_id.get(kid))
if usage_ratio is not None:
cost_scores[kid] = usage_ratio
if not cost_scores:
return rank_ascending(key_id, lru_scores, all_key_ids)
return rank_ascending(key_id, cost_scores, all_key_ids)
register_preset_dimension(CostFirstDimension())

View File

@@ -0,0 +1,51 @@
"""free_first preset dimension."""
from __future__ import annotations
from typing import Any
from ._helpers import extract_plan_type, plan_priority_score, rank_ascending
from .registry import PresetDimensionBase, register_preset_dimension
class FreeFirstDimension(PresetDimensionBase):
@property
def name(self) -> str:
return "free_first"
@property
def label(self) -> str:
return "Free 优先"
@property
def description(self) -> str:
return "优先消耗 Free 账号(依赖 plan_type"
@property
def evidence_hint(self) -> str | None:
return "依据 plan_typeFree 账号优先调度)"
@property
def providers(self) -> tuple[str, ...]:
return ("codex", "kiro")
def compute_metric(
self,
*,
key_id: str,
all_key_ids: list[str],
keys_by_id: dict[str, Any],
lru_scores: dict[str, Any],
context: dict[str, Any],
mode: str | None,
) -> float:
plan_scores = {
kid: plan_priority_score(extract_plan_type(keys_by_id.get(kid)), "free_only")
for kid in all_key_ids
}
if len(set(plan_scores.values())) <= 1:
return rank_ascending(key_id, lru_scores, all_key_ids)
return rank_ascending(key_id, plan_scores, all_key_ids)
register_preset_dimension(FreeFirstDimension())

View File

@@ -0,0 +1,63 @@
"""free_team_first preset dimension."""
from __future__ import annotations
from typing import Any
from ._helpers import extract_plan_type, plan_priority_score, rank_ascending
from .registry import PresetDimensionBase, register_preset_dimension
class FreeTeamFirstDimension(PresetDimensionBase):
@property
def name(self) -> str:
return "free_team_first"
@property
def label(self) -> str:
return "Free/Team 优先"
@property
def description(self) -> str:
return "优先消耗低档账号(依赖 plan_type"
@property
def evidence_hint(self) -> str | None:
return "依据 plan_typeoauth_plan_type 或 upstream_metadata"
@property
def providers(self) -> tuple[str, ...]:
return ("codex", "kiro")
@property
def modes(self) -> tuple[str, ...] | None:
return ("free_only", "team_only", "both")
@property
def default_mode(self) -> str | None:
return "both"
@property
def hidden(self) -> bool:
return True
def compute_metric(
self,
*,
key_id: str,
all_key_ids: list[str],
keys_by_id: dict[str, Any],
lru_scores: dict[str, Any],
context: dict[str, Any],
mode: str | None,
) -> float:
plan_scores = {
kid: plan_priority_score(extract_plan_type(keys_by_id.get(kid)), mode)
for kid in all_key_ids
}
if len(set(plan_scores.values())) <= 1:
return rank_ascending(key_id, lru_scores, all_key_ids)
return rank_ascending(key_id, plan_scores, all_key_ids)
register_preset_dimension(FreeTeamFirstDimension())

View File

@@ -0,0 +1,57 @@
"""health_first preset dimension."""
from __future__ import annotations
from typing import Any
from ._helpers import extract_health_score, rank_ascending, safe_float
from .registry import PresetDimensionBase, register_preset_dimension
class HealthFirstDimension(PresetDimensionBase):
@property
def name(self) -> str:
return "health_first"
@property
def label(self) -> str:
return "健康优先"
@property
def description(self) -> str:
return "优先选择健康分更高、失败更少的账号"
@property
def evidence_hint(self) -> str | None:
return "依据 health_by_format 聚合分(含熔断/失败衰减)"
def compute_metric(
self,
*,
key_id: str,
all_key_ids: list[str],
keys_by_id: dict[str, Any],
lru_scores: dict[str, Any],
context: dict[str, Any],
mode: str | None,
) -> float:
health_scores_ctx = context.get("health_scores")
if not isinstance(health_scores_ctx, dict):
health_scores_ctx = {}
penalty_scores: dict[str, float] = {}
for kid in all_key_ids:
score = safe_float(health_scores_ctx.get(kid))
if score is None:
score = extract_health_score(keys_by_id.get(kid))
if score is None:
continue
normalized = max(0.0, min(score, 1.0))
penalty_scores[kid] = 1.0 - normalized
if not penalty_scores:
return rank_ascending(key_id, lru_scores, all_key_ids)
return rank_ascending(key_id, penalty_scores, all_key_ids)
register_preset_dimension(HealthFirstDimension())

View File

@@ -0,0 +1,54 @@
"""latency_first preset dimension."""
from __future__ import annotations
from typing import Any
from ._helpers import rank_ascending, safe_float
from .registry import PresetDimensionBase, register_preset_dimension
class LatencyFirstDimension(PresetDimensionBase):
@property
def name(self) -> str:
return "latency_first"
@property
def label(self) -> str:
return "延迟优先"
@property
def description(self) -> str:
return "优先选择最近延迟更低的账号"
@property
def evidence_hint(self) -> str | None:
return "依据号池延迟窗口均值latency_window_seconds"
def compute_metric(
self,
*,
key_id: str,
all_key_ids: list[str],
keys_by_id: dict[str, Any],
lru_scores: dict[str, Any],
context: dict[str, Any],
mode: str | None,
) -> float:
latency_avgs = context.get("latency_avgs")
if not isinstance(latency_avgs, dict):
latency_avgs = {}
latency_scores: dict[str, float] = {}
for kid in all_key_ids:
latency = safe_float(latency_avgs.get(kid))
if latency is None or latency < 0:
continue
latency_scores[kid] = latency
if not latency_scores:
return rank_ascending(key_id, lru_scores, all_key_ids)
return rank_ascending(key_id, latency_scores, all_key_ids)
register_preset_dimension(LatencyFirstDimension())

View File

@@ -0,0 +1,45 @@
"""load_balance preset dimension."""
from __future__ import annotations
import random
from typing import Any
from .registry import PresetDimensionBase, register_preset_dimension
class LoadBalanceDimension(PresetDimensionBase):
@property
def name(self) -> str:
return "load_balance"
@property
def label(self) -> str:
return "负载均衡"
@property
def description(self) -> str:
return "随机分散 Key 使用,均匀分摊负载"
@property
def mutex_group(self) -> str | None:
return "distribution_mode"
@property
def evidence_hint(self) -> str | None:
return "每次随机分值,实现完全均匀分散"
def compute_metric(
self,
*,
key_id: str,
all_key_ids: list[str],
keys_by_id: dict[str, Any],
lru_scores: dict[str, Any],
context: dict[str, Any],
mode: str | None,
) -> float:
return random.random()
register_preset_dimension(LoadBalanceDimension())

View File

@@ -0,0 +1,51 @@
"""plus_first preset dimension."""
from __future__ import annotations
from typing import Any
from ._helpers import extract_plan_type, plan_priority_score, rank_ascending
from .registry import PresetDimensionBase, register_preset_dimension
class PlusFirstDimension(PresetDimensionBase):
@property
def name(self) -> str:
return "plus_first"
@property
def label(self) -> str:
return "Plus 优先"
@property
def description(self) -> str:
return "优先消耗 Plus/Pro 账号(依赖 plan_type"
@property
def evidence_hint(self) -> str | None:
return "依据 plan_typePlus/Pro 账号优先调度)"
@property
def providers(self) -> tuple[str, ...]:
return ("codex", "kiro")
def compute_metric(
self,
*,
key_id: str,
all_key_ids: list[str],
keys_by_id: dict[str, Any],
lru_scores: dict[str, Any],
context: dict[str, Any],
mode: str | None,
) -> float:
plan_scores = {
kid: plan_priority_score(extract_plan_type(keys_by_id.get(kid)), "plus_only")
for kid in all_key_ids
}
if len(set(plan_scores.values())) <= 1:
return rank_ascending(key_id, lru_scores, all_key_ids)
return rank_ascending(key_id, plan_scores, all_key_ids)
register_preset_dimension(PlusFirstDimension())

View File

@@ -0,0 +1,46 @@
"""priority_first preset dimension."""
from __future__ import annotations
from typing import Any
from ._helpers import extract_internal_priority, rank_ascending
from .registry import PresetDimensionBase, register_preset_dimension
class PriorityFirstDimension(PresetDimensionBase):
@property
def name(self) -> str:
return "priority_first"
@property
def label(self) -> str:
return "优先级优先"
@property
def description(self) -> str:
return "按账号优先级顺序调度(数字越小越优先)"
@property
def evidence_hint(self) -> str | None:
return "依据 internal_priority支持拖拽/手工编辑)"
def compute_metric(
self,
*,
key_id: str,
all_key_ids: list[str],
keys_by_id: dict[str, Any],
lru_scores: dict[str, Any],
context: dict[str, Any],
mode: str | None,
) -> float:
priority_scores = {
kid: float(extract_internal_priority(keys_by_id.get(kid))) for kid in all_key_ids
}
if len(set(priority_scores.values())) <= 1:
return rank_ascending(key_id, lru_scores, all_key_ids)
return rank_ascending(key_id, priority_scores, all_key_ids)
register_preset_dimension(PriorityFirstDimension())

View File

@@ -0,0 +1,61 @@
"""quota_balanced preset dimension."""
from __future__ import annotations
from typing import Any
from ._helpers import extract_usage_ratio, rank_ascending, safe_float
from .registry import PresetDimensionBase, register_preset_dimension
class QuotaBalancedDimension(PresetDimensionBase):
@property
def name(self) -> str:
return "quota_balanced"
@property
def label(self) -> str:
return "额度平均"
@property
def description(self) -> str:
return "优先选额度消耗最少的账号"
@property
def evidence_hint(self) -> str | None:
return "依据账号配额使用率;无配额时回退到窗口成本使用"
def compute_metric(
self,
*,
key_id: str,
all_key_ids: list[str],
keys_by_id: dict[str, Any],
lru_scores: dict[str, Any],
context: dict[str, Any],
mode: str | None,
) -> float:
usage_scores: dict[str, float] = {}
cost_totals = context.get("cost_totals")
if not isinstance(cost_totals, dict):
cost_totals = {}
cost_limit = safe_float(context.get("cost_limit_per_key_tokens"))
for kid in all_key_ids:
key_obj = keys_by_id.get(kid)
usage_ratio = extract_usage_ratio(key_obj)
if usage_ratio is None:
used = safe_float(cost_totals.get(kid))
if used is not None and used >= 0:
if cost_limit is not None and cost_limit > 0:
usage_ratio = max(0.0, min(used / cost_limit, 1.0))
else:
# 无明确上限时用 log 归一化,确保维度仍有区分能力。
usage_ratio = min(1.0, used / (used + 10000.0))
if usage_ratio is not None:
usage_scores[kid] = usage_ratio
if not usage_scores:
return rank_ascending(key_id, lru_scores, all_key_ids)
return rank_ascending(key_id, usage_scores, all_key_ids)
register_preset_dimension(QuotaBalancedDimension())

View File

@@ -0,0 +1,53 @@
"""recent_refresh preset dimension."""
from __future__ import annotations
from typing import Any
from ._helpers import extract_reset_seconds, rank_ascending
from .registry import PresetDimensionBase, register_preset_dimension
class RecentRefreshDimension(PresetDimensionBase):
@property
def name(self) -> str:
return "recent_refresh"
@property
def label(self) -> str:
return "额度刷新优先"
@property
def description(self) -> str:
return "优先选即将刷新额度的账号"
@property
def providers(self) -> tuple[str, ...]:
return ("codex", "kiro")
@property
def evidence_hint(self) -> str | None:
return "依据账号额度重置倒计时next_reset / reset_seconds"
def compute_metric(
self,
*,
key_id: str,
all_key_ids: list[str],
keys_by_id: dict[str, Any],
lru_scores: dict[str, Any],
context: dict[str, Any],
mode: str | None,
) -> float:
provider_type = context.get("provider_type")
reset_scores: dict[str, float] = {}
for kid in all_key_ids:
reset_seconds = extract_reset_seconds(keys_by_id.get(kid), provider_type=provider_type)
if reset_seconds is not None:
reset_scores[kid] = reset_seconds
if not reset_scores:
return rank_ascending(key_id, lru_scores, all_key_ids)
return rank_ascending(key_id, reset_scores, all_key_ids)
register_preset_dimension(RecentRefreshDimension())

View File

@@ -0,0 +1,265 @@
"""Preset dimension registry for pool multi-score scheduling."""
from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass
from threading import RLock
from typing import Any
@dataclass(frozen=True, slots=True)
class PresetDimensionMeta:
"""Serializable metadata for one preset dimension."""
name: str
label: str
description: str
providers: tuple[str, ...]
modes: tuple[str, ...] | None
default_mode: str | None
mutex_group: str | None
evidence_hint: str | None
class PresetDimensionBase(ABC):
"""Base class of one pool scheduling preset dimension."""
@property
@abstractmethod
def name(self) -> str:
"""Stable preset key, e.g. ``free_team_first``."""
@property
@abstractmethod
def label(self) -> str:
"""User-facing label."""
@property
@abstractmethod
def description(self) -> str:
"""User-facing description."""
@property
def providers(self) -> tuple[str, ...]:
"""Supported provider types.
Empty tuple means the dimension is universal and applies to all providers.
"""
return ()
@property
def modes(self) -> tuple[str, ...] | None:
"""Optional sub-modes for this dimension."""
return None
@property
def default_mode(self) -> str | None:
"""Default mode when mode is omitted."""
return None
@property
def mutex_group(self) -> str | None:
"""Optional mutual-exclusion group key.
Presets in the same group are expected to be mutually exclusive in UI.
"""
return None
@property
def evidence_hint(self) -> str | None:
"""Human-readable hint about which data this preset uses."""
return None
@property
def hidden(self) -> bool:
"""If True, this dimension is excluded from API metadata listings.
The dimension remains functional for backward compatibility but
will not appear in the scheduling dialog.
"""
return False
@abstractmethod
def compute_metric(
self,
*,
key_id: str,
all_key_ids: list[str],
keys_by_id: dict[str, Any],
lru_scores: dict[str, Any],
context: dict[str, Any],
mode: str | None,
) -> float:
"""Compute normalized metric in [0, 1], lower is better."""
def is_applicable(self, provider_type: str) -> bool:
"""Return whether this dimension applies to the given provider type."""
if not self.providers:
return True
normalized = _normalize_name(provider_type)
return normalized in self.providers
def _normalize_name(value: Any) -> str:
if not isinstance(value, str):
return ""
return value.strip().lower()
def _normalize_names(values: tuple[str, ...] | list[str]) -> tuple[str, ...]:
normalized = [_normalize_name(item) for item in values]
return tuple(item for item in normalized if item)
_registry_lock = RLock()
_registry: dict[str, PresetDimensionBase] = {}
def register_preset_dimension(dim: PresetDimensionBase) -> None:
"""Register or replace one preset dimension by name."""
name = _normalize_name(dim.name)
if not name:
raise ValueError("preset dimension name must be a non-empty string")
providers = _normalize_names(dim.providers)
modes = _normalize_names(dim.modes or ())
default_mode = _normalize_name(dim.default_mode)
if modes and default_mode and default_mode not in modes:
raise ValueError(f"default_mode must be one of modes for preset '{name}'")
class _NormalizedDimension(PresetDimensionBase):
# Lightweight wrapper to keep normalized metadata while preserving compute logic.
def __init__(self, wrapped: PresetDimensionBase) -> None:
self._wrapped = wrapped
@property
def name(self) -> str:
return name
@property
def label(self) -> str:
return self._wrapped.label
@property
def description(self) -> str:
return self._wrapped.description
@property
def providers(self) -> tuple[str, ...]:
return providers
@property
def modes(self) -> tuple[str, ...] | None:
return modes or None
@property
def default_mode(self) -> str | None:
if not modes:
return None
if default_mode:
return default_mode
return modes[0]
@property
def mutex_group(self) -> str | None:
raw = _normalize_name(self._wrapped.mutex_group)
return raw or None
@property
def evidence_hint(self) -> str | None:
raw = str(self._wrapped.evidence_hint or "").strip()
return raw or None
@property
def hidden(self) -> bool:
return self._wrapped.hidden
def compute_metric(
self,
*,
key_id: str,
all_key_ids: list[str],
keys_by_id: dict[str, Any],
lru_scores: dict[str, Any],
context: dict[str, Any],
mode: str | None,
) -> float:
return self._wrapped.compute_metric(
key_id=key_id,
all_key_ids=all_key_ids,
keys_by_id=keys_by_id,
lru_scores=lru_scores,
context=context,
mode=mode,
)
normalized = _NormalizedDimension(dim)
with _registry_lock:
_registry[name] = normalized
def get_preset_dimension(name: str) -> PresetDimensionBase | None:
"""Get one registered preset dimension by name."""
key = _normalize_name(name)
if not key:
return None
with _registry_lock:
return _registry.get(key)
def get_all_preset_dimensions() -> list[PresetDimensionBase]:
"""Get all registered preset dimensions in registration order."""
with _registry_lock:
return list(_registry.values())
def get_preset_names() -> set[str]:
"""Get all registered preset names."""
with _registry_lock:
return set(_registry.keys())
def get_preset_dimension_metas() -> list[PresetDimensionMeta]:
"""Get serializable metadata for all preset dimensions."""
metas: list[PresetDimensionMeta] = []
for dim in get_all_preset_dimensions():
if dim.hidden:
continue
metas.append(
PresetDimensionMeta(
name=dim.name,
label=dim.label,
description=dim.description,
providers=dim.providers,
modes=dim.modes,
default_mode=dim.default_mode,
mutex_group=dim.mutex_group,
evidence_hint=dim.evidence_hint,
)
)
return metas
__all__ = [
"PresetDimensionBase",
"PresetDimensionMeta",
"get_all_preset_dimensions",
"get_preset_dimension",
"get_preset_dimension_metas",
"get_preset_names",
"register_preset_dimension",
]

View File

@@ -0,0 +1,51 @@
"""single_account preset dimension."""
from __future__ import annotations
from typing import Any
from ._helpers import extract_internal_priority, rank_ascending, rank_descending
from .registry import PresetDimensionBase, register_preset_dimension
class SingleAccountDimension(PresetDimensionBase):
@property
def name(self) -> str:
return "single_account"
@property
def label(self) -> str:
return "单号优先"
@property
def description(self) -> str:
return "集中使用同一账号(反向 LRU"
@property
def mutex_group(self) -> str | None:
return "distribution_mode"
@property
def evidence_hint(self) -> str | None:
return "先按账号优先级internal_priority同级再按反向 LRU 集中"
def compute_metric(
self,
*,
key_id: str,
all_key_ids: list[str],
keys_by_id: dict[str, Any],
lru_scores: dict[str, Any],
context: dict[str, Any],
mode: str | None,
) -> float:
priority_scores = {
kid: float(extract_internal_priority(keys_by_id.get(kid))) for kid in all_key_ids
}
priority_rank = rank_ascending(key_id, priority_scores, all_key_ids)
lru_concentrate_rank = rank_descending(key_id, lru_scores, all_key_ids)
# 强化“单号优先”的可控性:优先级优先,反向 LRU 作为次级聚合。
return max(0.0, min(priority_rank * 0.75 + lru_concentrate_rank * 0.25, 1.0))
register_preset_dimension(SingleAccountDimension())

View File

@@ -0,0 +1,51 @@
"""team_first preset dimension."""
from __future__ import annotations
from typing import Any
from ._helpers import extract_plan_type, plan_priority_score, rank_ascending
from .registry import PresetDimensionBase, register_preset_dimension
class TeamFirstDimension(PresetDimensionBase):
@property
def name(self) -> str:
return "team_first"
@property
def label(self) -> str:
return "Team 优先"
@property
def description(self) -> str:
return "优先消耗 Team 账号(依赖 plan_type"
@property
def evidence_hint(self) -> str | None:
return "依据 plan_typeTeam 账号优先调度)"
@property
def providers(self) -> tuple[str, ...]:
return ("codex", "kiro")
def compute_metric(
self,
*,
key_id: str,
all_key_ids: list[str],
keys_by_id: dict[str, Any],
lru_scores: dict[str, Any],
context: dict[str, Any],
mode: str | None,
) -> float:
plan_scores = {
kid: plan_priority_score(extract_plan_type(keys_by_id.get(kid)), "team_only")
for kid in all_key_ids
}
if len(set(plan_scores.values())) <= 1:
return rank_ascending(key_id, lru_scores, all_key_ids)
return rank_ascending(key_id, plan_scores, all_key_ids)
register_preset_dimension(TeamFirstDimension())

View File

@@ -0,0 +1,104 @@
"""In-process pool health score cache.
This cache avoids recomputing per-key health aggregation on every request.
It does not replace persistent health storage; source data still comes from
``ProviderAPIKey.health_by_format`` carried on key objects.
"""
from __future__ import annotations
import threading
import time
from typing import Any
_TTL_SECONDS = 30.0
_MAX_PROVIDERS = 500 # 最大缓存 provider 数量,防止无界增长
_LOCK = threading.Lock()
_CACHE: dict[str, tuple[float, dict[str, float]]] = {}
def aggregate_health_score(health_by_format: Any) -> float:
"""Aggregate health score from ``health_by_format`` (lower-bound strategy)."""
if not isinstance(health_by_format, dict) or not health_by_format:
return 1.0
scores: list[float] = []
for item in health_by_format.values():
if not isinstance(item, dict):
continue
try:
score = float(item.get("health_score") or 1.0)
except (TypeError, ValueError):
score = 1.0
scores.append(max(0.0, min(score, 1.0)))
if not scores:
return 1.0
return min(scores)
def get_health_scores(provider_id: str, keys: list[Any]) -> dict[str, float]:
"""Return key health scores with per-provider TTL cache.
Uses incremental merge: if the cache is still valid but missing some keys,
only the missing keys are computed and merged into the existing cache entry.
"""
now = time.monotonic()
keys_by_id: dict[str, Any] = {}
for k in keys:
kid = str(getattr(k, "id", "") or "")
if kid:
keys_by_id[kid] = k
if not keys_by_id:
return {}
with _LOCK:
cached = _CACHE.get(provider_id)
if cached is not None:
expires_at, payload = cached
if now < expires_at:
missing_ids = [kid for kid in keys_by_id if kid not in payload]
if not missing_ids:
return {kid: payload[kid] for kid in keys_by_id}
# Compute only for missing keys, merge into existing cache
for kid in missing_ids:
payload[kid] = aggregate_health_score(
getattr(keys_by_id[kid], "health_by_format", None)
)
# Trim stale key entries that are no longer in the current key set
stale_ids = [sid for sid in payload if sid not in keys_by_id]
for sid in stale_ids:
del payload[sid]
return {kid: payload[kid] for kid in keys_by_id}
fresh: dict[str, float] = {}
for kid, key in keys_by_id.items():
fresh[kid] = aggregate_health_score(getattr(key, "health_by_format", None))
with _LOCK:
_CACHE[provider_id] = (now + _TTL_SECONDS, fresh)
# 超出上限时清理过期条目,仍超限则淘汰最旧条目
if len(_CACHE) > _MAX_PROVIDERS:
expired = [k for k, (exp, _) in _CACHE.items() if now >= exp]
for k in expired:
del _CACHE[k]
if len(_CACHE) > _MAX_PROVIDERS:
oldest_key = min(_CACHE, key=lambda k: _CACHE[k][0])
del _CACHE[oldest_key]
return dict(fresh)
def invalidate_provider_health_scores(provider_id: str) -> None:
"""Invalidate health-score cache for one provider."""
with _LOCK:
_CACHE.pop(provider_id, None)
def _clear_cache_for_tests() -> None:
with _LOCK:
_CACHE.clear()
__all__ = [
"aggregate_health_score",
"get_health_scores",
"invalidate_provider_health_scores",
]

View File

@@ -0,0 +1,359 @@
"""Account Pool health policy: error code classification and key state management.
Maps upstream HTTP status codes to pool-level actions:
| Code | Action |
|--------------|------------------------------------------------------------|
| 401 | Invalidate OAuth token cache; permanent (deactivated) 1h, else no cooldown |
| 402 | Long cooldown (payment issue) |
| 403 | Graded cooldown: severe (suspended/banned) 1h, else 300s+ |
| 400 | Check body for "organization has been disabled" -> cooldown |
| 429 | Cooldown (retry-after or rate_limit_cooldown_seconds) |
| 529 | Cooldown (overload_cooldown_seconds) |
| * | Check unschedulable_rules keyword matching |
| 408/5xx/etc | Transient cooldown (overload_cooldown_seconds) |
"""
from __future__ import annotations
import json
from typing import TYPE_CHECKING
from src.core.logger import logger
from src.services.provider.pool import redis_ops
if TYPE_CHECKING:
from src.services.provider.pool.config import PoolConfig
# Patterns in 400 error body that indicate account-level issues.
_ACCOUNT_DISABLE_PATTERNS = (
"organization has been disabled",
"organization_disabled",
"account has been disabled",
"account_disabled",
"account has been deactivated",
"account_deactivated",
"account deactivated",
)
# 需要更长冷却的账号异常语义403 body 关键字)。
_FORBIDDEN_ACCOUNT_PATTERNS = (
"account suspended",
"account banned",
"account deactivated",
"subscription inactive",
"suspended",
"banned",
"deactivated",
)
_TRANSIENT_STATUS_COOLDOWN_REASON: dict[int, str] = {
408: "request_timeout_408",
409: "conflict_409",
423: "locked_423",
425: "too_early_425",
500: "server_error_500",
502: "bad_gateway_502",
503: "service_unavailable_503",
504: "gateway_timeout_504",
}
def _parse_retry_after(headers: dict[str, str] | None) -> int | None:
"""Extract retry-after seconds from response headers."""
if not headers:
return None
raw = headers.get("retry-after") or headers.get("Retry-After")
if not raw:
return None
try:
val = int(raw)
return max(1, min(val, 3600))
except (ValueError, TypeError):
return None
def _extract_error_message(error_body: str | None) -> str:
"""Best-effort extraction of error message from JSON body."""
if not error_body:
return ""
try:
data = json.loads(error_body)
if isinstance(data, dict):
error_obj = data.get("error")
if isinstance(error_obj, dict):
return str(error_obj.get("message", ""))
if isinstance(error_obj, str):
return error_obj
return str(data.get("message", ""))
except (json.JSONDecodeError, TypeError):
pass
return error_body[:500]
def _resolve_transient_cooldown_ttl(
*,
status_code: int,
retry_after_seconds: int | None,
config: PoolConfig,
) -> int:
"""Resolve cooldown ttl for transient upstream status codes."""
if status_code in (429, 503):
if retry_after_seconds is not None:
return retry_after_seconds
if status_code == 429:
return config.rate_limit_cooldown_seconds
# 408/409/423/425/5xx: 统一走短时过载冷却,避免雪崩重试。
return config.overload_cooldown_seconds
def _parse_google_quota_cooldown(error_body: str | None) -> int | None:
"""Parse Google-specific quota cooldown from error body.
Safe to call for any provider: returns None unless the error body
contains Google-specific fields (quotaResetTimeStamp / quotaResetDelay /
"reset after" message pattern).
"""
if not error_body:
return None
try:
from src.services.provider.adapters.gemini_cli.quota import extract_quota_cooldown_seconds
return extract_quota_cooldown_seconds(error_body)
except Exception:
return None
async def apply_health_policy(
*,
provider_id: str,
key_id: str,
status_code: int,
error_body: str | None,
response_headers: dict[str, str] | None,
config: PoolConfig,
) -> None:
"""Apply health policy for an upstream error.
This is fire-and-forget; exceptions are caught and logged.
"""
if not config.health_policy_enabled:
return
try:
await _apply(
provider_id=provider_id,
key_id=key_id,
status_code=status_code,
error_body=error_body,
response_headers=response_headers,
config=config,
)
except Exception as exc:
logger.warning(
"Pool health policy failed for key {}: {}",
key_id[:8],
str(exc),
)
async def _apply(
*,
provider_id: str,
key_id: str,
status_code: int,
error_body: str | None,
response_headers: dict[str, str] | None,
config: PoolConfig,
) -> None:
error_msg = _extract_error_message(error_body)
# --- 401 Unauthorized ---------------------------------------------------
if status_code == 401:
await redis_ops.invalidate_oauth_token_cache(key_id)
# Check if the 401 body indicates a permanent account-level deactivation
# (e.g. OpenAI "account_deactivated"). These deserve a long cooldown.
error_lower = error_msg.lower()
is_permanent = any(p in error_lower for p in _ACCOUNT_DISABLE_PATTERNS)
if is_permanent:
await redis_ops.set_cooldown(provider_id, key_id, "account_deactivated_401", ttl=3600)
logger.warning(
"Pool[{}]: key {} got 401 with account deactivation, cooldown 1h",
provider_id[:8],
key_id[:8],
)
else:
# Transient auth failure (e.g. expired token) — token cache already
# invalidated above; the next request will trigger a token refresh.
# No cooldown needed: the key should be retried immediately after refresh.
logger.info(
"Pool[{}]: key {} got 401, token cache invalidated (no cooldown)",
provider_id[:8],
key_id[:8],
)
return
# --- 402 Payment Required ------------------------------------------------
if status_code == 402:
await redis_ops.set_cooldown(provider_id, key_id, "payment_required_402", ttl=3600)
logger.warning(
"Pool[{}]: key {} got 402 (payment required), cooldown 1h",
provider_id[:8],
key_id[:8],
)
return
# --- 403 Forbidden -------------------------------------------------------
if status_code == 403:
error_lower = error_msg.lower()
severe = any(pattern in error_lower for pattern in _FORBIDDEN_ACCOUNT_PATTERNS)
ttl = 3600 if severe else max(config.rate_limit_cooldown_seconds, 300)
await redis_ops.set_cooldown(provider_id, key_id, "forbidden_403", ttl=ttl)
logger.warning(
"Pool[{}]: key {} got 403 (forbidden), cooldown {}s",
provider_id[:8],
key_id[:8],
ttl,
)
return
# --- 400 with account-disable pattern ------------------------------------
if status_code == 400:
error_lower = error_msg.lower()
for pattern in _ACCOUNT_DISABLE_PATTERNS:
if pattern in error_lower:
await redis_ops.set_cooldown(
provider_id, key_id, f"account_disabled_400:{pattern}", ttl=3600
)
logger.warning(
"Pool[{}]: key {} got 400 with '{}', cooldown 1h",
provider_id[:8],
key_id[:8],
pattern,
)
return
# --- 429 Rate Limited ----------------------------------------------------
if status_code == 429:
retry_after = _parse_retry_after(response_headers)
if retry_after is None:
retry_after = _parse_google_quota_cooldown(error_body)
ttl = _resolve_transient_cooldown_ttl(
status_code=status_code,
retry_after_seconds=retry_after,
config=config,
)
await redis_ops.set_cooldown(provider_id, key_id, "rate_limited_429", ttl=ttl)
logger.info(
"Pool[{}]: key {} got 429, cooldown {}s",
provider_id[:8],
key_id[:8],
ttl,
)
return
# --- 529 Overloaded ------------------------------------------------------
if status_code == 529:
ttl = config.overload_cooldown_seconds
await redis_ops.set_cooldown(provider_id, key_id, "overloaded_529", ttl=ttl)
logger.info(
"Pool[{}]: key {} got 529, cooldown {}s",
provider_id[:8],
key_id[:8],
ttl,
)
return
# --- Keyword-based unschedulable rules -----------------------------------
if config.unschedulable_rules and error_msg:
error_lower = error_msg.lower()
for rule in config.unschedulable_rules:
if rule.keyword.lower() in error_lower:
ttl = max(60, rule.duration_minutes * 60)
await redis_ops.set_cooldown(
provider_id,
key_id,
f"rule:{rule.keyword}",
ttl=ttl,
)
logger.info(
"Pool[{}]: key {} matched rule '{}', cooldown {}m",
provider_id[:8],
key_id[:8],
rule.keyword,
rule.duration_minutes,
)
return
# --- Transient status bucket (408/409/423/425/5xx) ----------------------
reason = _TRANSIENT_STATUS_COOLDOWN_REASON.get(status_code)
if reason:
retry_after = _parse_retry_after(response_headers)
ttl = _resolve_transient_cooldown_ttl(
status_code=status_code,
retry_after_seconds=retry_after,
config=config,
)
await redis_ops.set_cooldown(provider_id, key_id, reason, ttl=ttl)
logger.info(
"Pool[{}]: key {} got {}, cooldown {}s ({})",
provider_id[:8],
key_id[:8],
status_code,
ttl,
reason,
)
return
async def apply_stream_timeout_policy(
*,
provider_id: str,
key_id: str,
config: PoolConfig,
) -> None:
"""Record a stream timeout event and apply cooldown if threshold is reached.
Called when an upstream stream response times out (no data within the
configured interval). Increments a per-key counter in Redis and sets
a cooldown if the count reaches the configured threshold.
"""
if not config.health_policy_enabled:
return
try:
count = await redis_ops.incr_stream_timeout_count(
provider_id,
key_id,
config.stream_timeout_window_seconds,
)
if count >= config.stream_timeout_threshold:
ttl = config.stream_timeout_cooldown_seconds
await redis_ops.set_cooldown(
provider_id,
key_id,
f"stream_timeout_x{count}",
ttl=ttl,
)
logger.warning(
"Pool[{}]: key {} stream timeout count {} >= threshold {}, cooldown {}s",
provider_id[:8],
key_id[:8],
count,
config.stream_timeout_threshold,
ttl,
)
else:
logger.info(
"Pool[{}]: key {} stream timeout count {}/{}",
provider_id[:8],
key_id[:8],
count,
config.stream_timeout_threshold,
)
except Exception as exc:
logger.warning(
"Pool stream timeout policy failed for key {}: {}",
key_id[:8],
str(exc),
)

View File

@@ -0,0 +1,91 @@
"""Pool scheduling hooks -- provider-type-specific pool behaviour.
Some provider types need custom logic during pool scheduling (e.g. extracting
a session UUID for sticky binding). This module provides a small Protocol +
registry so the pool layer stays generic while provider-specific behaviour
lives alongside each adapter.
"""
from __future__ import annotations
import threading
from typing import Any, Protocol, runtime_checkable
# ---------------------------------------------------------------------------
# Protocol
# ---------------------------------------------------------------------------
@runtime_checkable
class PoolSchedulingHook(Protocol):
"""Provider-type-specific pool scheduling behaviour.
Each provider type can optionally register a hook to customize:
- Session UUID extraction (for sticky sessions)
- Post-success / post-error callbacks
Optional methods (checked via ``hasattr`` by callers):
- ``on_pool_success``
- ``on_pool_error``
"""
name: str
def extract_session_uuid(self, request_body: dict[str, Any]) -> str | None:
"""Extract a session UUID for sticky binding from the request body."""
...
# -- Optional lifecycle callbacks -----------------------------------------
# These are checked via ``hasattr`` so existing implementations that
# don't define them will continue to work.
def on_pool_success(
self,
*,
key_id: str,
session_uuid: str | None,
context: dict[str, Any],
) -> None:
"""Called after a successful pool request (provider-specific logic)."""
...
def on_pool_error(
self,
*,
key_id: str,
status_code: int,
context: dict[str, Any],
) -> None:
"""Called after a failed pool request (provider-specific logic)."""
...
# ---------------------------------------------------------------------------
# Registry
# ---------------------------------------------------------------------------
_hook_registry: dict[str, PoolSchedulingHook] = {}
_registry_lock = threading.Lock()
def register_pool_hook(provider_type: str, hook: PoolSchedulingHook) -> None:
"""Register a pool scheduling hook for a provider type."""
from src.core.provider_types import normalize_provider_type
pt = normalize_provider_type(provider_type)
with _registry_lock:
_hook_registry[pt] = hook
def get_pool_hook(provider_type: str | None) -> PoolSchedulingHook | None:
"""Return the pool scheduling hook for a provider type, or ``None``."""
if not provider_type:
return None
from src.services.provider.envelope import ensure_providers_bootstrapped
ensure_providers_bootstrapped(provider_types=[provider_type])
from src.core.provider_types import normalize_provider_type
pt = normalize_provider_type(provider_type)
return _hook_registry.get(pt)

View File

@@ -0,0 +1,621 @@
"""Account Pool Manager (provider-agnostic).
Stateless facade that coordinates pool operations for any Provider with
pool configuration enabled. All state lives in Redis via :mod:`redis_ops`.
Usage::
mgr = PoolManager(provider_id, pool_config)
reordered = await mgr.reorder_candidates(session_uuid, candidates)
# ... execute request ...
await mgr.on_request_success(session_uuid=..., key_id=..., tokens_used=...)
"""
from __future__ import annotations
import asyncio
import random
from collections.abc import Callable
from typing import TYPE_CHECKING, Any, TypeVar
from src.core.logger import logger
from src.services.provider.pool import redis_ops
from src.services.provider.pool.account_state import resolve_pool_account_state
from src.services.provider.pool.config import PoolConfig
from src.services.provider.pool.health_cache import get_health_scores
from src.services.provider.pool.trace import PoolCandidateTrace, PoolSchedulingTrace
if TYPE_CHECKING:
from src.models.database import ProviderAPIKey
from src.services.scheduling.schemas import ProviderCandidate
class PoolManager:
"""Coordinate pool-level scheduling for a single Provider."""
__slots__ = ("provider_id", "config", "provider_type")
def __init__(
self,
provider_id: str,
config: PoolConfig,
provider_type: str | None = None,
) -> None:
self.provider_id = provider_id
self.config = config
self.provider_type = str(provider_type or "").strip().lower() or None
# ------------------------------------------------------------------
# Core scheduling: reorder candidate list for pool-aware selection
# ------------------------------------------------------------------
async def reorder_candidates(
self,
session_uuid: str | None,
candidates: list[ProviderCandidate],
) -> list[ProviderCandidate]:
"""Reorder *candidates* according to pool rules.
The returned list keeps the same elements but in a new order:
1. **Sticky session hit** -- if the session is already bound to a key
and that key appears in *candidates* and is not in cooldown, move it
to position 0.
2. **Filter** out keys in account-blocked / cooldown / cost-exhausted state (mark
``is_skipped``).
3. **LRU sort** -- among remaining candidates at the same priority
level, sort by least-recently-used.
4. **Random tiebreak** -- among candidates with identical LRU score.
Also builds a :class:`PoolSchedulingTrace` and attaches per-candidate
trace data via ``_pool_extra_data`` / ``_pool_scheduling_trace``
attributes on candidate objects.
"""
if not candidates:
return candidates
pid = self.provider_id
# Build trace
trace = PoolSchedulingTrace(
provider_id=pid,
total_keys=len(candidates),
session_uuid=session_uuid[:8] if session_uuid else None,
)
# --- Strategy: before_select ----------------------------------
strategies = _get_active_strategies(self.config)
key_ids = [str(c.key.id) for c in candidates]
strategy_context: dict[str, Any] = {"session_uuid": session_uuid}
for strategy in strategies:
if hasattr(strategy, "on_before_select"):
try:
filtered = strategy.on_before_select(
provider_id=pid,
key_ids=key_ids,
config=self.config,
context=strategy_context,
)
if filtered is not None:
key_ids = filtered
except Exception:
logger.opt(exception=True).debug(
"Pool[{}]: strategy before_select failed", pid[:8]
)
# --- 1. Sticky session ----------------------------------------
sticky_key_id: str | None = None
if session_uuid and self.config.sticky_session_ttl_seconds > 0:
sticky_key_id = await redis_ops.get_sticky_binding(
pid, session_uuid, self.config.sticky_session_ttl_seconds
)
provider_type = self.provider_type
if provider_type is None and candidates:
first_provider = getattr(candidates[0], "provider", None)
provider_type = str(getattr(first_provider, "provider_type", "") or "").strip().lower()
if not provider_type:
provider_type = None
# --- 2. Batch fetch pool state (parallel) ---------------------
all_key_ids = [str(c.key.id) for c in candidates]
# Fire independent Redis queries concurrently.
# Only fetch reason (no TTL) on the scheduling hot path -- TTL is only
# used for trace display and costs an extra pipeline command per key.
_cooldown_coro = redis_ops.batch_get_cooldowns(pid, all_key_ids, include_ttl=False)
_cost_coro = (
redis_ops.batch_get_cost_totals(pid, all_key_ids, self.config.cost_window_seconds)
if (
self.config.cost_limit_per_key_tokens is not None
or self.config.scheduling_mode == "multi_score"
)
else None
)
# LRU scores are needed both for plain LRU sorting and for multi_score
# dimensions (e.g. cache_affinity / single_account) that rely on
# lru_scores data.
_need_lru = self.config.lru_enabled or self.config.scheduling_mode == "multi_score"
_lru_coro = redis_ops.get_lru_scores(pid, all_key_ids) if _need_lru else None
_latency_coro = (
redis_ops.batch_get_latency_avgs(pid, all_key_ids, self.config.latency_window_seconds)
if self.config.scheduling_mode == "multi_score"
else None
)
# Gather all non-None coroutines in parallel.
coros: list[Any] = [_cooldown_coro]
_cost_idx = -1
_lru_idx = -1
_latency_idx = -1
if _cost_coro is not None:
_cost_idx = len(coros)
coros.append(_cost_coro)
if _lru_coro is not None:
_lru_idx = len(coros)
coros.append(_lru_coro)
if _latency_coro is not None:
_latency_idx = len(coros)
coros.append(_latency_coro)
gathered = await asyncio.gather(*coros)
cooldowns: dict[str, str | None] = gathered[0]
# Cost check
cost_exhausted: set[str] = set()
cost_soft: set[str] = set()
cost_totals: dict[str, int] = {}
if _cost_idx >= 0:
cost_totals = gathered[_cost_idx]
limit = self.config.cost_limit_per_key_tokens
if limit is not None:
for kid, total in cost_totals.items():
if total >= limit:
cost_exhausted.add(kid)
elif total >= limit * self.config.cost_soft_threshold_percent / 100:
cost_soft.add(kid)
# LRU scores
lru_scores: dict[str, float] = {}
if _lru_idx >= 0:
lru_scores = gathered[_lru_idx]
# Latency averages
latency_avgs: dict[str, float] = {}
if _latency_idx >= 0:
latency_avgs = gathered[_latency_idx]
# Health scores (TTL cached, no Redis round-trip) -- only needed for multi_score
health_scores: dict[str, float] = {}
if self.config.scheduling_mode == "multi_score":
health_scores = get_health_scores(pid, [c.key for c in candidates])
strategy_context.update(
{
"provider_type": provider_type,
"all_key_ids": all_key_ids,
"lru_scores": lru_scores,
"cost_totals": cost_totals,
"cost_limit_per_key_tokens": self.config.cost_limit_per_key_tokens,
"latency_avgs": latency_avgs,
"health_scores": health_scores,
"keys_by_id": {str(c.key.id): c.key for c in candidates},
}
)
# --- Strategy: compute_score ----------------------------------
custom_scores: dict[str, float] = {}
for strategy in strategies:
if hasattr(strategy, "compute_score"):
for kid in all_key_ids:
try:
custom = strategy.compute_score(
key_id=kid,
config=self.config,
context=strategy_context,
)
if custom is not None:
custom_scores[kid] = float(custom)
lru_scores[kid] = float(custom)
except Exception:
pass
# --- 3. Classify candidates -----------------------------------
# Use precomputed account states from CandidateBuilder when available
# (upstream_metadata is deferred on pool keys to save memory).
# Fall back to on-the-fly resolution for non-pool candidates.
account_states: dict[str, Any] = {}
for c in candidates:
kid = str(c.key.id)
if kid not in account_states:
precomputed = getattr(c.key, "_pool_account_state", None)
if precomputed is not None:
account_states[kid] = precomputed
else:
account_states[kid] = resolve_pool_account_state(
provider_type=provider_type,
upstream_metadata=getattr(c.key, "upstream_metadata", None),
oauth_invalid_reason=getattr(c.key, "oauth_invalid_reason", None),
)
sticky_candidate: ProviderCandidate | None = None
available: list[ProviderCandidate] = []
skipped: list[ProviderCandidate] = []
for c in candidates:
kid = str(c.key.id)
ct = PoolCandidateTrace(key_id=kid)
ct.scoring_mode = self.config.scheduling_mode
ct.latency_avg_ms = float(latency_avgs.get(kid, 0.0) or 0.0)
ct.health_score = float(health_scores.get(kid, 1.0) or 1.0)
if kid in custom_scores:
ct.composite_score = float(custom_scores[kid])
# Already skipped upstream?
if c.is_skipped:
skipped.append(c)
ct.skipped = True
ct.skip_type = "upstream"
trace.candidate_traces[kid] = ct
continue
# Account blocked?
account_state = account_states[kid]
if account_state.blocked:
c.is_skipped = True
skip_reason = account_state.reason or account_state.label or "account blocked"
c.skip_reason = f"pool account blocked: {skip_reason}"
skipped.append(c)
ct.skipped = True
ct.skip_type = "account_blocked"
ct.account_block_code = account_state.code
ct.account_block_label = account_state.label
ct.account_block_reason = account_state.reason
_attach_pool_extra(c, ct)
trace.candidate_traces[kid] = ct
continue
cd_reason = cooldowns.get(kid)
if cd_reason is not None:
c.is_skipped = True
c.skip_reason = f"pool cooldown: {cd_reason}"
skipped.append(c)
ct.skipped = True
ct.skip_type = "cooldown"
ct.cooldown_reason = cd_reason
ct.cooldown_ttl = None # TTL skipped on hot path for perf
_attach_pool_extra(c, ct)
trace.candidate_traces[kid] = ct
continue
# Cost exhausted?
if kid in cost_exhausted:
c.is_skipped = True
c.skip_reason = "pool cost limit reached"
skipped.append(c)
ct.skipped = True
ct.skip_type = "cost_exhausted"
ct.cost_window_usage = cost_totals.get(kid, 0)
ct.cost_limit = self.config.cost_limit_per_key_tokens
_attach_pool_extra(c, ct)
trace.candidate_traces[kid] = ct
continue
# Sticky hit?
if sticky_key_id and kid == sticky_key_id:
sticky_candidate = c
ct.reason = "sticky"
ct.sticky_hit = True
trace.sticky_session_used = True
else:
available.append(c)
if kid in custom_scores and self.config.scheduling_mode == "multi_score":
ct.reason = "multi_score"
else:
ct.reason = "lru" if lru_scores.get(kid, 0) > 0 else "random"
ct.lru_score = lru_scores.get(kid, 0.0)
ct.cost_window_usage = cost_totals.get(kid, 0)
ct.cost_limit = self.config.cost_limit_per_key_tokens
if kid in cost_soft:
ct.cost_soft_threshold = True
_attach_pool_extra(c, ct)
trace.candidate_traces[kid] = ct
# --- 4. Sort available by LRU ---------------------------------
if lru_scores and available:
available.sort(key=lambda c: lru_scores.get(str(c.key.id), 0.0))
# Random tiebreak among candidates with the same LRU score
if len(available) > 1 and lru_scores:
_shuffle_same_score_groups(available, lru_scores)
# --- 5. Assemble final order ----------------------------------
result: list[ProviderCandidate] = []
if sticky_candidate is not None:
result.append(sticky_candidate)
result.extend(available)
result.extend(skipped)
if sticky_candidate:
logger.debug(
"Pool[{}]: sticky hit key={}",
pid[:8],
sticky_key_id and sticky_key_id[:8],
)
# --- Strategy: after_select -----------------------------------
if result:
first_kid = str(result[0].key.id)
first_trace = trace.candidate_traces.get(first_kid)
for strategy in strategies:
if hasattr(strategy, "on_after_select") and first_trace:
try:
strategy.on_after_select(
provider_id=pid,
selected_key_id=first_kid,
trace=first_trace,
config=self.config,
context=strategy_context,
)
except Exception:
pass
# Attach the full trace to the first candidate for downstream use.
if result:
setattr(result[0], "_pool_scheduling_trace", trace)
return result
async def select_pool_keys(
self,
session_uuid: str | None,
keys: list[ProviderAPIKey],
*,
availability_checker: (
Callable[[ProviderAPIKey], tuple[bool, str | None, str | None]] | None
) = None,
page_size: int = 50,
) -> tuple[list[ProviderAPIKey], PoolSchedulingTrace]:
"""Select and order pool keys with trace output.
Reuses :meth:`reorder_candidates` logic by adapting keys to lightweight
candidate-like wrappers, then propagates skip/trace metadata back onto
each key object for downstream execution/recording.
When *availability_checker* is provided, post-sort availability checks
are performed lazily: only the top *page_size* non-skipped keys are
checked at a time; if all fail, the next page is checked, and so on.
Keys beyond the last checked page are marked as ``deferred`` (skipped
without checking) to avoid unnecessary CPU work on large pools.
"""
if not keys:
return (
[],
PoolSchedulingTrace(
provider_id=self.provider_id,
total_keys=0,
session_uuid=session_uuid[:8] if session_uuid else None,
),
)
class _KeyCandidate:
__slots__ = (
"key",
"is_skipped",
"skip_reason",
"_pool_extra_data",
"_pool_scheduling_trace",
)
def __init__(self, key: ProviderAPIKey) -> None:
self.key = key
self.is_skipped = False
self.skip_reason: str | None = None
self._pool_extra_data: dict | None = None
self._pool_scheduling_trace: PoolSchedulingTrace | None = None
wrappers = [_KeyCandidate(k) for k in keys]
reordered_wrappers = await self.reorder_candidates(session_uuid, wrappers) # type: ignore[arg-type]
trace: PoolSchedulingTrace | None = None
if reordered_wrappers:
maybe_trace = getattr(reordered_wrappers[0], "_pool_scheduling_trace", None)
if isinstance(maybe_trace, PoolSchedulingTrace):
trace = maybe_trace
if trace is None:
trace = PoolSchedulingTrace(
provider_id=self.provider_id,
total_keys=len(keys),
session_uuid=session_uuid[:8] if session_uuid else None,
)
ordered_keys: list[ProviderAPIKey] = []
for order_idx, wrapped in enumerate(reordered_wrappers):
key = wrapped.key
is_skipped = bool(getattr(wrapped, "is_skipped", False))
skip_reason = str(getattr(wrapped, "skip_reason", "") or "")
setattr(key, "_pool_skipped", is_skipped)
setattr(key, "_pool_skip_reason", skip_reason if skip_reason else None)
setattr(key, "_pool_order_index", order_idx)
pool_extra = getattr(wrapped, "_pool_extra_data", None)
setattr(
key, "_pool_extra_data", dict(pool_extra) if isinstance(pool_extra, dict) else {}
)
ordered_keys.append(key)
# -- 分页可用性检查 --
# 排序后对非 skipped key 分页调用 availability_checker
# 找到 page_size 个可用 key 后停止检查,剩余标记 deferred。
if availability_checker is not None:
available_count = 0
found_enough = False
for key in ordered_keys:
if getattr(key, "_pool_skipped", False):
continue
if found_enough:
setattr(key, "_pool_skipped", True)
setattr(key, "_pool_skip_reason", "deferred")
continue
is_available, skip_reason_check, mapping_model = availability_checker(key)
if not is_available:
setattr(key, "_pool_skipped", True)
setattr(key, "_pool_skip_reason", skip_reason_check)
else:
if mapping_model:
setattr(key, "_pool_mapping_matched_model", mapping_model)
available_count += 1
if available_count >= page_size:
found_enough = True
return ordered_keys, trace
# ------------------------------------------------------------------
# Post-request hooks
# ------------------------------------------------------------------
async def on_request_success(
self,
*,
session_uuid: str | None,
key_id: str,
tokens_used: int = 0,
ttfb_ms: int | None = None,
) -> None:
"""Called after a successful upstream request."""
pid = self.provider_id
# Bind sticky session
if session_uuid and self.config.sticky_session_ttl_seconds > 0:
await redis_ops.set_sticky_binding(
pid, session_uuid, key_id, self.config.sticky_session_ttl_seconds
)
# Touch LRU -- needed for both plain LRU mode and multi_score dimensions
# (e.g. cache_affinity) that rely on LRU timestamps.
if self.config.lru_enabled or self.config.scheduling_mode == "multi_score":
await redis_ops.touch_lru(pid, key_id)
# Record cost
if tokens_used > 0 and self.config.cost_limit_per_key_tokens is not None:
await redis_ops.add_cost_entry(
pid, key_id, tokens_used, self.config.cost_window_seconds
)
# Record latency sample for multi-score scheduling.
if self.config.scheduling_mode == "multi_score" and ttfb_ms is not None and ttfb_ms >= 0:
await redis_ops.record_latency(
pid,
key_id,
ttfb_ms,
self.config.latency_window_seconds,
self.config.latency_sample_limit,
)
async def on_request_error(
self,
*,
key_id: str,
status_code: int,
error_body: str | None = None,
response_headers: dict[str, str] | None = None,
) -> None:
"""Called after an upstream error. Delegates to health policy."""
# Import lazily to avoid circular deps
from src.services.provider.pool.health_policy import apply_health_policy
await apply_health_policy(
provider_id=self.provider_id,
key_id=key_id,
status_code=status_code,
error_body=error_body,
response_headers=response_headers,
config=self.config,
)
# ------------------------------------------------------------------
# Key schedulability check (used by candidate_builder)
# ------------------------------------------------------------------
async def is_key_schedulable(self, key_id: str) -> tuple[bool, str | None]:
"""Check if *key_id* is currently schedulable (not in cooldown, not
cost-exhausted). Returns ``(True, None)`` or ``(False, reason)``.
"""
pid = self.provider_id
# Cooldown check
cd = await redis_ops.get_cooldown(pid, key_id)
if cd is not None:
return False, f"pool cooldown: {cd}"
# Cost check
if self.config.cost_limit_per_key_tokens is not None:
total = await redis_ops.get_cost_window_total(
pid, key_id, self.config.cost_window_seconds
)
if total >= self.config.cost_limit_per_key_tokens:
return False, "pool cost limit reached"
return True, None
# Backward-compatible alias
ClaudeCodePoolManager = PoolManager
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
_T = TypeVar("_T")
def _attach_pool_extra(candidate: Any, ct: PoolCandidateTrace) -> None:
"""Attach pool trace extra_data onto a candidate object."""
existing = getattr(candidate, "_pool_extra_data", None) or {}
existing.update(ct.to_extra_data())
setattr(candidate, "_pool_extra_data", existing)
def _get_active_strategies(config: PoolConfig) -> list[Any]:
"""Get active strategies for the given config (lazy import)."""
if not config.strategies:
return []
try:
# Import triggers built-in strategy registration via module-level side effects.
from src.services.provider.pool import strategies as _builtin_strategies # noqa: F401
from src.services.provider.pool.strategy import get_active_strategies
return get_active_strategies(config.strategies)
except Exception:
return []
def _shuffle_same_score(
items: list[_T],
lru_scores: dict[str, float],
key_fn: Callable[[_T], str],
) -> None:
"""In-place random shuffle within groups that share the same LRU score."""
if len(items) <= 1:
return
i = 0
while i < len(items):
score_i = lru_scores.get(key_fn(items[i]), 0.0)
j = i + 1
while j < len(items) and lru_scores.get(key_fn(items[j]), 0.0) == score_i:
j += 1
if j - i > 1:
group = items[i:j]
random.shuffle(group)
items[i:j] = group
i = j
def _shuffle_same_score_groups(
candidates: list[ProviderCandidate],
lru_scores: dict[str, float],
) -> None:
_shuffle_same_score(candidates, lru_scores, lambda c: str(c.key.id))

View File

@@ -0,0 +1,42 @@
"""OAuth token Redis cache for the Account Pool.
Additions over the base ``auth.py`` refresh flow:
- **Redis token cache**: Avoids repeated DB decryption for hot keys.
Cache key: ``provider_oauth_token_cache:{key_id}``
- **Configurable proactive refresh skew**: Default 180 s (3 min) instead
of the base 120 s, configurable via ``PoolConfig.proactive_refresh_seconds``.
- **401 immediate invalidation**: Clears the Redis cache so the next request
triggers a fresh refresh.
This module does NOT replace ``auth.py``; it adds a caching layer that
``auth.py`` can consult before decrypting from DB.
"""
from __future__ import annotations
from src.core.logger import logger
from src.services.provider.pool import redis_ops
async def get_cached_token(key_id: str) -> str | None:
"""Return cached access token from Redis, or None."""
return await redis_ops.get_cached_oauth_token(key_id)
async def cache_token(key_id: str, token: str, expires_in_seconds: int) -> None:
"""Cache an access token in Redis.
*expires_in_seconds* is the remaining lifetime of the token. We shave
off 60 s so the cache expires slightly before the token itself, giving
the refresh flow time to act.
"""
ttl = max(1, expires_in_seconds - 60)
await redis_ops.cache_oauth_token(key_id, token, ttl)
logger.debug("Pool OAuth: cached token for key {} (TTL={}s)", key_id[:8], ttl)
async def invalidate_token(key_id: str) -> None:
"""Invalidate the cached token (e.g. after a 401)."""
await redis_ops.invalidate_oauth_token_cache(key_id)
logger.debug("Pool OAuth: invalidated token cache for key {}", key_id[:8])

View File

@@ -0,0 +1,694 @@
"""Redis operations for the Account Pool (provider-agnostic).
All pool transient state is stored in Redis. This module centralises key
naming, Lua scripts, and graceful fallbacks so that the rest of the pool
layer is free of Redis specifics.
Key schema
----------
ap:{pid}:sticky:{session_uuid} STRING -> key_id (TTL: config)
ap:{pid}:lru ZSET member=key_id, score=unix_ts
ap:{pid}:cooldown:{key_id} STRING -> reason (TTL: error-specific)
ap:{pid}:cost:{key_id} ZSET member=req_id, score=unix_ts
ap:{pid}:latency:{key_id} ZSET member=req_id:ttfb_ms, score=unix_ts
provider_oauth_token_cache:{key_id} STRING -> access_token (TTL: expires - 60)
"""
from __future__ import annotations
import time
import uuid
from typing import TYPE_CHECKING
from src.clients.redis_client import get_redis_client
from src.core.logger import logger
if TYPE_CHECKING:
import redis.asyncio as aioredis
PREFIX = "ap"
def _sticky_key(provider_id: str, session_uuid: str) -> str:
return f"{PREFIX}:{provider_id}:sticky:{session_uuid}"
def _lru_key(provider_id: str) -> str:
return f"{PREFIX}:{provider_id}:lru"
def _cooldown_key(provider_id: str, key_id: str) -> str:
return f"{PREFIX}:{provider_id}:cooldown:{key_id}"
def _cooldown_index_key(provider_id: str) -> str:
"""SET tracking which keys are in cooldown (for O(1) count queries)."""
return f"{PREFIX}:{provider_id}:cooldown_idx"
def _cost_key(provider_id: str, key_id: str) -> str:
return f"{PREFIX}:{provider_id}:cost:{key_id}"
def _latency_key(provider_id: str, key_id: str) -> str:
return f"{PREFIX}:{provider_id}:latency:{key_id}"
def _oauth_cache_key(key_id: str) -> str:
return f"provider_oauth_token_cache:{key_id}"
# ---------------------------------------------------------------------------
# Lua scripts
# ---------------------------------------------------------------------------
# Sticky-select: GET binding, verify it's not in cooldown, refresh TTL.
# KEYS[1] = sticky key, KEYS[2] = cooldown key prefix (ap:{pid}:cooldown:)
# ARGV[1] = ttl
# Returns: key_id or nil
_STICKY_SELECT_LUA = """
local binding = redis.call("GET", KEYS[1])
if not binding then
return nil
end
-- Check cooldown for the bound key
local cooldown_key = KEYS[2] .. binding
local in_cooldown = redis.call("EXISTS", cooldown_key)
if in_cooldown == 1 then
redis.call("DEL", KEYS[1])
return nil
end
redis.call("EXPIRE", KEYS[1], tonumber(ARGV[1]))
return binding
"""
# Cost window sum (read-only, no cleanup on read path for performance).
# KEYS[1] = cost zset key, ARGV[1] = window_start timestamp
# Returns total token count within the window.
_COST_WINDOW_SUM_LUA = """
local key = KEYS[1]
local window_start = tonumber(ARGV[1])
local members = redis.call("ZRANGEBYSCORE", key, window_start, "+inf")
local total = 0
for _, m in ipairs(members) do
local colon = string.find(m, ":", 1, true)
if colon then
local n = tonumber(string.sub(m, colon + 1))
if n then total = total + n end
end
end
return total
"""
# Latency window average (read-only, no cleanup on read path for performance).
# KEYS[1] = latency zset key, ARGV[1] = window_start timestamp
# Returns nil when there are no samples, or avg(ms) as number.
_LATENCY_WINDOW_AVG_LUA = """
local key = KEYS[1]
local window_start = tonumber(ARGV[1])
local members = redis.call("ZRANGEBYSCORE", key, window_start, "+inf")
local total = 0
local count = 0
for _, m in ipairs(members) do
local colon = string.find(m, ":", 1, true)
if colon then
local n = tonumber(string.sub(m, colon + 1))
if n then
total = total + n
count = count + 1
end
end
end
if count == 0 then
return nil
end
return total / count
"""
async def _get_redis() -> "aioredis.Redis | None":
return await get_redis_client(require_redis=False)
# ---------------------------------------------------------------------------
# Sticky session
# ---------------------------------------------------------------------------
async def get_sticky_binding(provider_id: str, session_uuid: str, ttl: int) -> str | None:
"""Get and refresh sticky session binding. Returns key_id or None."""
redis = await _get_redis()
if redis is None:
return None
try:
result = await redis.eval(
_STICKY_SELECT_LUA,
2,
_sticky_key(provider_id, session_uuid),
f"{PREFIX}:{provider_id}:cooldown:",
str(ttl),
)
if result:
return result.decode() if isinstance(result, bytes) else str(result)
return None
except Exception:
logger.debug("Pool: sticky GET failed for session {}", session_uuid[:8])
return None
async def set_sticky_binding(provider_id: str, session_uuid: str, key_id: str, ttl: int) -> None:
"""Create or update sticky session binding."""
redis = await _get_redis()
if redis is None:
return
try:
await redis.setex(_sticky_key(provider_id, session_uuid), ttl, key_id)
except Exception:
logger.debug("Pool: sticky SET failed for session {}", session_uuid[:8])
async def delete_sticky_binding(provider_id: str, session_uuid: str) -> None:
redis = await _get_redis()
if redis is None:
return
try:
await redis.delete(_sticky_key(provider_id, session_uuid))
except Exception:
pass
# ---------------------------------------------------------------------------
# LRU
# ---------------------------------------------------------------------------
async def get_lru_scores(provider_id: str, key_ids: list[str]) -> dict[str, float]:
"""Batch-fetch LRU timestamps. Missing keys get score 0 (highest priority)."""
redis = await _get_redis()
if redis is None:
return {}
try:
lru_k = _lru_key(provider_id)
scores = await redis.zmscore(lru_k, key_ids)
result: dict[str, float] = {}
for kid, score in zip(key_ids, scores):
result[kid] = float(score) if score is not None else 0.0
return result
except Exception:
logger.debug("Pool: LRU ZMSCORE failed for provider {}", provider_id[:8])
return {}
async def touch_lru(provider_id: str, key_id: str) -> None:
"""Update last-used timestamp."""
redis = await _get_redis()
if redis is None:
return
try:
await redis.zadd(_lru_key(provider_id), {key_id: time.time()})
except Exception:
pass
# ---------------------------------------------------------------------------
# Cooldown
# ---------------------------------------------------------------------------
async def set_cooldown(provider_id: str, key_id: str, reason: str, ttl: int) -> None:
redis = await _get_redis()
if redis is None:
return
try:
pipe = redis.pipeline()
pipe.setex(_cooldown_key(provider_id, key_id), ttl, reason)
# Track in index set for O(1) count queries.
idx_key = _cooldown_index_key(provider_id)
pipe.sadd(idx_key, key_id)
# Keep index alive at least as long as the longest cooldown entry.
# Each set_cooldown call refreshes the TTL so the SET won't expire
# while there are still active cooldowns.
pipe.expire(idx_key, ttl + 60)
await pipe.execute()
logger.info(
"Pool: key {} cooldown set: {} ({}s)",
key_id[:8],
reason,
ttl,
)
except Exception:
logger.debug("Pool: cooldown SET failed for key {}", key_id[:8])
async def get_cooldown(provider_id: str, key_id: str) -> str | None:
redis = await _get_redis()
if redis is None:
return None
try:
val = await redis.get(_cooldown_key(provider_id, key_id))
if val:
return val.decode() if isinstance(val, bytes) else str(val)
return None
except Exception:
return None
async def clear_cooldown(provider_id: str, key_id: str) -> None:
redis = await _get_redis()
if redis is None:
return
try:
pipe = redis.pipeline()
pipe.delete(_cooldown_key(provider_id, key_id))
pipe.srem(_cooldown_index_key(provider_id), key_id)
await pipe.execute()
except Exception:
pass
async def batch_get_cooldowns(
provider_id: str,
key_ids: list[str],
*,
include_ttl: bool = False,
) -> dict[str, str | None] | dict[str, tuple[str | None, int | None]]:
"""Batch check cooldown status for multiple keys.
When *include_ttl* is ``True``, each value is a ``(reason, ttl_seconds)``
tuple instead of a plain reason string. The TTL commands are batched in
the same pipeline so there is no extra round-trip.
"""
redis = await _get_redis()
if redis is None:
if include_ttl:
return {k: (None, None) for k in key_ids}
return {k: None for k in key_ids}
try:
pipe = redis.pipeline()
for kid in key_ids:
ck = _cooldown_key(provider_id, kid)
pipe.get(ck)
if include_ttl:
pipe.ttl(ck)
results = await pipe.execute()
if include_ttl:
out_ttl: dict[str, tuple[str | None, int | None]] = {}
# results interleave GET/TTL: [val0, ttl0, val1, ttl1, ...]
for i, kid in enumerate(key_ids):
val = results[i * 2]
ttl_val = results[i * 2 + 1]
reason: str | None = None
if val:
reason = val.decode() if isinstance(val, bytes) else str(val)
ttl_sec: int | None = None
if isinstance(ttl_val, int) and ttl_val > 0:
ttl_sec = ttl_val
out_ttl[kid] = (reason, ttl_sec)
return out_ttl
out: dict[str, str | None] = {}
for kid, val in zip(key_ids, results):
if val:
out[kid] = val.decode() if isinstance(val, bytes) else str(val)
else:
out[kid] = None
return out
except Exception:
if include_ttl:
return {k: (None, None) for k in key_ids}
return {k: None for k in key_ids}
# ---------------------------------------------------------------------------
# Cost tracking
# ---------------------------------------------------------------------------
async def add_cost_entry(provider_id: str, key_id: str, tokens: int, window_seconds: int) -> None:
"""Record a cost entry (tokens used) with automatic window expiry."""
redis = await _get_redis()
if redis is None:
return
try:
now = time.time()
cost_k = _cost_key(provider_id, key_id)
member = f"{uuid.uuid4().hex}:{tokens}"
window_start = now - max(int(window_seconds), 1)
pipe = redis.pipeline()
pipe.zadd(cost_k, {member: now})
# Prune expired entries on the write path (moved from read Lua script).
pipe.zremrangebyscore(cost_k, "-inf", window_start)
# Set a TTL slightly larger than the window to auto-clean abandoned keys.
pipe.expire(cost_k, window_seconds + 600)
await pipe.execute()
except Exception:
logger.debug("Pool: cost ADD failed for key {}", key_id[:8])
async def get_cost_window_total(provider_id: str, key_id: str, window_seconds: int) -> int:
"""Sum tokens used within the rolling window (single key)."""
redis = await _get_redis()
if redis is None:
return 0
try:
now = time.time()
window_start = now - window_seconds
cost_k = _cost_key(provider_id, key_id)
result = await redis.eval(_COST_WINDOW_SUM_LUA, 1, cost_k, str(window_start))
return int(result) if result else 0
except Exception:
logger.debug("Pool: cost SUM failed for key {}", key_id[:8])
return 0
async def batch_get_cost_totals(
provider_id: str, key_ids: list[str], window_seconds: int
) -> dict[str, int]:
"""Batch-fetch cost totals for multiple keys using pipeline + Lua."""
redis = await _get_redis()
if redis is None:
return {k: 0 for k in key_ids}
try:
now = time.time()
window_start = now - window_seconds
pipe = redis.pipeline()
for kid in key_ids:
cost_k = _cost_key(provider_id, kid)
pipe.eval(_COST_WINDOW_SUM_LUA, 1, cost_k, str(window_start))
results = await pipe.execute()
out: dict[str, int] = {}
for kid, val in zip(key_ids, results):
out[kid] = int(val) if val else 0
return out
except Exception:
logger.debug("Pool: batch cost SUM failed for provider {}", provider_id[:8])
return {k: 0 for k in key_ids}
async def record_latency(
provider_id: str,
key_id: str,
ttfb_ms: int,
window_seconds: int,
sample_limit: int,
) -> None:
"""Record one TTFB sample with rolling-window cleanup."""
redis = await _get_redis()
if redis is None:
return
try:
now = time.time()
latency_k = _latency_key(provider_id, key_id)
sample = max(int(ttfb_ms), 0)
member = f"{uuid.uuid4().hex}:{sample}"
pipe = redis.pipeline()
pipe.zadd(latency_k, {member: now})
window_start = now - max(int(window_seconds), 1)
pipe.zremrangebyscore(latency_k, "-inf", window_start)
capped_limit = max(int(sample_limit), 1)
pipe.zremrangebyrank(latency_k, 0, -(capped_limit + 1))
pipe.expire(latency_k, max(int(window_seconds), 1) + 600)
await pipe.execute()
except Exception:
logger.debug("Pool: latency ADD failed for key {}", key_id[:8])
async def batch_get_latency_avgs(
provider_id: str,
key_ids: list[str],
window_seconds: int,
) -> dict[str, float]:
"""Batch-fetch latency averages (ms) for keys in a rolling window."""
redis = await _get_redis()
if redis is None:
return {}
try:
now = time.time()
window_start = now - max(int(window_seconds), 1)
pipe = redis.pipeline()
for kid in key_ids:
pipe.eval(_LATENCY_WINDOW_AVG_LUA, 1, _latency_key(provider_id, kid), str(window_start))
results = await pipe.execute()
out: dict[str, float] = {}
for kid, val in zip(key_ids, results):
if val is None:
continue
try:
out[kid] = float(val)
except Exception:
continue
return out
except Exception:
logger.debug("Pool: batch latency AVG failed for provider {}", provider_id[:8])
return {}
async def clear_cost(provider_id: str, key_id: str) -> None:
redis = await _get_redis()
if redis is None:
return
try:
await redis.delete(_cost_key(provider_id, key_id))
except Exception:
pass
# ---------------------------------------------------------------------------
# OAuth token cache
# ---------------------------------------------------------------------------
async def cache_oauth_token(key_id: str, token: str, ttl: int) -> None:
redis = await _get_redis()
if redis is None:
return
try:
if ttl > 0:
await redis.setex(_oauth_cache_key(key_id), ttl, token)
except Exception:
pass
async def get_cached_oauth_token(key_id: str) -> str | None:
redis = await _get_redis()
if redis is None:
return None
try:
val = await redis.get(_oauth_cache_key(key_id))
if val:
return val.decode() if isinstance(val, bytes) else str(val)
return None
except Exception:
return None
async def invalidate_oauth_token_cache(key_id: str) -> None:
redis = await _get_redis()
if redis is None:
return
try:
await redis.delete(_oauth_cache_key(key_id))
except Exception:
pass
# ---------------------------------------------------------------------------
# Pool status query (admin)
# ---------------------------------------------------------------------------
async def get_sticky_session_count(provider_id: str) -> int:
"""Approximate count of active sticky sessions (via SCAN, for admin display only)."""
redis = await _get_redis()
if redis is None:
return 0
try:
pattern = f"{PREFIX}:{provider_id}:sticky:*"
count = 0
async for _ in redis.scan_iter(match=pattern, count=100):
count += 1
return count
except Exception:
return 0
async def get_key_sticky_count(provider_id: str, key_id: str) -> int:
"""Count sticky sessions bound to a specific key (admin only).
Uses batched SCAN + pipeline MGET to reduce Redis round-trips.
"""
redis = await _get_redis()
if redis is None:
return 0
try:
pattern = f"{PREFIX}:{provider_id}:sticky:*"
count = 0
batch: list[bytes | str] = []
async for k in redis.scan_iter(match=pattern, count=200):
batch.append(k)
if len(batch) >= 200:
vals = await redis.mget(batch)
for val in vals:
if val:
bound_id = val.decode() if isinstance(val, bytes) else str(val)
if bound_id == key_id:
count += 1
batch.clear()
if batch:
vals = await redis.mget(batch)
for val in vals:
if val:
bound_id = val.decode() if isinstance(val, bytes) else str(val)
if bound_id == key_id:
count += 1
return count
except Exception:
return 0
async def 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()
if redis is None:
return None
try:
ttl = await redis.ttl(_cooldown_key(provider_id, key_id))
return ttl if ttl > 0 else None
except Exception:
return None
async def batch_get_cooldown_ttls(provider_id: str, key_ids: list[str]) -> dict[str, int | None]:
"""Batch-fetch cooldown TTLs for multiple keys using pipeline."""
redis = await _get_redis()
if redis is None:
return {k: None for k in key_ids}
try:
pipe = redis.pipeline()
for kid in key_ids:
pipe.ttl(_cooldown_key(provider_id, kid))
results = await pipe.execute()
out: dict[str, int | None] = {}
for kid, ttl in zip(key_ids, results):
out[kid] = int(ttl) if isinstance(ttl, int) and ttl > 0 else None
return out
except Exception:
return {k: None for k in key_ids}
async def batch_count_provider_cooldowns(provider_ids: list[str]) -> dict[str, int]:
"""Count cooldown entries per provider using the cooldown index set.
Uses ``SCARD`` on the ``ap:{pid}:cooldown_idx`` set for O(1) count
instead of scanning the key-space. The index set is maintained by
:func:`set_cooldown` / :func:`clear_cooldown`.
Note: the index set may contain stale entries (expired cooldowns whose
TTL elapsed before an explicit clear). This over-count is acceptable
for admin display purposes -- precision is not critical here.
"""
if not provider_ids:
return {}
redis = await _get_redis()
if redis is None:
return {pid: 0 for pid in provider_ids}
try:
pipe = redis.pipeline()
for pid in provider_ids:
pipe.scard(_cooldown_index_key(pid))
results = await pipe.execute()
counts: dict[str, int] = {}
for pid, val in zip(provider_ids, results):
counts[pid] = max(int(val or 0), 0)
return counts
except Exception:
return {pid: 0 for pid in provider_ids}
# ---------------------------------------------------------------------------
# Stream timeout counter
# ---------------------------------------------------------------------------
_STREAM_TIMEOUT_KEY_FMT = f"{PREFIX}:{{}}:stream_timeout:{{}}"
def _stream_timeout_key(provider_id: str, key_id: str) -> str:
return _STREAM_TIMEOUT_KEY_FMT.format(provider_id, key_id)
async def incr_stream_timeout_count(
provider_id: str,
key_id: str,
window_seconds: int,
) -> int:
"""Increment stream timeout counter and return count within the window.
Uses a ZSET with timestamps as scores. Old entries beyond the window
are pruned on each call. Returns the count of timeouts in the window.
"""
redis = await _get_redis()
if redis is None:
return 0
try:
now = time.time()
window_start = now - window_seconds
key = _stream_timeout_key(provider_id, key_id)
member = f"{uuid.uuid4().hex}"
pipe = redis.pipeline()
pipe.zremrangebyscore(key, "-inf", window_start)
pipe.zadd(key, {member: now})
pipe.zcard(key)
pipe.expire(key, window_seconds + 60)
results = await pipe.execute()
count = int(results[2]) if results[2] else 0
return count
except Exception:
logger.debug("Pool: stream timeout INCR failed for key {}", key_id[:8])
return 0

View File

@@ -0,0 +1,466 @@
"""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
latency_avg_ms: float | None = None
account_blocked: bool = False
account_block_label: str | None = None
account_block_reason: str | None = None
@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
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 _AccountStateDimension:
code: str = "account_state"
label: str = "账号状态"
source: str = "policy"
weight: int = 10
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
if not snapshot.account_blocked:
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="ok",
score=1.0,
)
blocked_label = snapshot.account_block_label or "账号异常"
if blocked_label == "账号封禁":
blocked_code = "account_banned"
elif blocked_label == "工作区停用":
blocked_code = "workspace_deactivated"
elif blocked_label == "账号停用":
blocked_code = "account_disabled"
elif blocked_label == "需要验证":
blocked_code = "account_verification"
elif blocked_label == "访问受限":
blocked_code = "account_forbidden"
else:
blocked_code = "account_blocked"
return PoolSchedulingDimensionResult(
code=blocked_code,
label=blocked_label,
source=self.source,
weight=self.weight,
status="blocked",
blocking=True,
score=0.0,
detail=snapshot.account_block_reason or blocked_label,
)
@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,
)
@dataclass(frozen=True, slots=True)
class _LatencyDimension:
code: str = "latency"
label: str = "延迟"
source: str = "runtime"
weight: int = 3
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
latency = snapshot.latency_avg_ms
if latency is None:
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="ok",
score=1.0,
detail="-",
)
value = max(float(latency), 0.0)
detail = f"{value:.0f}ms"
if value >= 3000:
return PoolSchedulingDimensionResult(
code="latency_high",
label="延迟偏高",
source=self.source,
weight=self.weight,
status="degraded",
score=0.5,
detail=detail,
)
if value >= 1200:
return PoolSchedulingDimensionResult(
code="latency_slow",
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,
)
_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="可用",
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"]
if blocked:
primary = blocked[0]
return PoolSchedulingSummary(
status="blocked",
reason=primary.code,
label=primary.label,
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,
candidate_eligible=True,
blocked_count=0,
degraded_count=len(degraded),
)
return PoolSchedulingSummary(
status="available",
reason="available",
label="可用",
candidate_eligible=True,
blocked_count=0,
degraded_count=0,
)
def _register_default_dimensions() -> None:
register_pool_scheduling_dimension("account_state", _AccountStateDimension())
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("latency", _LatencyDimension())
register_pool_scheduling_dimension("health", _HealthDimension())
_register_default_dimensions()

View File

@@ -0,0 +1,8 @@
"""Built-in pool strategies."""
# Import side effects: register built-in strategies.
import src.services.provider.pool.dimensions # noqa: F401
from . import multi_score # noqa: F401
__all__ = ["multi_score"]

View File

@@ -0,0 +1,264 @@
"""Multi-dimension pool scoring strategy."""
from __future__ import annotations
from typing import Any
import src.services.provider.pool.dimensions # noqa: F401
from src.services.provider.pool.dimensions import get_preset_dimension, get_preset_names
from src.services.provider.pool.dimensions._helpers import rank_ascending, safe_float
from src.services.provider.pool.strategy import register_pool_strategy
def _normalize_mutex_group(value: Any) -> str | None:
if not isinstance(value, str):
return None
normalized = value.strip().lower()
return normalized or None
def _get_preset_mutex_group(preset_name: str) -> str | None:
# LRU is a built-in preset (not in registry) but shares the distribution mutex group.
if preset_name == "lru":
return "distribution_mode"
dim = get_preset_dimension(preset_name)
if dim is None:
return None
return _normalize_mutex_group(getattr(dim, "mutex_group", None))
def _normalize_presets_from_config(
config: Any,
*,
provider_type: str | None = None,
) -> tuple[tuple[str, str | None], ...]:
"""Extract enabled (preset_name, mode) tuples from config.scheduling_presets.
Supports both new SchedulingPreset objects and legacy string lists.
Excludes ``lru`` from output (LRU is a final tie-breaker only).
For mutex groups, enabled members inherit the group's first appearance index
so the selected member keeps the group's visible priority slot.
"""
raw = getattr(config, "scheduling_presets", ())
if not isinstance(raw, (list, tuple)):
return ()
normalized_provider_type = str(provider_type or "").strip().lower()
allowed = get_preset_names() | {"lru"}
entries: list[tuple[int, str, bool, str | None]] = []
seen: set[str] = set()
for idx, item in enumerate(raw):
preset_name: str | None = None
enabled = True
mode: str | None = None
if hasattr(item, "preset"):
preset_name = str(getattr(item, "preset", "")).strip().lower()
enabled = bool(getattr(item, "enabled", True))
raw_mode = getattr(item, "mode", None)
if isinstance(raw_mode, str):
mode = raw_mode.strip().lower() or None
elif isinstance(item, str):
preset_name = item.strip().lower()
else:
continue
if not preset_name or preset_name not in allowed or preset_name in seen:
continue
seen.add(preset_name)
entries.append((idx, preset_name, enabled, mode))
# Codex 默认启用额度刷新优先维度(除非显式配置了 recent_refresh
if (
normalized_provider_type == "codex"
and entries
and "recent_refresh" not in {name for _idx, name, _enabled, _mode in entries}
and "recent_refresh" in allowed
):
entries.append((len(entries), "recent_refresh", True, None))
if not entries:
return ()
group_anchor_index: dict[str, int] = {}
for idx, preset_name, _enabled, _mode in entries:
mutex_group = _get_preset_mutex_group(preset_name)
if mutex_group and mutex_group not in group_anchor_index:
group_anchor_index[mutex_group] = idx
ordered_enabled: list[tuple[int, int, str, str | None]] = []
group_enabled: dict[str, tuple[int, int, str, str | None]] = {}
for idx, preset_name, enabled, mode in entries:
if not enabled or preset_name == "lru":
continue
mutex_group = _get_preset_mutex_group(preset_name)
if not mutex_group:
ordered_enabled.append((idx, idx, preset_name, mode))
continue
anchor = group_anchor_index.get(mutex_group, idx)
existing = group_enabled.get(mutex_group)
if existing is None or idx < existing[1]:
group_enabled[mutex_group] = (anchor, idx, preset_name, mode)
ordered_enabled.extend(group_enabled.values())
ordered_enabled.sort(key=lambda item: (item[0], item[1]))
return tuple((preset_name, mode) for _anchor, _idx, preset_name, mode in ordered_enabled)
class MultiScoreStrategy:
name = "multi_score"
def compute_score(
self,
*,
key_id: str,
config: Any,
context: dict[str, Any],
) -> float | None:
mode = str(getattr(config, "scheduling_mode", "lru") or "lru").strip().lower()
if mode != "multi_score":
return None
all_key_ids = [str(k) for k in (context.get("all_key_ids") or []) if str(k)]
if not all_key_ids:
return None
lru_scores = context.get("lru_scores", {})
if not isinstance(lru_scores, dict):
lru_scores = {}
latency_avgs = context.get("latency_avgs", {})
if not isinstance(latency_avgs, dict):
latency_avgs = {}
health_scores = context.get("health_scores", {})
if not isinstance(health_scores, dict):
health_scores = {}
cost_totals = context.get("cost_totals", {})
if not isinstance(cost_totals, dict):
cost_totals = {}
keys_by_id = context.get("keys_by_id", {})
if not isinstance(keys_by_id, dict):
keys_by_id = {}
presets = _normalize_presets_from_config(
config,
provider_type=context.get("provider_type"),
)
lru_enabled = bool(getattr(config, "lru_enabled", True))
if presets:
return self._compute_preset_score(
key_id=key_id,
all_key_ids=all_key_ids,
presets=presets,
lru_enabled=lru_enabled,
lru_scores=lru_scores,
keys_by_id=keys_by_id,
context=context,
)
weights = getattr(config, "scoring_weights", None)
w_lru = safe_float(getattr(weights, "lru", 0.3)) or 0.0
w_latency = safe_float(getattr(weights, "latency", 0.25)) or 0.0
w_health = safe_float(getattr(weights, "health", 0.2)) or 0.0
w_cost = safe_float(getattr(weights, "cost_remaining", 0.25)) or 0.0
lru_rank = rank_ascending(key_id, lru_scores, all_key_ids)
latency_rank = rank_ascending(key_id, latency_avgs, all_key_ids)
health_raw = safe_float(health_scores.get(key_id))
if health_raw is None:
health_raw = 1.0
health_norm = 1.0 - max(0.0, min(health_raw, 1.0))
cost_limit = getattr(config, "cost_limit_per_key_tokens", None)
used = safe_float(cost_totals.get(key_id)) or 0.0
if cost_limit is None or int(cost_limit) <= 0:
cost_norm = 0.0
else:
cost_norm = max(0.0, min(used / float(cost_limit), 1.0))
return (
w_lru * lru_rank
+ w_latency * latency_rank
+ w_health * health_norm
+ w_cost * cost_norm
)
def _compute_preset_score(
self,
*,
key_id: str,
all_key_ids: list[str],
presets: tuple[tuple[str, str | None], ...],
lru_enabled: bool,
lru_scores: dict[str, Any],
keys_by_id: dict[str, Any],
context: dict[str, Any],
) -> float:
cache_signature = (tuple(all_key_ids), presets, bool(lru_enabled))
cache = context.get("_preset_hard_order_cache")
if (
isinstance(cache, dict)
and cache.get("signature") == cache_signature
and isinstance(cache.get("ranks"), dict)
):
cached_rank = safe_float(cache["ranks"].get(key_id))
if cached_rank is not None:
return max(0.0, min(cached_rank, 1.0))
# Hard-priority semantics:
# 1) Compare by preset[0] metric first;
# 2) only if tied, compare preset[1], preset[2], ...
# 3) if all preset metrics tie and LRU is enabled, use LRU as final tiebreak.
metric_vectors: dict[str, tuple[float, ...]] = {}
for kid in all_key_ids:
vector_parts: list[float] = []
for preset_name, mode in presets:
metric = 0.5
dim = get_preset_dimension(preset_name)
if dim is not None:
metric = dim.compute_metric(
key_id=kid,
all_key_ids=all_key_ids,
keys_by_id=keys_by_id,
lru_scores=lru_scores,
context=context,
mode=mode,
)
metric_value = safe_float(metric)
vector_parts.append(
max(0.0, min(metric_value, 1.0)) if metric_value is not None else 0.5
)
if lru_enabled:
vector_parts.append(rank_ascending(kid, lru_scores, all_key_ids))
metric_vectors[kid] = tuple(vector_parts)
decorated = [
(metric_vectors.get(kid, (0.5,)), idx, kid) for idx, kid in enumerate(all_key_ids)
]
decorated.sort(key=lambda item: (item[0], item[1]))
total = len(decorated)
ranks: dict[str, float] = {}
for rank_idx, (_vec, _idx, kid) in enumerate(decorated):
ranks[kid] = 0.0 if total <= 1 else rank_idx / float(total - 1)
context["_preset_hard_order_cache"] = {
"signature": cache_signature,
"ranks": ranks,
}
rank = safe_float(ranks.get(key_id))
if rank is None:
return 0.5
return max(0.0, min(rank, 1.0))
register_pool_strategy("multi_score", MultiScoreStrategy())
__all__ = [
"MultiScoreStrategy",
]

View File

@@ -0,0 +1,107 @@
"""Pluggable pool scheduling strategies.
Strategies allow customising pool-level candidate selection without
modifying the core :class:`PoolManager`. Each strategy is an object
that implements one or more optional methods defined by the
:class:`PoolSchedulingStrategy` protocol.
Registration uses a thread-safe global registry (same pattern as
:mod:`~src.services.provider.pool.hooks`).
Usage::
from src.services.provider.pool.strategy import register_pool_strategy
class MyStrategy:
name = "usage_weight"
def compute_score(self, *, key_id, config, context):
...
register_pool_strategy("usage_weight", MyStrategy())
"""
from __future__ import annotations
import threading
from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
if TYPE_CHECKING:
from src.services.provider.pool.config import PoolConfig
from src.services.provider.pool.trace import PoolCandidateTrace
@runtime_checkable
class PoolSchedulingStrategy(Protocol):
"""Pluggable pool scheduling strategy.
All methods are optional -- callers check via ``hasattr``.
Strategies are activated per-provider through ``PoolConfig.strategies``.
"""
name: str
def on_before_select(
self,
*,
provider_id: str,
key_ids: list[str],
config: PoolConfig,
context: dict[str, Any],
) -> list[str] | None:
"""Filter / reorder *key_ids* before selection.
Return ``None`` to leave the list unchanged.
"""
...
def on_after_select(
self,
*,
provider_id: str,
selected_key_id: str,
trace: PoolCandidateTrace,
config: PoolConfig,
context: dict[str, Any],
) -> None:
"""Called after a key has been selected (for logging / metrics)."""
...
def compute_score(
self,
*,
key_id: str,
config: PoolConfig,
context: dict[str, Any],
) -> float | None:
"""Return a custom sort score. ``None`` means "do not override"."""
...
# ---------------------------------------------------------------------------
# Registry
# ---------------------------------------------------------------------------
_strategy_registry: dict[str, PoolSchedulingStrategy] = {}
_strategy_lock = threading.Lock()
def register_pool_strategy(name: str, strategy: PoolSchedulingStrategy) -> None:
"""Register a pool scheduling strategy globally."""
with _strategy_lock:
_strategy_registry[name] = strategy
def get_pool_strategy(name: str) -> PoolSchedulingStrategy | None:
"""Return a registered strategy by *name*, or ``None``."""
return _strategy_registry.get(name)
def get_active_strategies(names: tuple[str, ...] | list[str]) -> list[PoolSchedulingStrategy]:
"""Return registered strategies whose names appear in *names*."""
result: list[PoolSchedulingStrategy] = []
for n in names:
s = _strategy_registry.get(n)
if s is not None:
result.append(s)
return result

View File

@@ -0,0 +1,137 @@
"""Pool scheduling trace -- per-candidate decision records.
Collects scheduling decisions made during pool-level candidate selection
without adding any extra Redis round-trips. Trace data is later written
to ``RequestCandidate.extra_data`` and ``Usage.request_metadata``.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any
@dataclass(slots=True)
class PoolCandidateTrace:
"""Single candidate scheduling decision in a pool context."""
key_id: str
reason: str = "" # sticky / lru / random / tiebreak
sticky_hit: bool = False
lru_score: float = 0.0
cost_window_usage: int = 0
cost_limit: int | None = None
cost_soft_threshold: bool = False
skipped: bool = False
skip_type: str | None = None # cooldown / cost_exhausted / account_blocked / upstream
cooldown_reason: str | None = None
cooldown_ttl: int | None = None
account_block_code: str | None = None
account_block_label: str | None = None
account_block_reason: str | None = None
latency_avg_ms: float = 0.0
health_score: float = 1.0
composite_score: float = 0.0
scoring_mode: str = "lru"
def to_extra_data(self) -> dict[str, Any]:
"""Build dict to merge into ``RequestCandidate.extra_data``."""
if self.skipped:
skip_info: dict[str, Any] = {"type": self.skip_type}
if self.cooldown_reason is not None:
skip_info["cooldown_reason"] = self.cooldown_reason
if self.cooldown_ttl is not None:
skip_info["cooldown_ttl"] = self.cooldown_ttl
if self.account_block_code is not None:
skip_info["account_block_code"] = self.account_block_code
if self.account_block_label is not None:
skip_info["account_block_label"] = self.account_block_label
if self.account_block_reason is not None:
skip_info["account_block_reason"] = self.account_block_reason
if self.cost_window_usage:
skip_info["cost_window_usage"] = self.cost_window_usage
if self.scoring_mode:
skip_info["scoring_mode"] = self.scoring_mode
return {"pool_skip": skip_info}
sel: dict[str, Any] = {"reason": self.reason}
if self.sticky_hit:
sel["sticky_hit"] = True
if self.lru_score:
sel["lru_score"] = self.lru_score
if self.cost_window_usage:
sel["cost_window_usage"] = self.cost_window_usage
if self.cost_limit is not None:
sel["cost_limit"] = self.cost_limit
if self.cost_soft_threshold:
sel["cost_soft_threshold"] = True
if self.latency_avg_ms > 0:
sel["latency_avg_ms"] = round(self.latency_avg_ms, 2)
if self.health_score < 1.0:
sel["health_score"] = round(self.health_score, 4)
if self.reason == "multi_score":
sel["composite_score"] = round(self.composite_score, 6)
if self.scoring_mode:
sel["scoring_mode"] = self.scoring_mode
return {"pool_selection": sel}
@dataclass(slots=True)
class PoolSchedulingTrace:
"""Aggregated scheduling trace for one pool-provider dispatch."""
provider_id: str
total_keys: int = 0
sticky_session_used: bool = False
session_uuid: str | None = None
candidate_traces: dict[str, PoolCandidateTrace] = field(default_factory=dict)
def build_summary(
self,
success_key_id: str | None = None,
*,
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
skipped_account_blocked = 0
attempted = 0
for t in self.candidate_traces.values():
if t.skipped:
if t.skip_type == "cooldown":
skipped_cooldown += 1
elif t.skip_type == "cost_exhausted":
skipped_cost += 1
elif t.skip_type == "account_blocked":
skipped_account_blocked += 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:
success_reason = self.candidate_traces[success_key_id].reason
summary: dict[str, Any] = {
"enabled": True,
"total_keys": self.total_keys,
"attempted": attempted,
"skipped_cooldown": skipped_cooldown,
"skipped_cost": skipped_cost,
"skipped_account_blocked": skipped_account_blocked,
"sticky_session": self.sticky_session_used,
}
if success_key_id:
summary["success_key_id"] = success_key_id[:8]
if success_reason:
summary["success_reason"] = success_reason
return summary