mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00: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:
@@ -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
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user