refactor: 拆分调度器为独立子模块,增强 Sub2API 多认证方式支持

调度器重构:
- 将 CacheAwareScheduler 拆分为 candidate_builder、candidate_sorter、
  concurrency_checker、restriction_checker、scheduling_config、schemas、utils 等独立模块
- 删除旧的 _candidate_builder.py 和 _candidate_sorter.py
- 新增调度并发拒绝 Prometheus 指标

Sub2API 架构增强:
- 支持账号密码登录和 Refresh Token 两种认证方式
- 实现 JWT 自动刷新和 Token Rotation 持久化
- 前端 ProviderAuthDialog 支持多认证方式切换和 credentials_schema 动态渲染
- 验证接口返回 updated_credentials 以同步轮换后的 token

其他改进:
- 并发管理器增加 RPM guard 和动态预留逻辑
- RequestCandidate 支持 mark_skipped 附加 extra_data
- TaskService 增强健壮性
- 补充相关单元测试和契约测试
This commit is contained in:
fawney19
2026-02-15 16:32:23 +08:00
parent 8a670f5524
commit 1c16b77a92
27 changed files with 2159 additions and 645 deletions

View File

@@ -0,0 +1,168 @@
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from src.services.cache.aware_scheduler import CacheAwareScheduler, ProviderCandidate
def _make_candidate(
*,
key_id: str,
global_priority: int,
needs_conversion: bool,
provider_keep_priority_on_conversion: bool,
is_skipped: bool = False,
) -> ProviderCandidate:
provider = SimpleNamespace(
id=f"p_{key_id}",
name=f"prov_{key_id}",
provider_priority=1,
keep_priority_on_conversion=provider_keep_priority_on_conversion,
)
endpoint = SimpleNamespace(id=f"e_{key_id}")
key = SimpleNamespace(
id=key_id,
internal_priority=1,
api_key="sk-test-1234567890",
global_priority_by_format={"openai:chat": global_priority},
)
return ProviderCandidate(
provider=provider,
endpoint=endpoint,
key=key,
is_cached=False,
is_skipped=is_skipped,
skip_reason="unhealthy" if is_skipped else None,
needs_conversion=needs_conversion,
provider_api_format="openai:chat",
) # type: ignore[arg-type]
@pytest.mark.asyncio
async def test_cache_affinity_hit_healthy_candidate_is_always_promoted_to_front() -> None:
"""契约:缓存亲和性命中且候选健康时,无条件置顶(覆盖降级分组/优先级)。"""
scheduler = CacheAwareScheduler()
scheduler.scheduling_mode = CacheAwareScheduler.SCHEDULING_MODE_CACHE_AFFINITY
scheduler.priority_mode = CacheAwareScheduler.PRIORITY_MODE_GLOBAL_KEY
db = MagicMock()
keep_1 = _make_candidate(
key_id="k_keep_1",
global_priority=1,
needs_conversion=False,
provider_keep_priority_on_conversion=False,
)
keep_2 = _make_candidate(
key_id="k_keep_2",
global_priority=2,
needs_conversion=False,
provider_keep_priority_on_conversion=False,
)
matched_demote = _make_candidate(
key_id="k_cached",
global_priority=0,
needs_conversion=True,
provider_keep_priority_on_conversion=False,
)
affinity = SimpleNamespace(
provider_id=matched_demote.provider.id,
endpoint_id=matched_demote.endpoint.id,
key_id=matched_demote.key.id,
request_count=7,
)
scheduler._affinity_manager = SimpleNamespace(get_affinity=AsyncMock(return_value=affinity))
with patch(
"src.services.system.config.SystemConfigService.is_keep_priority_on_conversion",
return_value=False,
):
result = await scheduler.reorder_candidates(
candidates=[keep_1, matched_demote, keep_2],
db=db,
affinity_key="a1",
api_format="openai:chat",
global_model_id="gm1",
)
assert [c.key.id for c in result] == ["k_cached", "k_keep_1", "k_keep_2"]
assert result[0] is matched_demote
assert result[0].is_cached is True
assert all(not c.is_cached for c in result[1:])
@pytest.mark.asyncio
async def test_cache_affinity_hit_skipped_candidate_is_promoted_within_its_group() -> None:
"""契约:缓存亲和性命中但候选被跳过时,只提升到其所属类别内最前面。"""
scheduler = CacheAwareScheduler()
scheduler.scheduling_mode = CacheAwareScheduler.SCHEDULING_MODE_CACHE_AFFINITY
scheduler.priority_mode = CacheAwareScheduler.PRIORITY_MODE_GLOBAL_KEY
db = MagicMock()
keep_1 = _make_candidate(
key_id="k_keep_1",
global_priority=1,
needs_conversion=False,
provider_keep_priority_on_conversion=False,
)
keep_2 = _make_candidate(
key_id="k_keep_2",
global_priority=2,
needs_conversion=False,
provider_keep_priority_on_conversion=False,
)
demote_other = _make_candidate(
key_id="k_demote_other",
global_priority=0,
needs_conversion=True,
provider_keep_priority_on_conversion=False,
)
matched_demote_skipped = _make_candidate(
key_id="k_cached",
global_priority=10,
needs_conversion=True,
provider_keep_priority_on_conversion=False,
is_skipped=True,
)
affinity = SimpleNamespace(
provider_id=matched_demote_skipped.provider.id,
endpoint_id=matched_demote_skipped.endpoint.id,
key_id=matched_demote_skipped.key.id,
request_count=3,
)
scheduler._affinity_manager = SimpleNamespace(get_affinity=AsyncMock(return_value=affinity))
with patch(
"src.services.system.config.SystemConfigService.is_keep_priority_on_conversion",
return_value=False,
):
result = await scheduler.reorder_candidates(
candidates=[keep_1, demote_other, keep_2, matched_demote_skipped],
db=db,
affinity_key="a1",
api_format="openai:chat",
global_model_id="gm1",
)
# keep 组exact整体在前matched 在 demote 组内置顶(即使 global_priority 更差)
assert [c.key.id for c in result] == [
"k_keep_1",
"k_keep_2",
"k_cached",
"k_demote_other",
]
assert result[2] is matched_demote_skipped
assert result[2].is_cached is True
assert all(not c.is_cached for i, c in enumerate(result) if i != 2)

