fix(pool): recent_refresh 按 provider_type 解析 reset_seconds 并锁定 Codex 周重置语义

- 为 `extract_reset_seconds` 增加 `provider_type` 解析链路(显式参数 -> key.provider_type -> key.provider.provider_type -> metadata 推断)。
- Codex 场景固定读取 `primary_reset_seconds`(周限额),避免误用 `secondary_reset_seconds`(5 小时窗口);Kiro/Antigravity 继续走统一 quota reader。
- 在 `RecentRefreshDimension` 和 `PoolManager` 的策略上下文中透传 `provider_type`,并在缺失时从 key provider 回退推断。
- 新增/补强测试覆盖:
  - multi_score `recent_refresh` 在 Codex 下按 weekly reset 排序;
  - quota reader helper 在 Codex provider 下返回 primary reset;
  - Codex 付费计划(plus/enterprise)在 headers 与 WHAM 响应中的窗口映射与主次窗口对齐(10080/300 分钟)。

降低 recent_refresh 评分偏差风险,并为 Codex 付费窗口解析提供回归保护。
This commit is contained in:
AAEE86
2026-03-10 14:55:48 +08:00
parent 2d846b2c58
commit f82964217e
7 changed files with 148 additions and 3 deletions

View File

@@ -5,7 +5,7 @@ from __future__ import annotations
import math
from typing import Any
from src.core.provider_types import ProviderType
from src.core.provider_types import ProviderType, normalize_provider_type
from src.services.provider_keys.quota_reader import get_quota_reader
@@ -108,8 +108,54 @@ def extract_plan_type(key_obj: Any) -> str | None:
return None
def extract_reset_seconds(key_obj: Any) -> float | 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
parsed = safe_float(codex.get("primary_reset_seconds"))
if parsed is None or parsed < 0:
return None
return parsed
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):

View File

@@ -39,9 +39,10 @@ class RecentRefreshDimension(PresetDimensionBase):
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))
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:

View File

@@ -193,6 +193,7 @@ class PoolManager:
strategy_context.update(
{
"provider_type": provider_type,
"all_key_ids": all_key_ids,
"lru_scores": lru_scores,
"cost_totals": cost_totals,
@@ -561,10 +562,18 @@ class PoolManager:
if self.config.scheduling_mode == "multi_score":
health_scores = get_health_scores(pid, keys)
provider_type = self.provider_type
if provider_type is None and keys:
first_provider = getattr(keys[0], "provider", None)
provider_type = str(getattr(first_provider, "provider_type", "") or "").strip().lower()
if not provider_type:
provider_type = None
# --- Strategy: compute_score ------------------------------------------
strategies = _get_active_strategies(self.config)
strategy_context: dict[str, Any] = {
"session_uuid": session_uuid,
"provider_type": provider_type,
"all_key_ids": key_ids,
"lru_scores": lru_scores,
"cost_totals": cost_totals if _cost_idx_sk >= 0 else {},

View File

@@ -161,6 +161,41 @@ def test_multi_score_preset_recent_refresh_prefers_nearer_reset() -> None:
assert s2 < s1
def test_multi_score_preset_recent_refresh_uses_codex_weekly_reset() -> None:
strategy = MultiScoreStrategy()
cfg = PoolConfig(
scheduling_mode="multi_score",
scheduling_presets=(SchedulingPreset(preset="recent_refresh", enabled=True),),
)
ctx = {
"provider_type": "codex",
"all_key_ids": ["k1", "k2"],
"lru_scores": {"k1": 100.0, "k2": 100.0},
"keys_by_id": {
"k1": _key_with_metadata(
{
"codex": {
"primary_reset_seconds": 600,
"secondary_reset_seconds": 120,
}
}
),
"k2": _key_with_metadata(
{
"codex": {
"primary_reset_seconds": 300,
"secondary_reset_seconds": 900,
}
}
),
},
}
s1 = strategy.compute_score(key_id="k1", config=cfg, context=ctx)
s2 = strategy.compute_score(key_id="k2", config=cfg, context=ctx)
assert s1 is not None and s2 is not None
assert s2 < s1
def test_multi_score_preset_single_account_prefers_internal_priority_then_reverse_lru() -> None:
strategy = MultiScoreStrategy()
cfg = PoolConfig(

View File

@@ -83,6 +83,17 @@ def test_parse_codex_usage_headers_paid_windows_mapping() -> None:
assert "credits_balance" not in parsed
@pytest.mark.parametrize("plan_type", ["plus", "enterprise"])
def test_parse_codex_usage_headers_paid_plan_windows_mapping(plan_type: str) -> None:
parsed = parse_codex_usage_headers(_paid_headers(**{"x-codex-plan-type": plan_type}))
assert parsed is not None
assert parsed["plan_type"] == plan_type
assert parsed["primary_used_percent"] == 64.0
assert parsed["secondary_used_percent"] == 3.0
assert parsed["primary_window_minutes"] == 10080
assert parsed["secondary_window_minutes"] == 300
def test_parse_codex_usage_headers_free_uses_primary_only() -> None:
parsed = parse_codex_usage_headers(
{

View File

@@ -542,6 +542,35 @@ def test_parse_codex_usage_missing_plan_type_infers_paid_windows() -> None:
assert parsed["secondary_window_minutes"] == 300
@pytest.mark.parametrize("plan_type", ["plus", "enterprise"])
def test_parse_codex_usage_paid_plan_windows_mapping(plan_type: str) -> None:
parsed = parse_codex_wham_usage_response(
{
"plan_type": plan_type,
"rate_limit": {
"primary_window": {
"used_percent": 25,
"reset_after_seconds": 600,
"reset_at": 1700000000,
"limit_window_seconds": 18000,
},
"secondary_window": {
"used_percent": 80,
"reset_after_seconds": 3600,
"reset_at": 1700003600,
"limit_window_seconds": 604800,
},
},
}
)
assert parsed is not None
assert parsed["plan_type"] == plan_type
assert parsed["primary_used_percent"] == 80.0
assert parsed["secondary_used_percent"] == 25.0
assert parsed["primary_window_minutes"] == 10080
assert parsed["secondary_window_minutes"] == 300
def test_parse_codex_usage_blank_credits_balance_is_ignored() -> None:
parsed = parse_codex_wham_usage_response(
{

View File

@@ -74,6 +74,20 @@ def test_dimension_helpers_delegate_to_unified_reader(monkeypatch: pytest.Monkey
assert extract_usage_ratio(key_obj) == pytest.approx(0.45)
def test_extract_reset_seconds_uses_codex_weekly_reset_for_codex_provider() -> None:
key_obj = SimpleNamespace(
provider_type="codex",
upstream_metadata={
"codex": {
"primary_reset_seconds": 1800.0,
"secondary_reset_seconds": 120.0,
}
},
)
assert extract_reset_seconds(key_obj) == pytest.approx(1800.0)
def test_resolve_pool_account_state_keeps_codex_metadata_block() -> None:
state = resolve_pool_account_state(
provider_type="codex",