mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
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:
@@ -5,7 +5,7 @@ from __future__ import annotations
|
|||||||
import math
|
import math
|
||||||
from typing import Any
|
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
|
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
|
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)
|
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] = []
|
candidates: list[float] = []
|
||||||
|
|
||||||
for provider_type in (ProviderType.CODEX, ProviderType.KIRO, ProviderType.ANTIGRAVITY):
|
for provider_type in (ProviderType.CODEX, ProviderType.KIRO, ProviderType.ANTIGRAVITY):
|
||||||
|
|||||||
@@ -39,9 +39,10 @@ class RecentRefreshDimension(PresetDimensionBase):
|
|||||||
context: dict[str, Any],
|
context: dict[str, Any],
|
||||||
mode: str | None,
|
mode: str | None,
|
||||||
) -> float:
|
) -> float:
|
||||||
|
provider_type = context.get("provider_type")
|
||||||
reset_scores: dict[str, float] = {}
|
reset_scores: dict[str, float] = {}
|
||||||
for kid in all_key_ids:
|
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:
|
if reset_seconds is not None:
|
||||||
reset_scores[kid] = reset_seconds
|
reset_scores[kid] = reset_seconds
|
||||||
if not reset_scores:
|
if not reset_scores:
|
||||||
|
|||||||
@@ -193,6 +193,7 @@ class PoolManager:
|
|||||||
|
|
||||||
strategy_context.update(
|
strategy_context.update(
|
||||||
{
|
{
|
||||||
|
"provider_type": provider_type,
|
||||||
"all_key_ids": all_key_ids,
|
"all_key_ids": all_key_ids,
|
||||||
"lru_scores": lru_scores,
|
"lru_scores": lru_scores,
|
||||||
"cost_totals": cost_totals,
|
"cost_totals": cost_totals,
|
||||||
@@ -561,10 +562,18 @@ class PoolManager:
|
|||||||
if self.config.scheduling_mode == "multi_score":
|
if self.config.scheduling_mode == "multi_score":
|
||||||
health_scores = get_health_scores(pid, keys)
|
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 ------------------------------------------
|
# --- Strategy: compute_score ------------------------------------------
|
||||||
strategies = _get_active_strategies(self.config)
|
strategies = _get_active_strategies(self.config)
|
||||||
strategy_context: dict[str, Any] = {
|
strategy_context: dict[str, Any] = {
|
||||||
"session_uuid": session_uuid,
|
"session_uuid": session_uuid,
|
||||||
|
"provider_type": provider_type,
|
||||||
"all_key_ids": key_ids,
|
"all_key_ids": key_ids,
|
||||||
"lru_scores": lru_scores,
|
"lru_scores": lru_scores,
|
||||||
"cost_totals": cost_totals if _cost_idx_sk >= 0 else {},
|
"cost_totals": cost_totals if _cost_idx_sk >= 0 else {},
|
||||||
|
|||||||
@@ -161,6 +161,41 @@ def test_multi_score_preset_recent_refresh_prefers_nearer_reset() -> None:
|
|||||||
assert s2 < s1
|
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:
|
def test_multi_score_preset_single_account_prefers_internal_priority_then_reverse_lru() -> None:
|
||||||
strategy = MultiScoreStrategy()
|
strategy = MultiScoreStrategy()
|
||||||
cfg = PoolConfig(
|
cfg = PoolConfig(
|
||||||
|
|||||||
@@ -83,6 +83,17 @@ def test_parse_codex_usage_headers_paid_windows_mapping() -> None:
|
|||||||
assert "credits_balance" not in parsed
|
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:
|
def test_parse_codex_usage_headers_free_uses_primary_only() -> None:
|
||||||
parsed = parse_codex_usage_headers(
|
parsed = parse_codex_usage_headers(
|
||||||
{
|
{
|
||||||
|
|||||||
@@ -542,6 +542,35 @@ def test_parse_codex_usage_missing_plan_type_infers_paid_windows() -> None:
|
|||||||
assert parsed["secondary_window_minutes"] == 300
|
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:
|
def test_parse_codex_usage_blank_credits_balance_is_ignored() -> None:
|
||||||
parsed = parse_codex_wham_usage_response(
|
parsed = parse_codex_wham_usage_response(
|
||||||
{
|
{
|
||||||
|
|||||||
@@ -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)
|
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:
|
def test_resolve_pool_account_state_keeps_codex_metadata_block() -> None:
|
||||||
state = resolve_pool_account_state(
|
state = resolve_pool_account_state(
|
||||||
provider_type="codex",
|
provider_type="codex",
|
||||||
|
|||||||
Reference in New Issue
Block a user