View File

@@ -20,7 +20,7 @@ def _make_key(
@patch(
"src.services.cache._candidate_builder.health_monitor.get_circuit_breaker_status",
"src.services.cache.candidate_builder.health_monitor.get_circuit_breaker_status",
return_value=(True, None),
)
def test_kiro_quota_remaining_zero_skips(_mock_cb: MagicMock) -> None:
@@ -39,7 +39,7 @@ def test_kiro_quota_remaining_zero_skips(_mock_cb: MagicMock) -> None:
@patch(
"src.services.cache._candidate_builder.health_monitor.get_circuit_breaker_status",
"src.services.cache.candidate_builder.health_monitor.get_circuit_breaker_status",
return_value=(True, None),
)
def test_kiro_quota_remaining_positive_allows(_mock_cb: MagicMock) -> None:
@@ -58,7 +58,7 @@ def test_kiro_quota_remaining_positive_allows(_mock_cb: MagicMock) -> None:
@patch(
"src.services.cache._candidate_builder.health_monitor.get_circuit_breaker_status",
"src.services.cache.candidate_builder.health_monitor.get_circuit_breaker_status",
return_value=(True, None),
)
def test_codex_weekly_quota_exhausted_skips(_mock_cb: MagicMock) -> None:
@@ -85,7 +85,7 @@ def test_codex_weekly_quota_exhausted_skips(_mock_cb: MagicMock) -> None:
@patch(
"src.services.cache._candidate_builder.health_monitor.get_circuit_breaker_status",
"src.services.cache.candidate_builder.health_monitor.get_circuit_breaker_status",
return_value=(True, None),
)
def test_codex_5h_quota_exhausted_skips(_mock_cb: MagicMock) -> None:
@@ -111,7 +111,7 @@ def test_codex_5h_quota_exhausted_skips(_mock_cb: MagicMock) -> None:
@patch(
"src.services.cache._candidate_builder.health_monitor.get_circuit_breaker_status",
"src.services.cache.candidate_builder.health_monitor.get_circuit_breaker_status",
return_value=(True, None),
)
def test_codex_ignores_code_review_quota(_mock_cb: MagicMock) -> None:
@@ -138,7 +138,7 @@ def test_codex_ignores_code_review_quota(_mock_cb: MagicMock) -> None:
@patch(
"src.services.cache._candidate_builder.health_monitor.get_circuit_breaker_status",
"src.services.cache.candidate_builder.health_monitor.get_circuit_breaker_status",
return_value=(True, None),
)
def test_antigravity_model_quota_exhausted_skips(_mock_cb: MagicMock) -> None:
@@ -166,7 +166,7 @@ def test_antigravity_model_quota_exhausted_skips(_mock_cb: MagicMock) -> None:
@patch(
"src.services.cache._candidate_builder.health_monitor.get_circuit_breaker_status",
"src.services.cache.candidate_builder.health_monitor.get_circuit_breaker_status",
return_value=(True, None),
)
def test_antigravity_other_model_not_exhausted_allows(_mock_cb: MagicMock) -> None:
@@ -194,7 +194,7 @@ def test_antigravity_other_model_not_exhausted_allows(_mock_cb: MagicMock) -> No
@patch(
"src.services.cache._candidate_builder.health_monitor.get_circuit_breaker_status",
"src.services.cache.candidate_builder.health_monitor.get_circuit_breaker_status",
return_value=(True, None),
)
def test_antigravity_quota_uses_mapping_matched_model(_mock_cb: MagicMock) -> None:

View File

@@ -2,11 +2,11 @@ from __future__ import annotations
from types import SimpleNamespace
from typing import Any
from unittest.mock import MagicMock
from unittest.mock import AsyncMock, MagicMock
import pytest
from src.services.cache.aware_scheduler import CacheAwareScheduler
from src.services.cache.aware_scheduler import CacheAwareScheduler, ProviderCandidate
from src.services.orchestration.candidate_resolver import CandidateResolver
@@ -76,6 +76,27 @@ class _FakeScheduler:
return candidates
def _make_global_key_candidate(*, key_id: str, priority: int) -> ProviderCandidate:
provider = SimpleNamespace(
id=f"p_{key_id}",
name=f"prov_{key_id}",
provider_priority=1,
)
endpoint = SimpleNamespace(id=f"e_{key_id}")
key = SimpleNamespace(
id=key_id,
internal_priority=1,
global_priority_by_format={"openai:chat": priority},
)
return ProviderCandidate(
provider=provider,
endpoint=endpoint,
key=key,
needs_conversion=False,
provider_api_format="openai:chat",
) # type: ignore[arg-type]
@pytest.mark.asyncio
async def test_candidate_resolver_pagination_continues_on_empty_candidate_batch() -> None:
db = MagicMock()
@@ -95,3 +116,47 @@ async def test_candidate_resolver_pagination_continues_on_empty_candidate_batch(
assert global_model_id == "gm1"
assert len(candidates) == 1
assert scheduler.calls == [0, 20]
@pytest.mark.asyncio
async def test_candidate_resolver_applies_global_reorder_after_pagination() -> None:
db = MagicMock()
scheduler = CacheAwareScheduler()
scheduler.scheduling_mode = CacheAwareScheduler.SCHEDULING_MODE_FIXED_ORDER
scheduler.priority_mode = CacheAwareScheduler.PRIORITY_MODE_GLOBAL_KEY
c10 = _make_global_key_candidate(key_id="k10", priority=10)
c1 = _make_global_key_candidate(key_id="k1", priority=1)
async def _list_all_candidates(**kwargs: Any) -> tuple[list[Any], str, int]:
provider_offset = int(kwargs.get("provider_offset", 0))
provider_limit = kwargs.get("provider_limit")
assert provider_limit is not None
if provider_offset == 0:
# First provider page returns a worse candidate first.
return [c10], "gm1", int(provider_limit)
if provider_offset == int(provider_limit):
# Second page returns a better candidate and is also the last provider page.
return [c1], "gm1", 5
return [], "gm1", 0
scheduler.list_all_candidates = AsyncMock(side_effect=_list_all_candidates) # type: ignore[method-assign]
resolver = CandidateResolver(db=db, cache_scheduler=scheduler)
candidates, global_model_id = await resolver.fetch_candidates(
api_format="openai:chat",
model_name="gpt-4o",
affinity_key="a1",
user_api_key=None,
request_id="r1",
is_stream=False,
capability_requirements=None,
)
assert global_model_id == "gm1"
assert [c.key.id for c in candidates] == ["k1", "k10"]

View File

@@ -0,0 +1,201 @@
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
from src.services.cache.candidate_sorter import CandidateSorter
from src.services.cache.scheduling_config import SchedulingConfig
from src.services.cache.schemas import ProviderCandidate
def _make_candidate(
*,
key_id: str,
global_priority: int,
needs_conversion: bool,
provider_keep_priority_on_conversion: bool,
) -> ProviderCandidate:
provider = SimpleNamespace(
id=f"p_{key_id}",
name=f"prov_{key_id}",
provider_priority=1,
keep_priority_on_conversion=provider_keep_priority_on_conversion,
)
endpoint = SimpleNamespace(id=f"e_{key_id}")
key = SimpleNamespace(
id=key_id,
internal_priority=1,
global_priority_by_format={"openai:chat": global_priority},
)
return ProviderCandidate(
provider=provider,
endpoint=endpoint,
key=key,
needs_conversion=needs_conversion,
provider_api_format="openai:chat",
) # type: ignore[arg-type]
def test_priority_sort_global_key_does_not_demote_when_global_keep_priority_enabled() -> None:
db = MagicMock()
config = SchedulingConfig(
priority_mode=SchedulingConfig.PRIORITY_MODE_GLOBAL_KEY,
scheduling_mode=SchedulingConfig.SCHEDULING_MODE_FIXED_ORDER,
)
sorter = CandidateSorter(config)
exact = _make_candidate(
key_id="k_exact",
global_priority=10,
needs_conversion=False,
provider_keep_priority_on_conversion=False,
)
demoted = _make_candidate(
key_id="k_demote",
global_priority=1,
needs_conversion=True,
provider_keep_priority_on_conversion=False,
)
# 全局 keep_priority_on_conversion=True不做 needs_conversion 降级分组,纯按 global_priority 排序
with patch(
"src.services.cache.candidate_sorter.SystemConfigService.is_keep_priority_on_conversion",
return_value=True,
):
result = sorter._apply_priority_mode_sort([exact, demoted], db, None, "openai:chat")
assert [c.key.id for c in result] == ["k_demote", "k_exact"]
def test_priority_sort_global_key_demotes_convertible_when_global_keep_priority_disabled() -> None:
db = MagicMock()
config = SchedulingConfig(
priority_mode=SchedulingConfig.PRIORITY_MODE_GLOBAL_KEY,
scheduling_mode=SchedulingConfig.SCHEDULING_MODE_FIXED_ORDER,
)
sorter = CandidateSorter(config)
exact = _make_candidate(
key_id="k_exact",
global_priority=10,
needs_conversion=False,
provider_keep_priority_on_conversion=False,
)
demoted = _make_candidate(
key_id="k_demote",
global_priority=1,
needs_conversion=True,
provider_keep_priority_on_conversion=False,
)
# 全局 keep_priority_on_conversion=False需要降级的 convertible 候选整体排后
with patch(
"src.services.cache.candidate_sorter.SystemConfigService.is_keep_priority_on_conversion",
return_value=False,
):
result = sorter._apply_priority_mode_sort([exact, demoted], db, None, "openai:chat")
assert [c.key.id for c in result] == ["k_exact", "k_demote"]
def test_priority_sort_global_key_provider_keep_priority_overrides_demotion_group() -> None:
db = MagicMock()
config = SchedulingConfig(
priority_mode=SchedulingConfig.PRIORITY_MODE_GLOBAL_KEY,
scheduling_mode=SchedulingConfig.SCHEDULING_MODE_FIXED_ORDER,
)
sorter = CandidateSorter(config)
exact = _make_candidate(
key_id="k_exact",
global_priority=10,
needs_conversion=False,
provider_keep_priority_on_conversion=False,
)
convertible_keep = _make_candidate(
key_id="k_keep",
global_priority=1,
needs_conversion=True,
provider_keep_priority_on_conversion=True,
)
convertible_demote = _make_candidate(
key_id="k_demote",
global_priority=0,
needs_conversion=True,
provider_keep_priority_on_conversion=False,
)
with patch(
"src.services.cache.candidate_sorter.SystemConfigService.is_keep_priority_on_conversion",
return_value=False,
):
result = sorter._apply_priority_mode_sort(
[exact, convertible_demote, convertible_keep],
db,
None,
"openai:chat",
)
assert [c.key.id for c in result] == ["k_keep", "k_exact", "k_demote"]
def test_priority_sort_provider_mode_demotes_convertible_when_global_keep_priority_disabled() -> (
None
):
db = MagicMock()
config = SchedulingConfig(
priority_mode=SchedulingConfig.PRIORITY_MODE_PROVIDER,
scheduling_mode=SchedulingConfig.SCHEDULING_MODE_FIXED_ORDER,
)
sorter = CandidateSorter(config)
exact = _make_candidate(
key_id="k_exact",
global_priority=10,
needs_conversion=False,
provider_keep_priority_on_conversion=False,
)
demoted = _make_candidate(
key_id="k_demote",
global_priority=1,
needs_conversion=True,
provider_keep_priority_on_conversion=False,
)
with patch(
"src.services.cache.candidate_sorter.SystemConfigService.is_keep_priority_on_conversion",
return_value=False,
):
result = sorter._apply_priority_mode_sort([demoted, exact], db, None, "openai:chat")
assert [c.key.id for c in result] == ["k_exact", "k_demote"]
def test_priority_sort_provider_mode_does_not_demote_when_global_keep_priority_enabled() -> None:
db = MagicMock()
config = SchedulingConfig(
priority_mode=SchedulingConfig.PRIORITY_MODE_PROVIDER,
scheduling_mode=SchedulingConfig.SCHEDULING_MODE_FIXED_ORDER,
)
sorter = CandidateSorter(config)
exact = _make_candidate(
key_id="k_exact",
global_priority=10,
needs_conversion=False,
provider_keep_priority_on_conversion=False,
)
demoted = _make_candidate(
key_id="k_demote",
global_priority=1,
needs_conversion=True,
provider_keep_priority_on_conversion=False,
)
with patch(
"src.services.cache.candidate_sorter.SystemConfigService.is_keep_priority_on_conversion",
return_value=True,
):
result = sorter._apply_priority_mode_sort([demoted, exact], db, None, "openai:chat")
assert [c.key.id for c in result] == ["k_demote", "k_exact"]

View File

@@ -0,0 +1,40 @@
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import MagicMock
from src.services.request.candidate import RequestCandidateService
def test_mark_candidate_skipped_merges_extra_data_and_sets_fields() -> None:
candidate = SimpleNamespace(
status="available",
skip_reason=None,
finished_at=None,
status_code=None,
concurrent_requests=None,
extra_data={"needs_conversion": True},
)
query = MagicMock()
query.filter.return_value.first.return_value = candidate
db = MagicMock()
db.query.return_value = query
RequestCandidateService.mark_candidate_skipped(
db=db,
candidate_id="c1",
skip_reason="并发限制",
status_code=429,
concurrent_requests=12,
extra_data={"concurrency_denied": True},
)
assert candidate.status == "skipped"
assert candidate.skip_reason == "并发限制"
assert candidate.status_code == 429
assert candidate.concurrent_requests == 12
assert candidate.extra_data == {"needs_conversion": True, "concurrency_denied": True}
assert candidate.finished_at is not None
db.flush.assert_called_once()