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:
fawney19
2026-03-05 15:15:26 +08:00
parent fdb50a065b
commit b1be413dc0
29 changed files with 3098 additions and 852 deletions

View File

@@ -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",

View 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] ")

View File

@@ -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:

View File

@@ -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

View 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"]

View File

@@ -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")]