Von
e86317aa2b
- 新增 auth_group/auth_entry 配置模型、运行态管理与 SQLite 持久化 - 支持 provider 绑定 auth_group，并在代理链路中按 least_inflight 选择凭证 - 增加 429 冷却、401/403 禁用、请求数与 token 配额统计及重置能力 - 扩展 HookContext，暴露 auth_group_name 和 auth_entry_id - 重构 Provider 管理页，新增 Auth Group 管理、运行态面板与 Auth Entries YAML 编辑 - 同步更新 README、架构文档、示例配置和相关测试
453 lines
17 KiB
Python
453 lines
17 KiB
Python
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from flask import Flask
|
|
import yaml
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
|
|
from src.application.app_context import AppContext
|
|
from src.config import build_auth_group_schemas, build_provider_schemas
|
|
from src.config.auth_group_manager import AuthGroupManager, AuthGroupSelectionError
|
|
from src.config.provider_config import AuthGroupSchema, ProviderConfigSchema
|
|
from src.config.config_manager import ConfigManager
|
|
from src.repositories import AuthGroupRepository
|
|
from src.services.auth_group_service import AuthGroupService
|
|
from src.utils.database import create_connection_factory
|
|
from src.utils.local_time import parse_local_datetime
|
|
|
|
|
|
class FakeLogger:
|
|
def info(self, msg: str, *args) -> None:
|
|
del msg, args
|
|
|
|
def warning(self, msg: str, *args) -> None:
|
|
del msg, args
|
|
|
|
def error(self, msg: str, *args) -> None:
|
|
del msg, args
|
|
|
|
def debug(self, msg: str, *args) -> None:
|
|
del msg, args
|
|
|
|
|
|
class AuthGroupConfigTests(unittest.TestCase):
|
|
def test_provider_rejects_multiple_auth_bindings(self) -> None:
|
|
with self.assertRaisesRegex(ValueError, "either auth_group or api_key, not both"):
|
|
ProviderConfigSchema.from_mapping(
|
|
{
|
|
"name": "demo",
|
|
"api": "https://example.com/v1/chat/completions",
|
|
"auth_group": "pool-a",
|
|
"api_key": "sk-demo",
|
|
"model_list": ["gpt-4.1"],
|
|
}
|
|
)
|
|
|
|
def test_provider_allows_empty_auth_binding_for_hook_or_public_upstream(self) -> None:
|
|
schema = ProviderConfigSchema.from_mapping(
|
|
{
|
|
"name": "demo",
|
|
"api": "https://example.com/v1/chat/completions",
|
|
"model_list": ["gpt-4.1"],
|
|
}
|
|
)
|
|
|
|
self.assertIsNone(schema.auth_group)
|
|
self.assertIsNone(schema.api_key)
|
|
|
|
def test_build_provider_schemas_rejects_unknown_auth_group_even_when_no_groups_exist(self) -> None:
|
|
with self.assertRaisesRegex(ValueError, "unknown auth_group: pool-a"):
|
|
build_provider_schemas(
|
|
[
|
|
{
|
|
"name": "demo",
|
|
"api": "https://example.com/v1/chat/completions",
|
|
"auth_group": "pool-a",
|
|
"model_list": ["gpt-4.1"],
|
|
}
|
|
],
|
|
available_auth_group_names=set(),
|
|
)
|
|
|
|
|
|
class AuthGroupManagerTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.temp_dir = tempfile.TemporaryDirectory()
|
|
db_path = Path(self.temp_dir.name) / "auth-groups.db"
|
|
self.ctx = AppContext(
|
|
logger=FakeLogger(),
|
|
config_manager=None, # type: ignore[arg-type]
|
|
root_path=Path(self.temp_dir.name),
|
|
flask_app=Flask(__name__),
|
|
)
|
|
self.repository = AuthGroupRepository(create_connection_factory(db_path))
|
|
self.manager = AuthGroupManager(self.ctx, self.repository)
|
|
self.group = AuthGroupSchema.from_mapping(
|
|
{
|
|
"name": "pool-a",
|
|
"strategy": "least_inflight",
|
|
"cooldown_seconds_on_429": 30,
|
|
"entries": [
|
|
{
|
|
"id": "key-a",
|
|
"headers": {"Authorization": "Bearer sk-a"},
|
|
"max_concurrency": 2,
|
|
"cooldown_seconds_on_429": 45,
|
|
},
|
|
{
|
|
"id": "key-b",
|
|
"headers": {"Authorization": "Bearer sk-b"},
|
|
"max_concurrency": 2,
|
|
},
|
|
],
|
|
}
|
|
)
|
|
self.manager.load_auth_groups((self.group,))
|
|
|
|
def tearDown(self) -> None:
|
|
self.temp_dir.cleanup()
|
|
|
|
def test_least_inflight_selection_uses_rotation_to_spread_equal_load(self) -> None:
|
|
first = self.manager.acquire("pool-a")
|
|
second = self.manager.acquire("pool-a")
|
|
|
|
self.assertEqual("key-a", first.entry_id)
|
|
self.assertEqual("key-b", second.entry_id)
|
|
|
|
def test_finish_success_persists_request_and_token_usage(self) -> None:
|
|
selection = self.manager.acquire("pool-a")
|
|
self.manager.mark_request_dispatched(selection)
|
|
self.manager.finish(
|
|
selection,
|
|
status_code=200,
|
|
usage={
|
|
"prompt_tokens": 12,
|
|
"completion_tokens": 8,
|
|
"total_tokens": 20,
|
|
},
|
|
)
|
|
|
|
runtime = self.manager.get_auth_group_runtime("pool-a")
|
|
entry = next(item for item in runtime["entries"] if item["id"] == "key-a")
|
|
|
|
self.assertEqual("available", entry["status"])
|
|
self.assertEqual(1, entry["minute_request_count"])
|
|
self.assertEqual(20, entry["minute_total_tokens"])
|
|
|
|
def test_429_cools_down_only_current_entry_using_retry_after_priority(self) -> None:
|
|
selection = self.manager.acquire("pool-a")
|
|
self.manager.mark_request_dispatched(selection)
|
|
self.manager.finish(
|
|
selection,
|
|
status_code=429,
|
|
response_headers={"Retry-After": "120"},
|
|
error_message="rate limited",
|
|
)
|
|
|
|
runtime = self.manager.get_auth_group_runtime("pool-a")
|
|
cooled_entry = next(item for item in runtime["entries"] if item["id"] == "key-a")
|
|
other_entry = next(item for item in runtime["entries"] if item["id"] == "key-b")
|
|
|
|
self.assertEqual("cooldown", cooled_entry["status"])
|
|
cooldown_until = parse_local_datetime(cooled_entry["cooldown_until"])
|
|
self.assertIsNotNone(cooldown_until)
|
|
self.assertEqual("available", other_entry["status"])
|
|
|
|
next_selection = self.manager.acquire("pool-a")
|
|
self.assertEqual("key-b", next_selection.entry_id)
|
|
|
|
def test_clear_entry_cooldown_restores_entry_to_available(self) -> None:
|
|
selection = self.manager.acquire("pool-a")
|
|
self.manager.mark_request_dispatched(selection)
|
|
self.manager.finish(
|
|
selection,
|
|
status_code=429,
|
|
response_headers={"Retry-After": "120"},
|
|
error_message="rate limited",
|
|
)
|
|
|
|
self.manager.clear_entry_cooldown("pool-a", "key-a")
|
|
runtime = self.manager.get_auth_group_runtime("pool-a")
|
|
entry = next(item for item in runtime["entries"] if item["id"] == "key-a")
|
|
|
|
self.assertEqual("available", entry["status"])
|
|
self.assertIsNone(entry["cooldown_until"])
|
|
|
|
def test_401_disables_entry_until_restore(self) -> None:
|
|
selection = self.manager.acquire("pool-a")
|
|
self.manager.mark_request_dispatched(selection)
|
|
self.manager.finish(selection, status_code=401, error_message="bad key")
|
|
|
|
runtime = self.manager.get_auth_group_runtime("pool-a")
|
|
disabled_entry = next(item for item in runtime["entries"] if item["id"] == "key-a")
|
|
self.assertEqual("disabled", disabled_entry["status"])
|
|
|
|
next_selection = self.manager.acquire("pool-a")
|
|
self.assertEqual("key-b", next_selection.entry_id)
|
|
|
|
self.manager.restore_entry("pool-a", "key-a")
|
|
restored_runtime = self.manager.get_auth_group_runtime("pool-a")
|
|
restored_entry = next(item for item in restored_runtime["entries"] if item["id"] == "key-a")
|
|
self.assertEqual("available", restored_entry["status"])
|
|
|
|
def test_manual_disable_and_enable_entry_updates_runtime_status(self) -> None:
|
|
self.manager.set_entry_disabled("pool-a", "key-a", disabled=True)
|
|
disabled_runtime = self.manager.get_auth_group_runtime("pool-a")
|
|
disabled_entry = next(item for item in disabled_runtime["entries"] if item["id"] == "key-a")
|
|
|
|
self.assertEqual("disabled", disabled_entry["status"])
|
|
self.assertEqual("manual_disabled", disabled_entry["disabled_reason"])
|
|
|
|
self.manager.set_entry_disabled("pool-a", "key-a", disabled=False)
|
|
enabled_runtime = self.manager.get_auth_group_runtime("pool-a")
|
|
enabled_entry = next(item for item in enabled_runtime["entries"] if item["id"] == "key-a")
|
|
|
|
self.assertEqual("available", enabled_entry["status"])
|
|
self.assertIsNone(enabled_entry["disabled_reason"])
|
|
|
|
def test_reset_entry_minute_usage_clears_only_current_minute_bucket(self) -> None:
|
|
selection = self.manager.acquire("pool-a")
|
|
self.manager.mark_request_dispatched(selection)
|
|
self.manager.finish(
|
|
selection,
|
|
status_code=200,
|
|
usage={
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 5,
|
|
"total_tokens": 15,
|
|
},
|
|
)
|
|
|
|
before_reset = self.manager.get_auth_group_runtime("pool-a")
|
|
before_entry = next(item for item in before_reset["entries"] if item["id"] == "key-a")
|
|
self.assertEqual(1, before_entry["minute_request_count"])
|
|
self.assertEqual(15, before_entry["minute_total_tokens"])
|
|
self.assertEqual(1, before_entry["day_request_count"])
|
|
self.assertEqual(15, before_entry["day_total_tokens"])
|
|
|
|
self.manager.reset_entry_minute_usage("pool-a", "key-a")
|
|
after_reset = self.manager.get_auth_group_runtime("pool-a")
|
|
after_entry = next(item for item in after_reset["entries"] if item["id"] == "key-a")
|
|
|
|
self.assertEqual(0, after_entry["minute_request_count"])
|
|
self.assertEqual(0, after_entry["minute_total_tokens"])
|
|
self.assertEqual(1, after_entry["day_request_count"])
|
|
self.assertEqual(15, after_entry["day_total_tokens"])
|
|
|
|
def test_reset_entry_runtime_clears_runtime_flags_and_current_usage(self) -> None:
|
|
selection = self.manager.acquire("pool-a")
|
|
self.manager.mark_request_dispatched(selection)
|
|
self.manager.finish(
|
|
selection,
|
|
status_code=429,
|
|
response_headers={"Retry-After": "120"},
|
|
error_message="rate limited",
|
|
)
|
|
self.manager.set_entry_disabled("pool-a", "key-a", disabled=True)
|
|
|
|
before_reset = self.manager.get_auth_group_runtime("pool-a")
|
|
before_entry = next(item for item in before_reset["entries"] if item["id"] == "key-a")
|
|
self.assertEqual("disabled", before_entry["status"])
|
|
self.assertEqual("manual_disabled", before_entry["disabled_reason"])
|
|
self.assertIsNotNone(before_entry["cooldown_until"])
|
|
self.assertEqual(429, before_entry["last_status_code"])
|
|
self.assertEqual("rate limited", before_entry["last_error_message"])
|
|
self.assertEqual(1, before_entry["minute_request_count"])
|
|
self.assertEqual(1, before_entry["day_request_count"])
|
|
|
|
self.manager.reset_entry_runtime("pool-a", "key-a")
|
|
|
|
after_reset = self.manager.get_auth_group_runtime("pool-a")
|
|
after_entry = next(item for item in after_reset["entries"] if item["id"] == "key-a")
|
|
self.assertEqual("available", after_entry["status"])
|
|
self.assertFalse(after_entry["disabled"])
|
|
self.assertIsNone(after_entry["disabled_reason"])
|
|
self.assertIsNone(after_entry["cooldown_until"])
|
|
self.assertIsNone(after_entry["last_status_code"])
|
|
self.assertIsNone(after_entry["last_error_type"])
|
|
self.assertIsNone(after_entry["last_error_message"])
|
|
self.assertEqual(0, after_entry["minute_request_count"])
|
|
self.assertEqual(0, after_entry["day_request_count"])
|
|
self.assertEqual(0, after_entry["minute_total_tokens"])
|
|
self.assertEqual(0, after_entry["day_total_tokens"])
|
|
|
|
|
|
class AuthGroupServiceTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.temp_dir = tempfile.TemporaryDirectory()
|
|
self.root_path = Path(self.temp_dir.name)
|
|
self.config_path = self.root_path / "config.yaml"
|
|
self.db_path = self.root_path / "requests.db"
|
|
self._write_config({"auth_groups": [], "providers": []})
|
|
self.config_manager = ConfigManager(self.config_path, self.root_path)
|
|
self.ctx = AppContext(
|
|
logger=FakeLogger(),
|
|
config_manager=self.config_manager,
|
|
root_path=self.root_path,
|
|
flask_app=Flask(__name__),
|
|
)
|
|
self.repository = AuthGroupRepository(create_connection_factory(self.db_path))
|
|
self.manager = AuthGroupManager(self.ctx, self.repository)
|
|
self.reload_count = 0
|
|
|
|
def reload_callback() -> None:
|
|
self.reload_count += 1
|
|
self.config_manager.reload()
|
|
raw = self.config_manager.get_raw_config()
|
|
self.manager.load_auth_groups(build_auth_group_schemas(raw.get("auth_groups", []) or []))
|
|
|
|
self.reload_callback = reload_callback
|
|
self.service = AuthGroupService(self.ctx, self.reload_callback, self.manager)
|
|
self.reload_callback()
|
|
self.reload_count = 0
|
|
|
|
def tearDown(self) -> None:
|
|
self.temp_dir.cleanup()
|
|
|
|
def _write_config(self, payload) -> None:
|
|
with open(self.config_path, "w", encoding="utf-8") as handle:
|
|
yaml.safe_dump(payload, handle, allow_unicode=True, sort_keys=False)
|
|
|
|
def test_create_auth_group_persists_and_lists_summary(self) -> None:
|
|
created = self.service.create_auth_group(
|
|
{
|
|
"name": "pool-a",
|
|
"strategy": "least_inflight",
|
|
"cooldown_seconds_on_429": 60,
|
|
"entries": [
|
|
{
|
|
"id": "key-a",
|
|
"headers": {"Authorization": "Bearer sk-a"},
|
|
}
|
|
],
|
|
}
|
|
)
|
|
|
|
self.assertEqual("pool-a", created["name"])
|
|
self.assertEqual(1, self.reload_count)
|
|
listed = self.service.list_auth_groups()
|
|
self.assertEqual(["pool-a"], [item["name"] for item in listed])
|
|
|
|
def test_import_auth_entries_accepts_yaml_list(self) -> None:
|
|
entries = self.service.import_auth_entries(
|
|
"""
|
|
- id: key-a
|
|
headers:
|
|
Authorization: Bearer sk-a
|
|
max_concurrency: 3
|
|
- id: key-b
|
|
headers:
|
|
x-api-key: abc
|
|
"""
|
|
)
|
|
|
|
self.assertEqual(2, len(entries))
|
|
self.assertEqual("key-a", entries[0]["id"])
|
|
self.assertEqual("Bearer sk-a", entries[0]["headers"]["Authorization"])
|
|
self.assertEqual("key-b", entries[1]["id"])
|
|
|
|
def test_import_auth_entries_accepts_entries_wrapper(self) -> None:
|
|
entries = self.service.import_auth_entries(
|
|
"""
|
|
entries:
|
|
- id: key-a
|
|
headers:
|
|
Authorization: Bearer sk-a
|
|
"""
|
|
)
|
|
|
|
self.assertEqual(1, len(entries))
|
|
self.assertEqual("key-a", entries[0]["id"])
|
|
|
|
def test_import_auth_entries_rejects_invalid_yaml_shape(self) -> None:
|
|
with self.assertRaisesRegex(ValueError, "entries 必须是列表"):
|
|
self.service.import_auth_entries(
|
|
"""
|
|
entries:
|
|
id: key-a
|
|
headers:
|
|
Authorization: Bearer sk-a
|
|
"""
|
|
)
|
|
|
|
def test_import_auth_entries_rejects_duplicate_entry_ids(self) -> None:
|
|
with self.assertRaisesRegex(ValueError, "检测到重复的 Auth Entry ID: key-a"):
|
|
self.service.import_auth_entries(
|
|
"""
|
|
entries:
|
|
- id: key-a
|
|
headers:
|
|
Authorization: Bearer sk-a
|
|
- id: key-a
|
|
headers:
|
|
Authorization: Bearer sk-b
|
|
"""
|
|
)
|
|
|
|
def test_update_auth_group_renames_provider_reference(self) -> None:
|
|
self._write_config(
|
|
{
|
|
"auth_groups": [
|
|
{
|
|
"name": "pool-a",
|
|
"strategy": "least_inflight",
|
|
"entries": [{"id": "key-a", "headers": {"Authorization": "Bearer sk-a"}}],
|
|
}
|
|
],
|
|
"providers": [
|
|
{
|
|
"name": "demo",
|
|
"api": "https://example.com/v1/chat/completions",
|
|
"auth_group": "pool-a",
|
|
"model_list": ["gpt-4.1"],
|
|
}
|
|
],
|
|
}
|
|
)
|
|
self.reload_callback()
|
|
|
|
updated = self.service.update_auth_group(
|
|
"pool-a",
|
|
{
|
|
"name": "pool-b",
|
|
"strategy": "least_inflight",
|
|
"entries": [{"id": "key-a", "headers": {"Authorization": "Bearer sk-a"}}],
|
|
},
|
|
)
|
|
|
|
self.assertEqual("pool-b", updated["name"])
|
|
current_config = self.config_manager.get_raw_config()
|
|
self.assertEqual("pool-b", current_config["providers"][0]["auth_group"])
|
|
|
|
def test_delete_referenced_auth_group_is_rejected(self) -> None:
|
|
self._write_config(
|
|
{
|
|
"auth_groups": [
|
|
{
|
|
"name": "pool-a",
|
|
"strategy": "least_inflight",
|
|
"entries": [{"id": "key-a", "headers": {"Authorization": "Bearer sk-a"}}],
|
|
}
|
|
],
|
|
"providers": [
|
|
{
|
|
"name": "demo",
|
|
"api": "https://example.com/v1/chat/completions",
|
|
"auth_group": "pool-a",
|
|
"model_list": ["gpt-4.1"],
|
|
}
|
|
],
|
|
}
|
|
)
|
|
self.reload_callback()
|
|
|
|
with self.assertRaisesRegex(ValueError, "still referenced"):
|
|
self.service.delete_auth_group("pool-a")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|