i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 09:20:22 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(usage): 支持 Codex 配额响应头实时异步同步

Browse Source 
- 新增 parse_codex_usage_headers，统一解析响应头中的 Codex 配额信息
- 新增实时配额同步与异步调度器，按 provider_api_key_id 去重并后台落库
- 在 usage 记录与结算流程中接入配额同步投递，并在应用生命周期中启动/停止调度器
- 补充 codex_realtime_quota 与 codex_quota_sync_dispatcher 相关单元测试
This commit is contained in:
AAEE86
2026-02-28 16:35:56 +08:00
parent 10ccbf109c
commit 92e9caf57e
8 changed files with 1066 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

139
tests/services/test_provider_keys_codex_quota_sync_dispatcher.py Normal file
View File

@@ -0,0 +1,139 @@
from __future__ import annotations
import asyncio
from typing import Any, cast
import pytest
from src.services.provider_keys import codex_quota_sync_dispatcher as dispatcher_module
def test_dispatch_falls_back_to_sync_when_dispatcher_not_running(
monkeypatch: pytest.MonkeyPatch,
) -> None:
dispatcher = dispatcher_module.CodexQuotaSyncDispatcher(flush_interval_seconds=0.01)
monkeypatch.setattr(dispatcher_module, "_dispatcher_instance", dispatcher)
captured: dict[str, Any] = {}
def _fake_sync(
*,
db: Any,
provider_api_key_id: str | None,
response_headers: dict[str, Any] | None,
) -> bool:
captured["db"] = db
captured["provider_api_key_id"] = provider_api_key_id
captured["response_headers"] = response_headers
return False
monkeypatch.setattr(dispatcher_module, "sync_codex_quota_from_response_headers", _fake_sync)
fake_db = cast(Any, object())
dispatcher_module.dispatch_codex_quota_sync_from_response_headers(
provider_api_key_id="fallback-key",
response_headers={"x-codex-primary-used-percent": "12"},
db=fake_db,
)
assert captured["db"] is fake_db
assert captured["provider_api_key_id"] == "fallback-key"
assert captured["response_headers"] == {"x-codex-primary-used-percent": "12"}
@pytest.mark.asyncio
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:
flushed_batches.append(batch)
dispatcher._flush_batch_sync = _fake_flush # type: ignore[method-assign]
await dispatcher.start()
try:
assert dispatcher.enqueue(
provider_api_key_id="key-a",
response_headers={"x-codex-primary-used-percent": "1"},
)
assert dispatcher.enqueue(
provider_api_key_id="key-a",
response_headers={"x-codex-primary-used-percent": "2"},
)
assert dispatcher.enqueue(
provider_api_key_id="key-b",
response_headers={"x-codex-secondary-used-percent": "8"},
)
await asyncio.sleep(0.06)
finally:
await dispatcher.stop()
merged: dict[str, dict[str, Any]] = {}
for batch in flushed_batches:
merged.update(batch)
assert merged["key-a"] == {"x-codex-primary-used-percent": "2"}
assert merged["key-b"] == {"x-codex-secondary-used-percent": "8"}
@pytest.mark.asyncio
async def test_dispatcher_retries_batch_after_flush_error() -> None:
dispatcher = dispatcher_module.CodexQuotaSyncDispatcher(flush_interval_seconds=0.01)
flushed_batches: list[dict[str, dict[str, Any]]] = []
flush_attempts = 0
def _flaky_flush(batch: dict[str, dict[str, Any]]) -> None:
nonlocal flush_attempts
flush_attempts += 1
if flush_attempts == 1:
raise RuntimeError("temporary flush error")
flushed_batches.append(batch)
dispatcher._flush_batch_sync = _flaky_flush # type: ignore[method-assign]
await dispatcher.start()
try:
assert dispatcher.enqueue(
provider_api_key_id="retry-key",
response_headers={"x-codex-primary-used-percent": "7"},
)
await asyncio.sleep(0.12)
finally:
await dispatcher.stop()
merged: dict[str, dict[str, Any]] = {}
for batch in flushed_batches:
merged.update(batch)
assert flush_attempts >= 2
assert merged["retry-key"] == {"x-codex-primary-used-percent": "7"}
@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:
flushed_batches.append(batch)
dispatcher._flush_batch_sync = _fake_flush # type: ignore[method-assign]
await dispatcher.start()
try:
dispatcher_module.dispatch_codex_quota_sync_from_response_headers(
provider_api_key_id="async-key",
response_headers={"x-codex-plan-type": "team"},
db=None,
)
await asyncio.sleep(0.06)
finally:
await dispatcher.stop()
dispatcher_module._dispatcher_instance = None
merged: dict[str, dict[str, Any]] = {}
for batch in flushed_batches:
merged.update(batch)
assert merged["async-key"] == {"x-codex-plan-type": "team"}

