mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
feat(pool): 号池额度主动探测、封禁自动清除、调度硬优先级与前端重构
- 新增 PoolQuotaProbeScheduler,按 probing_interval_minutes 主动探测静默 Key 额度 - pool_advanced 增加 probing_enabled / auto_remove_banned_keys 配置项 - error_handler 和 quota_service 支持封禁 Key 自动删除及缓存清理 - multi_score 策略从加权混合重构为硬优先级排序,引入 mutex_group 互斥组 - 指纹注入从 handler 层下移至 ClaudeCode envelope 层 - OAuth 批量导入支持 concurrency 并发参数 - 前端号池管理拆分高级设置/账号批量/代理设置为独立组件 - 号池总览接口精简,仅返回已启用调度的 Provider
This commit is contained in:
@@ -12,6 +12,10 @@ async def test_enrich_auth_config_antigravity_adds_project_id_and_email() -> Non
|
||||
auth_config: dict[str, object] = {}
|
||||
token_response: dict[str, object] = {}
|
||||
|
||||
from src.services.provider.envelope import ensure_providers_bootstrapped
|
||||
|
||||
ensure_providers_bootstrapped()
|
||||
|
||||
with (
|
||||
patch(
|
||||
"src.core.provider_oauth_utils.fetch_google_email",
|
||||
|
||||
65
tests/services/test_error_handler_auto_remove_banned.py
Normal file
65
tests/services/test_error_handler_auto_remove_banned.py
Normal file
@@ -0,0 +1,65 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, cast
|
||||
|
||||
from src.services.orchestration.error_handler import ErrorHandlerService
|
||||
|
||||
|
||||
class _FakeDB:
|
||||
def __init__(self) -> None:
|
||||
self.deleted: list[object] = []
|
||||
self.commit_count = 0
|
||||
|
||||
def delete(self, obj: object) -> None:
|
||||
self.deleted.append(obj)
|
||||
|
||||
def commit(self) -> None:
|
||||
self.commit_count += 1
|
||||
|
||||
|
||||
def _build_key() -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
id="k1",
|
||||
provider_id="p1",
|
||||
name="k1-name",
|
||||
auth_type="oauth",
|
||||
auth_config=None,
|
||||
oauth_invalid_at=None,
|
||||
oauth_invalid_reason=None,
|
||||
is_active=True,
|
||||
)
|
||||
|
||||
|
||||
def test_mark_oauth_key_blocked_auto_remove_enabled(monkeypatch: Any) -> None:
|
||||
db = _FakeDB()
|
||||
service = ErrorHandlerService(db=cast(Any, db))
|
||||
key = _build_key()
|
||||
provider = SimpleNamespace(config={"pool_advanced": {"auto_remove_banned_keys": True}})
|
||||
|
||||
monkeypatch.setattr(
|
||||
ErrorHandlerService,
|
||||
"_schedule_auto_cleanup_after_delete",
|
||||
staticmethod(lambda **kwargs: None),
|
||||
)
|
||||
|
||||
service._mark_oauth_key_blocked(cast(Any, key), "req-1", provider=cast(Any, provider))
|
||||
|
||||
assert db.commit_count == 1
|
||||
assert db.deleted == [key]
|
||||
assert key.is_active is False
|
||||
assert str(key.oauth_invalid_reason).startswith("[ACCOUNT_BLOCK] ")
|
||||
|
||||
|
||||
def test_mark_oauth_key_blocked_auto_remove_disabled() -> None:
|
||||
db = _FakeDB()
|
||||
service = ErrorHandlerService(db=cast(Any, db))
|
||||
key = _build_key()
|
||||
provider = SimpleNamespace(config={"pool_advanced": {"auto_remove_banned_keys": False}})
|
||||
|
||||
service._mark_oauth_key_blocked(cast(Any, key), "req-1", provider=cast(Any, provider))
|
||||
|
||||
assert db.commit_count == 1
|
||||
assert db.deleted == []
|
||||
assert key.is_active is False
|
||||
assert str(key.oauth_invalid_reason).startswith("[ACCOUNT_BLOCK] ")
|
||||
@@ -39,6 +39,9 @@ def test_parse_pool_config_returns_defaults_for_empty_advanced() -> None:
|
||||
assert cfg.proactive_refresh_seconds == 180
|
||||
assert cfg.health_policy_enabled is True
|
||||
assert cfg.unschedulable_rules == []
|
||||
assert cfg.probing_enabled is False
|
||||
assert cfg.probing_interval_minutes == 10
|
||||
assert cfg.auto_remove_banned_keys is False
|
||||
|
||||
|
||||
def test_parse_pool_config_overrides_values_legacy_string_list() -> None:
|
||||
@@ -66,6 +69,9 @@ def test_parse_pool_config_overrides_values_legacy_string_list() -> None:
|
||||
"overload_cooldown_seconds": 60,
|
||||
"proactive_refresh_seconds": 300,
|
||||
"health_policy_enabled": False,
|
||||
"probing_enabled": True,
|
||||
"probing_interval_minutes": 15,
|
||||
"auto_remove_banned_keys": True,
|
||||
}
|
||||
}
|
||||
)
|
||||
@@ -94,6 +100,9 @@ def test_parse_pool_config_overrides_values_legacy_string_list() -> None:
|
||||
assert cfg.overload_cooldown_seconds == 60
|
||||
assert cfg.proactive_refresh_seconds == 300
|
||||
assert cfg.health_policy_enabled is False
|
||||
assert cfg.probing_enabled is True
|
||||
assert cfg.probing_interval_minutes == 15
|
||||
assert cfg.auto_remove_banned_keys is True
|
||||
|
||||
|
||||
def test_parse_pool_config_new_object_list_format() -> None:
|
||||
|
||||
@@ -257,3 +257,58 @@ def test_multi_score_disabled_presets_are_skipped() -> None:
|
||||
# quota_balanced only: k2 (20%) should score lower (better) than k1 (80%)
|
||||
# free_team_first is disabled so plan_type should not matter
|
||||
assert s2 < s1
|
||||
|
||||
|
||||
def test_multi_score_preset_hard_priority_overrides_later_presets() -> None:
|
||||
"""Earlier preset should dominate later presets (lexicographic hard priority)."""
|
||||
strategy = MultiScoreStrategy()
|
||||
cfg = PoolConfig(
|
||||
scheduling_mode="multi_score",
|
||||
lru_enabled=False,
|
||||
scheduling_presets=(
|
||||
SchedulingPreset(preset="priority_first", enabled=True),
|
||||
SchedulingPreset(preset="quota_balanced", 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({"codex": {"primary_used_percent": 90}}, internal_priority=1),
|
||||
"k2": _key_with_metadata({"codex": {"primary_used_percent": 10}}, internal_priority=2),
|
||||
"k3": _key_with_metadata({"codex": {"primary_used_percent": 50}}, internal_priority=3),
|
||||
},
|
||||
}
|
||||
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
|
||||
# k1 has better priority_first rank even though quota_balanced is worse.
|
||||
assert s1 < s2
|
||||
|
||||
|
||||
def test_multi_score_mutex_group_selected_member_uses_group_priority_slot() -> None:
|
||||
"""Selecting single_account should keep distribution group's first priority slot."""
|
||||
strategy = MultiScoreStrategy()
|
||||
cfg = PoolConfig(
|
||||
scheduling_mode="multi_score",
|
||||
lru_enabled=False,
|
||||
scheduling_presets=(
|
||||
SchedulingPreset(preset="lru", enabled=False),
|
||||
SchedulingPreset(preset="quota_balanced", enabled=True),
|
||||
SchedulingPreset(preset="single_account", 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({"codex": {"primary_used_percent": 90}}, internal_priority=1),
|
||||
"k2": _key_with_metadata({"codex": {"primary_used_percent": 10}}, internal_priority=2),
|
||||
"k3": _key_with_metadata({"codex": {"primary_used_percent": 20}}, internal_priority=3),
|
||||
},
|
||||
}
|
||||
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
|
||||
# single_account is selected in distribution_mode and should outrank quota_balanced.
|
||||
assert s1 < s2
|
||||
|
||||
88
tests/services/test_pool_quota_probe_scheduler.py
Normal file
88
tests/services/test_pool_quota_probe_scheduler.py
Normal file
@@ -0,0 +1,88 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import SimpleNamespace
|
||||
|
||||
from src.services.provider_keys.pool_quota_probe_scheduler import _select_probe_key_ids
|
||||
|
||||
|
||||
def _key(
|
||||
key_id: str,
|
||||
*,
|
||||
last_used_at: datetime | None = None,
|
||||
upstream_metadata: dict | None = None,
|
||||
) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
id=key_id,
|
||||
last_used_at=last_used_at,
|
||||
upstream_metadata=upstream_metadata or {},
|
||||
)
|
||||
|
||||
|
||||
def test_select_probe_key_ids_selects_silent_keys_only() -> None:
|
||||
now = datetime(2026, 3, 5, 12, 0, 0, tzinfo=timezone.utc)
|
||||
now_ts = int(now.timestamp())
|
||||
|
||||
keys = [
|
||||
_key("k1"), # never used, should be probed
|
||||
_key("k2", last_used_at=now - timedelta(minutes=2)), # recently used, skip
|
||||
_key(
|
||||
"k3",
|
||||
upstream_metadata={"codex": {"updated_at": now_ts - (20 * 60)}},
|
||||
), # long-time no refresh, should be probed
|
||||
]
|
||||
|
||||
selected = _select_probe_key_ids(
|
||||
keys=keys, # type: ignore[arg-type]
|
||||
provider_type="codex",
|
||||
now_ts=now_ts,
|
||||
interval_seconds=10 * 60,
|
||||
last_probe_timestamps={},
|
||||
limit=0,
|
||||
)
|
||||
assert selected == ["k1", "k3"]
|
||||
|
||||
|
||||
def test_select_probe_key_ids_resets_probe_window_after_key_usage() -> None:
|
||||
now = datetime(2026, 3, 5, 12, 0, 0, tzinfo=timezone.utc)
|
||||
now_ts = int(now.timestamp())
|
||||
|
||||
keys = [
|
||||
_key(
|
||||
"k1",
|
||||
last_used_at=now - timedelta(seconds=30),
|
||||
upstream_metadata={"codex": {"updated_at": now_ts - (40 * 60)}},
|
||||
)
|
||||
]
|
||||
|
||||
# 上一次主动探测非常早,但 key 刚刚被真实流量使用,应跳过本次探测
|
||||
selected = _select_probe_key_ids(
|
||||
keys=keys, # type: ignore[arg-type]
|
||||
provider_type="codex",
|
||||
now_ts=now_ts,
|
||||
interval_seconds=10 * 60,
|
||||
last_probe_timestamps={"k1": now_ts - (25 * 60)},
|
||||
limit=0,
|
||||
)
|
||||
assert selected == []
|
||||
|
||||
|
||||
def test_select_probe_key_ids_applies_limit_by_oldest_anchor_first() -> None:
|
||||
now = datetime(2026, 3, 5, 12, 0, 0, tzinfo=timezone.utc)
|
||||
now_ts = int(now.timestamp())
|
||||
|
||||
keys = [
|
||||
_key("k1", last_used_at=now - timedelta(minutes=60)),
|
||||
_key("k2", last_used_at=now - timedelta(minutes=50)),
|
||||
_key("k3", last_used_at=now - timedelta(minutes=40)),
|
||||
]
|
||||
|
||||
selected = _select_probe_key_ids(
|
||||
keys=keys, # type: ignore[arg-type]
|
||||
provider_type="codex",
|
||||
now_ts=now_ts,
|
||||
interval_seconds=10 * 60,
|
||||
last_probe_timestamps={},
|
||||
limit=2,
|
||||
)
|
||||
assert selected == ["k1", "k2"]
|
||||
@@ -50,6 +50,7 @@ class _FakeDB:
|
||||
self._provider = provider
|
||||
self._keys = keys
|
||||
self.added: list[object] = []
|
||||
self.deleted: list[object] = []
|
||||
self.commit_count = 0
|
||||
|
||||
def query(self, model: Any) -> _FakeQuery:
|
||||
@@ -63,6 +64,9 @@ class _FakeDB:
|
||||
def add(self, obj: object) -> None:
|
||||
self.added.append(obj)
|
||||
|
||||
def delete(self, obj: object) -> None:
|
||||
self.deleted.append(obj)
|
||||
|
||||
def commit(self) -> None:
|
||||
self.commit_count += 1
|
||||
|
||||
@@ -148,6 +152,7 @@ async def test_refresh_provider_quota_no_active_keys_returns_empty() -> None:
|
||||
"total": 0,
|
||||
"results": [],
|
||||
"message": "没有可刷新的 Key",
|
||||
"auto_removed": 0,
|
||||
}
|
||||
|
||||
|
||||
@@ -174,6 +179,7 @@ async def test_refresh_provider_quota_empty_key_ids_returns_empty() -> None:
|
||||
"total": 0,
|
||||
"results": [],
|
||||
"message": "未提供可刷新的 Key",
|
||||
"auto_removed": 0,
|
||||
}
|
||||
|
||||
|
||||
@@ -324,3 +330,84 @@ async def test_refresh_provider_quota_handler_exception_returns_error(
|
||||
assert result["total"] == 1
|
||||
assert result["results"][0]["status"] == "error"
|
||||
assert "unit-test boom" in result["results"][0]["message"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refresh_provider_quota_auto_removes_banned_keys_when_enabled(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
provider = SimpleNamespace(
|
||||
id="p1",
|
||||
provider_type=ProviderType.CODEX,
|
||||
endpoints=[SimpleNamespace(api_format="openai:cli", is_active=True)],
|
||||
config={"pool_advanced": {"auto_remove_banned_keys": True}},
|
||||
)
|
||||
key1 = SimpleNamespace(
|
||||
id="k1",
|
||||
name="K1",
|
||||
provider_id="p1",
|
||||
allowed_models=["gpt-4o"],
|
||||
upstream_metadata={},
|
||||
is_active=True,
|
||||
oauth_invalid_reason=None,
|
||||
)
|
||||
key2 = SimpleNamespace(
|
||||
id="k2",
|
||||
name="K2",
|
||||
provider_id="p1",
|
||||
allowed_models=None,
|
||||
upstream_metadata={},
|
||||
is_active=True,
|
||||
oauth_invalid_reason=None,
|
||||
)
|
||||
db = _FakeDB(provider=provider, keys=[key1, key2])
|
||||
deleted_side_effect_calls: list[tuple[str | None, list[str] | None]] = []
|
||||
redis_cleared: list[tuple[str, str]] = []
|
||||
|
||||
async def _fake_handler(**kwargs: Any) -> dict[str, Any]:
|
||||
key = kwargs["key"]
|
||||
state_updates = kwargs["state_updates"]
|
||||
if key.id == "k1":
|
||||
state_updates[key.id] = {"is_active": False, "oauth_invalid_reason": "账户已封禁: test"}
|
||||
return {"key_id": key.id, "key_name": key.name, "status": "banned"}
|
||||
return {"key_id": key.id, "key_name": key.name, "status": "success"}
|
||||
|
||||
async def _fake_run_delete_side_effects(
|
||||
*,
|
||||
db: Any,
|
||||
provider_id: str | None,
|
||||
deleted_key_allowed_models: list[str] | None,
|
||||
) -> None:
|
||||
_ = db
|
||||
deleted_side_effect_calls.append((provider_id, deleted_key_allowed_models))
|
||||
|
||||
async def _fake_clear(provider_id: str, key_id: str) -> None:
|
||||
redis_cleared.append((provider_id, key_id))
|
||||
|
||||
monkeypatch.setattr(
|
||||
quota_service_module,
|
||||
"_select_refresh_endpoint",
|
||||
lambda provider, provider_type: provider.endpoints[0],
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
quota_service_module, "_resolve_quota_refresh_handler", lambda _: _fake_handler
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
quota_service_module, "run_delete_key_side_effects", _fake_run_delete_side_effects
|
||||
)
|
||||
monkeypatch.setattr(quota_service_module.pool_redis, "clear_cooldown", _fake_clear)
|
||||
monkeypatch.setattr(quota_service_module.pool_redis, "clear_cost", _fake_clear)
|
||||
|
||||
result = await refresh_provider_quota_for_provider(
|
||||
db=cast(Any, db),
|
||||
provider_id="p1",
|
||||
codex_wham_usage_url="https://example.test/wham/usage",
|
||||
)
|
||||
|
||||
assert result["auto_removed"] == 1
|
||||
assert db.commit_count == 1
|
||||
assert db.deleted == [key1]
|
||||
assert key2 not in db.deleted
|
||||
assert result["results"][0]["auto_removed"] is True
|
||||
assert deleted_side_effect_calls == [("p1", ["gpt-4o"])]
|
||||
assert redis_cleared == [("p1", "k1"), ("p1", "k1")]
|
||||
|
||||
Reference in New Issue
Block a user