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

@@ -118,7 +118,7 @@ async def test_402_sets_long_cooldown(config: PoolConfig) -> None:
@pytest.mark.asyncio
async def test_403_sets_long_cooldown(config: PoolConfig) -> None:
async def test_403_default_sets_medium_cooldown(config: PoolConfig) -> None:
with patch(
"src.services.provider.pool.redis_ops.set_cooldown",
new_callable=AsyncMock,
@@ -132,6 +132,24 @@ async def test_403_sets_long_cooldown(config: PoolConfig) -> None:
config=config,
)
mock_cd.assert_called_once_with(PID, KID, "forbidden_403", ttl=300)
@pytest.mark.asyncio
async def test_403_suspended_body_sets_long_cooldown(config: PoolConfig) -> None:
with patch(
"src.services.provider.pool.redis_ops.set_cooldown",
new_callable=AsyncMock,
) as mock_cd:
await apply_health_policy(
provider_id=PID,
key_id=KID,
status_code=403,
error_body=json.dumps({"error": {"message": "account suspended"}}),
response_headers=None,
config=config,
)
mock_cd.assert_called_once_with(PID, KID, "forbidden_403", ttl=3600)
@@ -228,6 +246,42 @@ async def test_529_uses_overload_cooldown(config: PoolConfig) -> None:
mock_cd.assert_called_once_with(PID, KID, "overloaded_529", ttl=30)
@pytest.mark.asyncio
async def test_503_uses_retry_after_when_present(config: PoolConfig) -> None:
with patch(
"src.services.provider.pool.redis_ops.set_cooldown",
new_callable=AsyncMock,
) as mock_cd:
await apply_health_policy(
provider_id=PID,
key_id=KID,
status_code=503,
error_body=None,
response_headers={"retry-after": "45"},
config=config,
)
mock_cd.assert_called_once_with(PID, KID, "service_unavailable_503", ttl=45)
@pytest.mark.asyncio
async def test_500_uses_overload_cooldown(config: PoolConfig) -> None:
with patch(
"src.services.provider.pool.redis_ops.set_cooldown",
new_callable=AsyncMock,
) as mock_cd:
await apply_health_policy(
provider_id=PID,
key_id=KID,
status_code=500,
error_body=None,
response_headers=None,
config=config,
)
mock_cd.assert_called_once_with(PID, KID, "server_error_500", ttl=30)
# ---------------------------------------------------------------------------
# Unschedulable keyword rules
# ---------------------------------------------------------------------------

View File

@@ -19,8 +19,8 @@ def _context() -> dict:
}
def _key_with_metadata(metadata: dict) -> SimpleNamespace:
return SimpleNamespace(upstream_metadata=metadata)
def _key_with_metadata(metadata: dict, **kwargs: object) -> SimpleNamespace:
return SimpleNamespace(upstream_metadata=metadata, **kwargs)
def test_multi_score_returns_none_when_mode_not_enabled() -> None:
@@ -159,7 +159,7 @@ def test_multi_score_preset_recent_refresh_prefers_nearer_reset() -> None:
assert s2 < s1
def test_multi_score_preset_single_account_prefers_latest_used() -> None:
def test_multi_score_preset_single_account_prefers_internal_priority_then_reverse_lru() -> None:
strategy = MultiScoreStrategy()
cfg = PoolConfig(
scheduling_mode="multi_score",
@@ -168,7 +168,33 @@ def test_multi_score_preset_single_account_prefers_latest_used() -> None:
ctx = {
"all_key_ids": ["k1", "k2", "k3"],
"lru_scores": {"k1": 100.0, "k2": 900.0, "k3": 400.0},
"keys_by_id": {},
"keys_by_id": {
"k1": _key_with_metadata({}, internal_priority=30),
"k2": _key_with_metadata({}, internal_priority=1),
"k3": _key_with_metadata({}, internal_priority=10),
},
}
s1 = strategy.compute_score(key_id="k1", config=cfg, context=ctx)
s2 = strategy.compute_score(key_id="k2", config=cfg, context=ctx)
s3 = strategy.compute_score(key_id="k3", config=cfg, context=ctx)
assert s1 is not None and s2 is not None and s3 is not None
assert s2 < s3 < s1
def test_multi_score_preset_priority_first_prefers_low_internal_priority() -> None:
strategy = MultiScoreStrategy()
cfg = PoolConfig(
scheduling_mode="multi_score",
scheduling_presets=(SchedulingPreset(preset="priority_first", enabled=True),),
)
ctx = {
"all_key_ids": ["k1", "k2", "k3"],
"lru_scores": {"k1": 100.0, "k2": 100.0, "k3": 100.0},
"keys_by_id": {
"k1": _key_with_metadata({}, internal_priority=20),
"k2": _key_with_metadata({}, internal_priority=3),
"k3": _key_with_metadata({}, internal_priority=11),
},
}
s1 = strategy.compute_score(key_id="k1", config=cfg, context=ctx)
s2 = strategy.compute_score(key_id="k2", config=cfg, context=ctx)

View File

@@ -14,7 +14,16 @@ def _key(metadata: dict, *, plan_type: str | None = None) -> SimpleNamespace:
def test_registry_discovers_builtin_dimensions() -> None:
names = get_preset_names()
assert {"free_team_first", "recent_refresh", "quota_balanced", "single_account"}.issubset(names)
assert {
"free_team_first",
"recent_refresh",
"quota_balanced",
"single_account",
"priority_first",
"health_first",
"latency_first",
"cost_first",
}.issubset(names)
def test_universal_dimensions_are_applicable_to_any_provider() -> None:
@@ -78,6 +87,7 @@ def test_builtin_dimensions_compute_metric_in_range() -> None:
all_key_ids=all_key_ids,
keys_by_id=keys_by_id,
lru_scores=lru_scores,
context={},
mode=mode,
)
assert 0.0 <= metric <= 1.0