259
tests/services/test_provider_keys_codex_realtime_quota.py Normal file
View File

@@ -0,0 +1,259 @@
from __future__ import annotations
from types import SimpleNamespace
from typing import Any, cast
import pytest
from src.services.provider_keys import codex_realtime_quota as realtime_module
from src.services.provider_keys.codex_realtime_quota import (
sync_codex_quota_from_response_headers,
)
from src.services.provider_keys.codex_usage_parser import (
CodexUsageParseError,
parse_codex_usage_headers,
)
class _FakeQuery:
def __init__(self, db: "_FakeDB") -> None:
self._db = db
def options(self, *args: Any, **kwargs: Any) -> "_FakeQuery":
_ = args, kwargs
return self
def filter(self, *args: Any, **kwargs: Any) -> "_FakeQuery":
_ = args, kwargs
return self
def first(self) -> Any:
return self._db.key
class _FakeDB:
def __init__(self, key: Any) -> None:
self.key = key
self.added: list[Any] = []
self.query_count = 0
def query(self, _model: Any) -> _FakeQuery:
self.query_count += 1
return _FakeQuery(self)
def add(self, obj: Any) -> None:
self.added.append(obj)
def _paid_headers(**overrides: Any) -> dict[str, Any]:
base = {
"x-codex-plan-type": "team",
"x-codex-primary-used-percent": "3",
"x-codex-secondary-used-percent": "64",
"x-codex-primary-window-minutes": "300",
"x-codex-secondary-window-minutes": "10080",
"x-codex-primary-reset-after-seconds": "411",
"x-codex-secondary-reset-after-seconds": "267545",
"x-codex-primary-reset-at": "1772259405",
"x-codex-secondary-reset-at": "1772526539",
"x-codex-credits-has-credits": "False",
"x-codex-credits-balance": "",
"x-codex-credits-unlimited": "False",
}
base.update(overrides)
return base
def _as_session(db: _FakeDB) -> Any:
# 测试中使用最小假对象模拟 Session 接口
return cast(Any, db)
def test_parse_codex_usage_headers_paid_windows_mapping() -> None:
parsed = parse_codex_usage_headers(_paid_headers())
assert parsed is not None
assert parsed["plan_type"] == "team"
# 与 wham 解析保持一致：primary=周限额，secondary=5H 限额
assert parsed["primary_used_percent"] == 64.0
assert parsed["secondary_used_percent"] == 3.0
assert parsed["primary_window_minutes"] == 10080
assert parsed["secondary_window_minutes"] == 300
assert parsed["has_credits"] is False
assert parsed["credits_unlimited"] is False
assert "credits_balance" not in parsed
def test_parse_codex_usage_headers_free_uses_primary_only() -> None:
parsed = parse_codex_usage_headers(
{
"x-codex-plan-type": "FREE",
"x-codex-primary-used-percent": "12.5",
"x-codex-primary-window-minutes": "10080",
"x-codex-primary-reset-after-seconds": "120",
"x-codex-primary-reset-at": "1700000000",
}
)
assert parsed is not None
assert parsed["plan_type"] == "free"
assert parsed["primary_used_percent"] == 12.5
assert parsed["primary_window_minutes"] == 10080
assert "secondary_used_percent" not in parsed
def test_parse_codex_usage_headers_invalid_type_raises() -> None:
with pytest.raises(CodexUsageParseError, match="headers.primary_window.limit_window_minutes"):
parse_codex_usage_headers({"x-codex-primary-window-minutes": "abc"})
def test_sync_codex_quota_from_headers_updates_and_preserves_existing_fields() -> None:
realtime_module._header_fingerprint_cache.clear()
key = SimpleNamespace(
id="sync-update-key",
provider=SimpleNamespace(provider_type="codex"),
upstream_metadata={
"codex": {
"primary_used_percent": 50.0,
"code_review_used_percent": 12.0,
}
},
)
db = _FakeDB(key)
updated = sync_codex_quota_from_response_headers(
db=_as_session(db),
provider_api_key_id="sync-update-key",
response_headers=_paid_headers(),
)
assert updated is True
assert db.added == [key]
codex_meta = key.upstream_metadata["codex"]
assert codex_meta["primary_used_percent"] == 64.0
assert codex_meta["secondary_used_percent"] == 3.0
# 旧的 code_review 字段应被保留（wham/usage 补充信息）
assert codex_meta["code_review_used_percent"] == 12.0
def test_sync_codex_quota_from_headers_skips_when_only_reset_seconds_changed() -> None:
realtime_module._header_fingerprint_cache.clear()
key = SimpleNamespace(
id="sync-volatile-key",
provider=SimpleNamespace(provider_type="codex"),
upstream_metadata={
"codex": {
"plan_type": "team",
"primary_used_percent": 64.0,
"secondary_used_percent": 3.0,
"primary_window_minutes": 10080,
"secondary_window_minutes": 300,
"primary_reset_at": 1772526539,
"secondary_reset_at": 1772259405,
"primary_reset_seconds": 999999,
"secondary_reset_seconds": 999999,
"has_credits": False,
"credits_unlimited": False,
}
},
)
db = _FakeDB(key)
updated = sync_codex_quota_from_response_headers(
db=_as_session(db),
provider_api_key_id="sync-volatile-key",
response_headers=_paid_headers(),
)
assert updated is False
assert db.added == []
def test_sync_codex_quota_from_headers_cache_hit_skips_second_query() -> None:
realtime_module._header_fingerprint_cache.clear()
key = SimpleNamespace(
id="sync-cache-key",
provider=SimpleNamespace(provider_type="codex"),
upstream_metadata={},
)
db = _FakeDB(key)
headers = _paid_headers()
first_updated = sync_codex_quota_from_response_headers(
db=_as_session(db),
provider_api_key_id="sync-cache-key",
response_headers=headers,
)
second_updated = sync_codex_quota_from_response_headers(
db=_as_session(db),
provider_api_key_id="sync-cache-key",
response_headers=headers,
)
assert first_updated is True
assert second_updated is False
assert db.query_count == 1
def test_sync_codex_quota_from_headers_non_codex_key_is_ignored() -> None:
realtime_module._header_fingerprint_cache.clear()
key = SimpleNamespace(
id="sync-kiro-key",
provider=SimpleNamespace(provider_type="kiro"),
upstream_metadata={},
)
db = _FakeDB(key)
updated = sync_codex_quota_from_response_headers(
db=_as_session(db),
provider_api_key_id="sync-kiro-key",
response_headers=_paid_headers(),
)
assert updated is False
assert db.added == []
def test_sync_codex_quota_from_headers_parse_error_is_non_blocking() -> None:
realtime_module._header_fingerprint_cache.clear()
key = SimpleNamespace(
id="sync-parse-error-key",
provider=SimpleNamespace(provider_type="codex"),
upstream_metadata={},
)
db = _FakeDB(key)
updated = sync_codex_quota_from_response_headers(
db=_as_session(db),
provider_api_key_id="sync-parse-error-key",
response_headers={
"x-codex-primary-window-minutes": "abc",
},
)
assert updated is False
assert db.query_count == 0
assert db.added == []
def test_header_fingerprint_cache_prunes_expired_entries() -> None:
realtime_module._header_fingerprint_cache.clear()
realtime_module._set_cached_fingerprint("expired-key", "fp-1", now_ts=10.0)
realtime_module._set_cached_fingerprint("fresh-key", "fp-2", now_ts=10.0 + 31.0)
assert "expired-key" not in realtime_module._header_fingerprint_cache
assert "fresh-key" in realtime_module._header_fingerprint_cache
def test_header_fingerprint_cache_respects_max_entries() -> None:
realtime_module._header_fingerprint_cache.clear()
max_entries = realtime_module._CACHE_MAX_ENTRIES
for i in range(max_entries + 8):
realtime_module._set_cached_fingerprint(f"key-{i}", f"fp-{i}", now_ts=1000.0)
assert len(realtime_module._header_fingerprint_cache) == max_entries