mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
feat(pool): 新增调度维度、互斥组机制、健康策略扩展与配额刷新增强
- 新增 priority_first/health_first/latency_first/cost_first 四个调度维度 - 引入 mutex_group 互斥组机制,lru 与 single_account 归入 distribution_mode - 维度 compute_metric 签名扩展 context 参数,支持获取 cost_totals 等上下文 - 各维度增加 evidence_hint 字段描述评分依据 - 健康策略扩展 408/409/423/425/5xx 瞬态状态码冷却,403 按 body 分级冷却 - Codex 配额刷新增强 401/402/403 错误处理,402 生成 fallback 元数据 - 前端号池管理支持账号优先级内联编辑与互斥维度切换 UI - 列表排序改为 internal_priority + created_at,移除 sticky_counts 查询
This commit is contained in:
@@ -5,7 +5,11 @@ Importing this package registers all built-in preset dimensions.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from . import cost_first # noqa: F401
|
||||
from . import free_team_first # noqa: F401
|
||||
from . import health_first # noqa: F401
|
||||
from . import latency_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
|
||||
|
||||
@@ -187,6 +187,37 @@ def extract_usage_ratio(key_obj: Any) -> float | None:
|
||||
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 free_team_first mode."""
|
||||
|
||||
@@ -215,6 +246,8 @@ def plan_priority_score(plan_type: str | None, mode: str | None = None) -> float
|
||||
|
||||
|
||||
__all__ = [
|
||||
"extract_health_score",
|
||||
"extract_internal_priority",
|
||||
"extract_plan_type",
|
||||
"extract_reset_seconds",
|
||||
"extract_usage_ratio",
|
||||
|
||||
62
src/services/provider/pool/dimensions/cost_first.py
Normal file
62
src/services/provider/pool/dimensions/cost_first.py
Normal 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())
|
||||
@@ -21,6 +21,10 @@ class FreeTeamFirstDimension(PresetDimensionBase):
|
||||
def description(self) -> str:
|
||||
return "优先消耗低档账号(依赖 plan_type)"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "依据 plan_type(oauth_plan_type 或 upstream_metadata)"
|
||||
|
||||
@property
|
||||
def providers(self) -> tuple[str, ...]:
|
||||
return ("codex", "kiro")
|
||||
@@ -40,12 +44,15 @@ class FreeTeamFirstDimension(PresetDimensionBase):
|
||||
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)
|
||||
|
||||
|
||||
|
||||
57
src/services/provider/pool/dimensions/health_first.py
Normal file
57
src/services/provider/pool/dimensions/health_first.py
Normal 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())
|
||||
54
src/services/provider/pool/dimensions/latency_first.py
Normal file
54
src/services/provider/pool/dimensions/latency_first.py
Normal 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())
|
||||
46
src/services/provider/pool/dimensions/priority_first.py
Normal file
46
src/services/provider/pool/dimensions/priority_first.py
Normal 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())
|
||||
@@ -4,7 +4,7 @@ from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import extract_usage_ratio, rank_ascending
|
||||
from ._helpers import extract_usage_ratio, rank_ascending, safe_float
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
@@ -21,6 +21,10 @@ class QuotaBalancedDimension(PresetDimensionBase):
|
||||
def description(self) -> str:
|
||||
return "优先选额度消耗最少的账号"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "依据账号配额使用率;无配额时回退到窗口成本使用"
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
@@ -28,13 +32,29 @@ class QuotaBalancedDimension(PresetDimensionBase):
|
||||
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:
|
||||
usage_ratio = extract_usage_ratio(keys_by_id.get(kid))
|
||||
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)
|
||||
|
||||
|
||||
|
||||
@@ -25,6 +25,10 @@ class RecentRefreshDimension(PresetDimensionBase):
|
||||
def providers(self) -> tuple[str, ...]:
|
||||
return ("codex", "kiro")
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "依据账号额度重置倒计时(next_reset / reset_seconds)"
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
@@ -32,6 +36,7 @@ class RecentRefreshDimension(PresetDimensionBase):
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
reset_scores: dict[str, float] = {}
|
||||
@@ -39,6 +44,8 @@ class RecentRefreshDimension(PresetDimensionBase):
|
||||
reset_seconds = extract_reset_seconds(keys_by_id.get(kid))
|
||||
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)
|
||||
|
||||
|
||||
|
||||
@@ -18,6 +18,8 @@ class PresetDimensionMeta:
|
||||
providers: tuple[str, ...]
|
||||
modes: tuple[str, ...] | None
|
||||
default_mode: str | None
|
||||
mutex_group: str | None
|
||||
evidence_hint: str | None
|
||||
|
||||
|
||||
class PresetDimensionBase(ABC):
|
||||
@@ -59,6 +61,21 @@ class PresetDimensionBase(ABC):
|
||||
|
||||
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
|
||||
|
||||
@abstractmethod
|
||||
def compute_metric(
|
||||
self,
|
||||
@@ -67,6 +84,7 @@ class PresetDimensionBase(ABC):
|
||||
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."""
|
||||
@@ -142,6 +160,16 @@ def register_preset_dimension(dim: PresetDimensionBase) -> None:
|
||||
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
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
@@ -149,6 +177,7 @@ def register_preset_dimension(dim: PresetDimensionBase) -> None:
|
||||
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(
|
||||
@@ -156,6 +185,7 @@ def register_preset_dimension(dim: PresetDimensionBase) -> None:
|
||||
all_key_ids=all_key_ids,
|
||||
keys_by_id=keys_by_id,
|
||||
lru_scores=lru_scores,
|
||||
context=context,
|
||||
mode=mode,
|
||||
)
|
||||
|
||||
@@ -201,6 +231,8 @@ def get_preset_dimension_metas() -> list[PresetDimensionMeta]:
|
||||
providers=dim.providers,
|
||||
modes=dim.modes,
|
||||
default_mode=dim.default_mode,
|
||||
mutex_group=dim.mutex_group,
|
||||
evidence_hint=dim.evidence_hint,
|
||||
)
|
||||
)
|
||||
return metas
|
||||
|
||||
@@ -4,7 +4,7 @@ from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import rank_descending
|
||||
from ._helpers import extract_internal_priority, rank_ascending, rank_descending
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
@@ -21,6 +21,14 @@ class SingleAccountDimension(PresetDimensionBase):
|
||||
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,
|
||||
*,
|
||||
@@ -28,9 +36,16 @@ class SingleAccountDimension(PresetDimensionBase):
|
||||
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)
|
||||
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())
|
||||
|
||||
@@ -2,15 +2,16 @@
|
||||
|
||||
Maps upstream HTTP status codes to pool-level actions:
|
||||
|
||||
| Code | Action |
|
||||
|------|--------------------------------------------------------------|
|
||||
| 401 | Invalidate OAuth token cache -> attempt refresh -> disable |
|
||||
| 402 | Auto-disable key (payment issue) |
|
||||
| 403 | Auto-disable key (suspended/banned) |
|
||||
| 400 | Check body for "organization has been disabled" -> disable |
|
||||
| 429 | Set cooldown (retry-after header or config default) |
|
||||
| 529 | Set cooldown (config default) |
|
||||
| * | Check unschedulable_rules keyword matching |
|
||||
| Code | Action |
|
||||
|--------------|------------------------------------------------------------|
|
||||
| 401 | Invalidate OAuth token cache + short 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
|
||||
@@ -32,6 +33,26 @@ _ACCOUNT_DISABLE_PATTERNS = (
|
||||
"account_disabled",
|
||||
)
|
||||
|
||||
# 需要更长冷却的账号异常语义(403 body 关键字)。
|
||||
_FORBIDDEN_ACCOUNT_PATTERNS = (
|
||||
"account suspended",
|
||||
"account banned",
|
||||
"subscription inactive",
|
||||
"suspended",
|
||||
"banned",
|
||||
)
|
||||
|
||||
_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."""
|
||||
@@ -65,6 +86,22 @@ def _extract_error_message(error_body: str | None) -> str:
|
||||
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
|
||||
|
||||
|
||||
async def apply_health_policy(
|
||||
*,
|
||||
provider_id: str,
|
||||
@@ -133,11 +170,15 @@ async def _apply(
|
||||
|
||||
# --- 403 Forbidden -------------------------------------------------------
|
||||
if status_code == 403:
|
||||
await redis_ops.set_cooldown(provider_id, key_id, "forbidden_403", ttl=3600)
|
||||
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/suspended), cooldown 1h",
|
||||
"Pool[{}]: key {} got 403 (forbidden), cooldown {}s",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
ttl,
|
||||
)
|
||||
return
|
||||
|
||||
@@ -160,7 +201,11 @@ async def _apply(
|
||||
# --- 429 Rate Limited ----------------------------------------------------
|
||||
if status_code == 429:
|
||||
retry_after = _parse_retry_after(response_headers)
|
||||
ttl = retry_after or config.rate_limit_cooldown_seconds
|
||||
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",
|
||||
@@ -203,6 +248,26 @@ async def _apply(
|
||||
)
|
||||
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(
|
||||
*,
|
||||
|
||||
@@ -124,7 +124,10 @@ class PoolManager:
|
||||
_cooldown_coro = redis_ops.batch_get_cooldowns(pid, all_key_ids, include_ttl=True)
|
||||
_cost_coro = (
|
||||
redis_ops.batch_get_cost_totals(pid, all_key_ids, self.config.cost_window_seconds)
|
||||
if self.config.cost_limit_per_key_tokens is not None
|
||||
if (
|
||||
self.config.cost_limit_per_key_tokens is not None
|
||||
or self.config.scheduling_mode == "multi_score"
|
||||
)
|
||||
else None
|
||||
)
|
||||
_lru_coro = redis_ops.get_lru_scores(pid, all_key_ids) if self.config.lru_enabled else None
|
||||
@@ -170,12 +173,12 @@ class PoolManager:
|
||||
if _cost_idx >= 0:
|
||||
cost_totals = gathered[_cost_idx]
|
||||
limit = self.config.cost_limit_per_key_tokens
|
||||
assert limit is not None # guarded by _cost_idx >= 0
|
||||
for kid, total in cost_totals.items():
|
||||
if total >= limit:
|
||||
cost_exhausted.add(kid)
|
||||
elif total >= limit * self.config.cost_soft_threshold_percent / 100:
|
||||
cost_soft.add(kid)
|
||||
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] = {}
|
||||
@@ -197,6 +200,7 @@ class PoolManager:
|
||||
"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},
|
||||
@@ -460,7 +464,10 @@ class PoolManager:
|
||||
_cooldown_coro = redis_ops.batch_get_cooldowns(pid, key_ids)
|
||||
_cost_coro = (
|
||||
redis_ops.batch_get_cost_totals(pid, key_ids, self.config.cost_window_seconds)
|
||||
if self.config.cost_limit_per_key_tokens is not None
|
||||
if (
|
||||
self.config.cost_limit_per_key_tokens is not None
|
||||
or self.config.scheduling_mode == "multi_score"
|
||||
)
|
||||
else None
|
||||
)
|
||||
_lru_coro = redis_ops.get_lru_scores(pid, key_ids) if self.config.lru_enabled else None
|
||||
@@ -492,9 +499,10 @@ class PoolManager:
|
||||
cost_totals: dict[str, int] = {}
|
||||
if _cost_idx_sk >= 0:
|
||||
cost_totals = gathered_sk[_cost_idx_sk]
|
||||
for kid, total in cost_totals.items():
|
||||
if total >= self.config.cost_limit_per_key_tokens: # type: ignore[operator]
|
||||
cost_exhausted.add(kid)
|
||||
if self.config.cost_limit_per_key_tokens is not None:
|
||||
for kid, total in cost_totals.items():
|
||||
if total >= self.config.cost_limit_per_key_tokens:
|
||||
cost_exhausted.add(kid)
|
||||
|
||||
lru_scores: dict[str, float] = {}
|
||||
if _lru_idx_sk >= 0:
|
||||
@@ -515,6 +523,7 @@ class PoolManager:
|
||||
"all_key_ids": key_ids,
|
||||
"lru_scores": lru_scores,
|
||||
"cost_totals": cost_totals if _cost_idx_sk >= 0 else {},
|
||||
"cost_limit_per_key_tokens": self.config.cost_limit_per_key_tokens,
|
||||
"latency_avgs": latency_avgs,
|
||||
"health_scores": health_scores,
|
||||
"keys_by_id": {str(k.id): k for k in keys},
|
||||
|
||||
@@ -101,6 +101,7 @@ class MultiScoreStrategy:
|
||||
lru_enabled=lru_enabled,
|
||||
lru_scores=lru_scores,
|
||||
keys_by_id=keys_by_id,
|
||||
context=context,
|
||||
)
|
||||
|
||||
weights = getattr(config, "scoring_weights", None)
|
||||
@@ -140,6 +141,7 @@ class MultiScoreStrategy:
|
||||
lru_enabled: bool,
|
||||
lru_scores: dict[str, Any],
|
||||
keys_by_id: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
) -> float:
|
||||
lru_rank_asc = rank_ascending(key_id, lru_scores, all_key_ids)
|
||||
|
||||
@@ -155,6 +157,7 @@ class MultiScoreStrategy:
|
||||
all_key_ids=all_key_ids,
|
||||
keys_by_id=keys_by_id,
|
||||
lru_scores=lru_scores,
|
||||
context=context,
|
||||
mode=mode,
|
||||
)
|
||||
weight = 1.0 / (1.0 + _POSITIONAL_DECAY * idx)
|
||||
|
||||
@@ -5,6 +5,8 @@ Codex 配额刷新策略。
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
@@ -14,7 +16,10 @@ from src.api.handlers.base.request_builder import get_provider_auth
|
||||
from src.core.crypto import crypto_service
|
||||
from src.models.database import Provider, ProviderAPIKey, ProviderEndpoint
|
||||
from src.services.provider_keys.auth_type import normalize_auth_type
|
||||
from src.services.provider_keys.codex_usage_parser import parse_codex_wham_usage_response
|
||||
from src.services.provider_keys.codex_usage_parser import (
|
||||
parse_codex_usage_headers,
|
||||
parse_codex_wham_usage_response,
|
||||
)
|
||||
|
||||
|
||||
def _normalize_plan_type(value: Any) -> str | None:
|
||||
@@ -24,6 +29,40 @@ def _normalize_plan_type(value: Any) -> str | None:
|
||||
return normalized or None
|
||||
|
||||
|
||||
def _build_quota_exhausted_fallback_metadata(plan_type: str | None) -> dict[str, Any]:
|
||||
"""Build conservative Codex quota metadata when wham/usage returns 402."""
|
||||
normalized_plan = _normalize_plan_type(plan_type)
|
||||
metadata: dict[str, Any] = {"updated_at": int(time.time())}
|
||||
if normalized_plan:
|
||||
metadata["plan_type"] = normalized_plan
|
||||
# primary_* = weekly, secondary_* = 5H (aligned with parser semantics)
|
||||
metadata["primary_used_percent"] = 100.0
|
||||
if normalized_plan != "free":
|
||||
metadata["secondary_used_percent"] = 100.0
|
||||
return metadata
|
||||
|
||||
|
||||
def _extract_error_message_from_response(response: httpx.Response) -> str:
|
||||
"""Best-effort extraction of upstream error message for diagnostics."""
|
||||
try:
|
||||
payload = response.json()
|
||||
if isinstance(payload, dict):
|
||||
err = payload.get("error")
|
||||
if isinstance(err, dict):
|
||||
message = str(err.get("message", "")).strip()
|
||||
if message:
|
||||
return message
|
||||
if isinstance(err, str) and err.strip():
|
||||
return err.strip()
|
||||
message = str(payload.get("message", "")).strip()
|
||||
if message:
|
||||
return message
|
||||
except Exception:
|
||||
pass
|
||||
text = str(getattr(response, "text", "") or "").strip()
|
||||
return text[:300] if text else ""
|
||||
|
||||
|
||||
async def refresh_codex_key_quota(
|
||||
*,
|
||||
db: Session,
|
||||
@@ -96,12 +135,68 @@ async def refresh_codex_key_quota(
|
||||
response = await client.get(codex_wham_usage_url, headers=headers)
|
||||
|
||||
if response.status_code != 200:
|
||||
status_code = int(response.status_code)
|
||||
err_msg = _extract_error_message_from_response(response)
|
||||
|
||||
header_quota = parse_codex_usage_headers(dict(response.headers) if response.headers else {})
|
||||
if isinstance(header_quota, dict) and header_quota:
|
||||
metadata_updates[key.id] = {"codex": header_quota}
|
||||
|
||||
if status_code == 401:
|
||||
state_updates[key.id] = {
|
||||
"is_active": False,
|
||||
"oauth_invalid_at": datetime.now(timezone.utc),
|
||||
"oauth_invalid_reason": "Codex Token 无效或已过期 (401)",
|
||||
}
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "auth_invalid",
|
||||
"message": f"wham/usage API 返回状态码 401{f': {err_msg}' if err_msg else ''}",
|
||||
"status_code": 401,
|
||||
"auto_disabled": True,
|
||||
}
|
||||
|
||||
if status_code == 402:
|
||||
if key.id not in metadata_updates:
|
||||
metadata_updates[key.id] = {
|
||||
"codex": _build_quota_exhausted_fallback_metadata(oauth_plan_type)
|
||||
}
|
||||
state_updates[key.id] = {
|
||||
"oauth_invalid_at": None,
|
||||
"oauth_invalid_reason": None,
|
||||
}
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "quota_exhausted",
|
||||
"message": f"wham/usage API 返回状态码 402{f': {err_msg}' if err_msg else ''}",
|
||||
"status_code": 402,
|
||||
}
|
||||
|
||||
if status_code == 403:
|
||||
state_updates[key.id] = {
|
||||
"is_active": False,
|
||||
"oauth_invalid_at": datetime.now(timezone.utc),
|
||||
"oauth_invalid_reason": "Codex 账户访问受限 (403)",
|
||||
}
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "forbidden",
|
||||
"message": f"wham/usage API 返回状态码 403{f': {err_msg}' if err_msg else ''}",
|
||||
"status_code": 403,
|
||||
"auto_disabled": True,
|
||||
}
|
||||
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": f"wham/usage API 返回状态码 {response.status_code}",
|
||||
"status_code": response.status_code,
|
||||
"message": (
|
||||
f"wham/usage API 返回状态码 {status_code}{f': {err_msg}' if err_msg else ''}"
|
||||
),
|
||||
"status_code": status_code,
|
||||
}
|
||||
|
||||
# 解析 JSON 响应
|
||||
|
||||
Reference in New Issue
Block a user