View File

@@ -33,10 +33,20 @@ class _FakeResponse:
status_code: int,
payload: Any = None,
json_exc: Exception | None = None,
*,
headers: dict[str, str] | None = None,
text: str | None = None,
) -> None:
self.status_code = status_code
self._payload = payload
self._json_exc = json_exc
self.headers = headers or {}
if text is not None:
self.text = text
elif isinstance(payload, (dict, list)):
self.text = json.dumps(payload, ensure_ascii=False)
else:
self.text = ""
def json(self) -> Any:
if self._json_exc:
@@ -139,6 +149,121 @@ async def test_codex_refresher_http_non_200_returns_error(
assert result["status_code"] == 503
@pytest.mark.asyncio
async def test_codex_refresher_http_401_marks_auth_invalid_and_disables(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from src.services.provider_keys.quota_refresh import codex_refresher as module
key = SimpleNamespace(
id="k1", name="K1", api_key="enc", auth_type="api_key", auth_config=None, proxy=None
)
provider = SimpleNamespace(proxy=None)
endpoint = SimpleNamespace()
metadata_updates: dict[str, dict[str, Any]] = {}
state_updates: dict[str, dict[str, Any]] = {}
async def _fake_auth_info(_endpoint: Any, _key: Any) -> Any:
return None
_install_module(
monkeypatch,
"src.services.proxy_node.resolver",
{
"resolve_effective_proxy": lambda provider_proxy, key_proxy: None,
"build_proxy_client_kwargs": lambda proxy, timeout: {"timeout": timeout},
},
)
monkeypatch.setattr(module, "get_provider_auth", _fake_auth_info)
monkeypatch.setattr(module.crypto_service, "decrypt", lambda _v: "sk-test")
response = _FakeResponse(status_code=401, payload={"error": {"message": "token expired"}})
monkeypatch.setattr(
module.httpx, "AsyncClient", lambda **kwargs: _FakeAsyncClient(response, **kwargs)
)
result = await refresh_codex_key_quota(
db=cast(Any, _FakeDB()),
provider=cast(Any, provider),
key=cast(Any, key),
endpoint=cast(Any, endpoint),
codex_wham_usage_url="https://example.test",
metadata_updates=metadata_updates,
state_updates=state_updates,
)
assert result["status"] == "auth_invalid"
assert result["status_code"] == 401
assert result["auto_disabled"] is True
assert metadata_updates == {}
assert state_updates["k1"]["is_active"] is False
assert "401" in str(state_updates["k1"]["oauth_invalid_reason"])
@pytest.mark.asyncio
async def test_codex_refresher_http_402_sets_quota_exhausted_metadata(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from src.services.provider_keys.quota_refresh import codex_refresher as module
key = SimpleNamespace(
id="k1",
name="K1",
api_key="enc-key",
auth_type="oauth",
auth_config="enc-config",
proxy=None,
)
provider = SimpleNamespace(proxy=None)
endpoint = SimpleNamespace()
metadata_updates: dict[str, dict[str, Any]] = {}
state_updates: dict[str, dict[str, Any]] = {}
async def _fake_auth_info(_endpoint: Any, _key: Any) -> Any:
return None
_install_module(
monkeypatch,
"src.services.proxy_node.resolver",
{
"resolve_effective_proxy": lambda provider_proxy, key_proxy: None,
"build_proxy_client_kwargs": lambda proxy, timeout: {"timeout": timeout},
},
)
monkeypatch.setattr(module, "get_provider_auth", _fake_auth_info)
monkeypatch.setattr(
module.crypto_service,
"decrypt",
lambda value: (
"sk-test"
if value == "enc-key"
else json.dumps({"plan_type": "team", "account_id": "acc-1"})
),
)
response = _FakeResponse(status_code=402, payload={"error": {"message": "payment required"}})
monkeypatch.setattr(
module.httpx, "AsyncClient", lambda **kwargs: _FakeAsyncClient(response, **kwargs)
)
result = await refresh_codex_key_quota(
db=cast(Any, _FakeDB()),
provider=cast(Any, provider),
key=cast(Any, key),
endpoint=cast(Any, endpoint),
codex_wham_usage_url="https://example.test",
metadata_updates=metadata_updates,
state_updates=state_updates,
)
assert result["status"] == "quota_exhausted"
assert result["status_code"] == 402
codex_meta = metadata_updates["k1"]["codex"]
assert codex_meta["plan_type"] == "team"
assert codex_meta["primary_used_percent"] == 100.0
assert codex_meta["secondary_used_percent"] == 100.0
assert state_updates["k1"]["oauth_invalid_at"] is None
assert state_updates["k1"]["oauth_invalid_reason"] is None
@pytest.mark.asyncio
async def test_codex_refresher_success_updates_metadata(
monkeypatch: pytest.MonkeyPatch,