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:
fawney19
2026-03-05 09:53:23 +08:00
parent 32ccf61baa
commit 1ac59d4894
24 changed files with 1172 additions and 88 deletions

View File

@@ -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

View File

@@ -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",

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

@@ -21,6 +21,10 @@ class FreeTeamFirstDimension(PresetDimensionBase):
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")
@@ -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)

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,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

@@ -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)

View File

@@ -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)

View File

@@ -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

View File

@@ -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())

View File

@@ -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(
*,

View File

@@ -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},

View File

@@ -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)

View File

@@ -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 响应