mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
fix(provider_keys): 提升 Codex 配额异步同步的可靠性与可观测性
- 为 flush 引入 FlushResult,统一返回更新数与重试批次 - 增加指数退避与失败日志限流,成功后重置 backoff - 批量提交失败时回退到单条提交,降低整批失败风险 - 补充测试,覆盖 flush 重试与提交失败回退场景
This commit is contained in:
@@ -8,6 +8,45 @@ import pytest
|
||||
from src.services.provider_keys import codex_quota_sync_dispatcher as dispatcher_module
|
||||
|
||||
|
||||
def _flush_ok(batch: dict[str, dict[str, Any]]) -> dispatcher_module.FlushResult:
|
||||
return dispatcher_module.FlushResult(
|
||||
queued_count=len(batch),
|
||||
updated_count=len(batch),
|
||||
retry_batch={},
|
||||
)
|
||||
|
||||
|
||||
class _FakeNestedTxn:
|
||||
def __enter__(self) -> "_FakeNestedTxn":
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
|
||||
_ = exc_type, exc, tb
|
||||
return None
|
||||
|
||||
|
||||
class _FakeSession:
|
||||
def __init__(self, *, fail_commit: bool = False) -> None:
|
||||
self.fail_commit = fail_commit
|
||||
self.commit_calls = 0
|
||||
self.rollback_calls = 0
|
||||
self.closed = False
|
||||
|
||||
def begin_nested(self) -> _FakeNestedTxn:
|
||||
return _FakeNestedTxn()
|
||||
|
||||
def commit(self) -> None:
|
||||
self.commit_calls += 1
|
||||
if self.fail_commit and self.commit_calls == 1:
|
||||
raise RuntimeError("commit failed once")
|
||||
|
||||
def rollback(self) -> None:
|
||||
self.rollback_calls += 1
|
||||
|
||||
def close(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
|
||||
def test_dispatch_falls_back_to_sync_when_dispatcher_not_running(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
@@ -46,8 +85,9 @@ async def test_dispatcher_deduplicates_headers_by_provider_api_key_id() -> None:
|
||||
dispatcher = dispatcher_module.CodexQuotaSyncDispatcher(flush_interval_seconds=0.01)
|
||||
flushed_batches: list[dict[str, dict[str, Any]]] = []
|
||||
|
||||
def _fake_flush(batch: dict[str, dict[str, Any]]) -> None:
|
||||
def _fake_flush(batch: dict[str, dict[str, Any]]) -> dispatcher_module.FlushResult:
|
||||
flushed_batches.append(batch)
|
||||
return _flush_ok(batch)
|
||||
|
||||
dispatcher._flush_batch_sync = _fake_flush # type: ignore[method-assign]
|
||||
|
||||
@@ -83,12 +123,13 @@ async def test_dispatcher_retries_batch_after_flush_error() -> None:
|
||||
flushed_batches: list[dict[str, dict[str, Any]]] = []
|
||||
flush_attempts = 0
|
||||
|
||||
def _flaky_flush(batch: dict[str, dict[str, Any]]) -> None:
|
||||
def _flaky_flush(batch: dict[str, dict[str, Any]]) -> dispatcher_module.FlushResult:
|
||||
nonlocal flush_attempts
|
||||
flush_attempts += 1
|
||||
if flush_attempts == 1:
|
||||
raise RuntimeError("temporary flush error")
|
||||
flushed_batches.append(batch)
|
||||
return _flush_ok(batch)
|
||||
|
||||
dispatcher._flush_batch_sync = _flaky_flush # type: ignore[method-assign]
|
||||
|
||||
@@ -110,14 +151,54 @@ async def test_dispatcher_retries_batch_after_flush_error() -> None:
|
||||
assert merged["retry-key"] == {"x-codex-primary-used-percent": "7"}
|
||||
|
||||
|
||||
def test_flush_batch_falls_back_to_single_item_commit(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
dispatcher = dispatcher_module.CodexQuotaSyncDispatcher(flush_interval_seconds=0.01)
|
||||
main_session = _FakeSession(fail_commit=True)
|
||||
fallback_session_a = _FakeSession()
|
||||
fallback_session_b = _FakeSession()
|
||||
sessions: list[_FakeSession] = [main_session, fallback_session_a, fallback_session_b]
|
||||
|
||||
def _fake_create_session() -> _FakeSession:
|
||||
return sessions.pop(0)
|
||||
|
||||
def _fake_sync(
|
||||
*,
|
||||
db: Any,
|
||||
provider_api_key_id: str | None,
|
||||
response_headers: dict[str, Any] | None,
|
||||
) -> bool:
|
||||
_ = db, provider_api_key_id, response_headers
|
||||
return True
|
||||
|
||||
monkeypatch.setattr(dispatcher_module, "create_session", _fake_create_session)
|
||||
monkeypatch.setattr(dispatcher_module, "sync_codex_quota_from_response_headers", _fake_sync)
|
||||
|
||||
result = dispatcher._flush_batch_sync(
|
||||
{
|
||||
"key-a": {"x-codex-primary-used-percent": "1"},
|
||||
"key-b": {"x-codex-primary-used-percent": "2"},
|
||||
}
|
||||
)
|
||||
|
||||
assert result.updated_count == 2
|
||||
assert result.retry_batch == {}
|
||||
assert main_session.rollback_calls == 1
|
||||
assert fallback_session_a.commit_calls == 1
|
||||
assert fallback_session_b.commit_calls == 1
|
||||
assert main_session.closed is True
|
||||
assert fallback_session_a.closed is True
|
||||
assert fallback_session_b.closed is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dispatch_uses_async_queue_when_dispatcher_running() -> None:
|
||||
dispatcher = dispatcher_module.CodexQuotaSyncDispatcher(flush_interval_seconds=0.01)
|
||||
flushed_batches: list[dict[str, dict[str, Any]]] = []
|
||||
dispatcher_module._dispatcher_instance = dispatcher
|
||||
|
||||
def _fake_flush(batch: dict[str, dict[str, Any]]) -> None:
|
||||
def _fake_flush(batch: dict[str, dict[str, Any]]) -> dispatcher_module.FlushResult:
|
||||
flushed_batches.append(batch)
|
||||
return _flush_ok(batch)
|
||||
|
||||
dispatcher._flush_batch_sync = _fake_flush # type: ignore[method-assign]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user