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

- 新增 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

View File

@@ -151,6 +151,15 @@ async def lifespan(app: FastAPI) -> Any:
await init_batch_committer()
logger.info("[OK] 批量提交器已启动,数据库写入性能优化已启用")
# 初始化 Codex 配额异步同步器(请求路径仅投递事件)
logger.info("初始化 Codex 配额异步同步器...")
from src.services.provider_keys.codex_quota_sync_dispatcher import (
init_codex_quota_sync_dispatcher,
)
await init_codex_quota_sync_dispatcher()
logger.info("[OK] Codex 配额异步同步器已启动")
# 初始化 Usage 队列消费者(可选)
if config.usage_queue_enabled:
logger.info("初始化 Usage 队列消费者...")
@@ -288,6 +297,15 @@ async def lifespan(app: FastAPI) -> Any:
# 关闭时执行
logger.info("正在关闭服务...")
# 停止 Codex 配额异步同步器(停止前会 flush 待同步事件)
logger.info("停止 Codex 配额异步同步器...")
from src.services.provider_keys.codex_quota_sync_dispatcher import (
shutdown_codex_quota_sync_dispatcher,
)
await shutdown_codex_quota_sync_dispatcher()
logger.info("[OK] Codex 配额异步同步器已停止")
# 停止批量提交器(确保所有待提交的数据都被保存)
logger.info("停止批量提交器...")
from src.core.batch_committer import shutdown_batch_committer

View File

@@ -0,0 +1,234 @@
"""
Codex 配额实时同步调度器(异步去重版)。
目标:
- 请求主路径只投递事件,不阻塞在解析/查询/提交上
- 同一 provider_api_key_id 在短窗口内仅保留最后一份响应头
- 后台批量 flush 到数据库,降低请求路径抖动
"""
from __future__ import annotations
import asyncio
from threading import Lock
from typing import Any
from sqlalchemy.orm import Session
from src.core.logger import logger
from src.database.database import create_session
from src.services.provider_keys.codex_realtime_quota import sync_codex_quota_from_response_headers
class CodexQuotaSyncDispatcher:
"""Codex 配额同步异步调度器。"""
def __init__(self, flush_interval_seconds: float = 0.5) -> None:
self.flush_interval_seconds = flush_interval_seconds
self._pending: dict[str, dict[str, Any]] = {}
self._pending_lock = Lock()
self._event: asyncio.Event | None = None
self._loop: asyncio.AbstractEventLoop | None = None
self._task: asyncio.Task[None] | None = None
self._running = False
async def start(self) -> None:
if self._running:
return
self._loop = asyncio.get_running_loop()
self._event = asyncio.Event()
self._task = asyncio.create_task(self._run(), name="codex-quota-sync-dispatcher")
self._running = True
logger.info(
"Codex 配额异步同步器已启动flush_interval={}s",
self.flush_interval_seconds,
)
async def stop(self) -> None:
if not self._running:
return
task = self._task
self._running = False
self._task = None
self._loop = None
self._event = None
if task is not None:
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
logger.info("Codex 配额异步同步器已停止")
def enqueue(
self,
*,
provider_api_key_id: str | None,
response_headers: dict[str, Any] | None,
) -> bool:
"""
投递配额同步事件。
返回:
- True: 已进入异步队列
- False: 调度器未运行或参数无效(调用方可回退同步路径)
"""
loop = self._loop
event = self._event
if (
not self._running
or loop is None
or event is None
or not provider_api_key_id
or not isinstance(response_headers, dict)
):
return False
payload = dict(response_headers)
with self._pending_lock:
# 同一 key 仅保留最新响应头,天然去重
self._pending[provider_api_key_id] = payload
try:
loop.call_soon_threadsafe(event.set)
except RuntimeError as exc:
# 事件循环关闭/切换时回退同步路径,避免请求路径抛异常。
with self._pending_lock:
if self._pending.get(provider_api_key_id) is payload:
self._pending.pop(provider_api_key_id, None)
logger.warning("Codex 配额异步同步器投递失败,已回退同步路径: {}", exc)
return False
return True
async def _run(self) -> None:
assert self._event is not None
event = self._event
try:
while True:
await event.wait()
await asyncio.sleep(self.flush_interval_seconds)
batch = self._drain_pending()
if batch:
try:
await asyncio.to_thread(self._flush_batch_sync, batch)
except Exception as exc:
logger.warning("Codex 配额异步同步器 flush 失败,将在下一轮重试: {}", exc)
self._merge_back_pending(batch)
with self._pending_lock:
if not self._pending:
event.clear()
except asyncio.CancelledError:
batch = self._drain_pending()
if batch:
try:
await asyncio.to_thread(self._flush_batch_sync, batch)
except Exception as exc:
logger.warning("Codex 配额异步同步器停止时 flush 失败: {}", exc)
raise
def _drain_pending(self) -> dict[str, dict[str, Any]]:
with self._pending_lock:
if not self._pending:
return {}
batch = dict(self._pending)
self._pending.clear()
return batch
def _merge_back_pending(self, batch: dict[str, dict[str, Any]]) -> None:
if not batch:
return
with self._pending_lock:
self._pending.update(batch)
def _flush_batch_sync(self, batch: dict[str, dict[str, Any]]) -> None:
if not batch:
return
db: Session = create_session()
updated_count = 0
try:
for provider_api_key_id, response_headers in batch.items():
try:
updated = sync_codex_quota_from_response_headers(
db=db,
provider_api_key_id=provider_api_key_id,
response_headers=response_headers,
)
if updated:
db.commit()
updated_count += 1
else:
db.rollback()
except Exception as exc:
db.rollback()
logger.warning(
"异步同步 Codex 配额失败,已跳过: provider_api_key_id={}, error={}",
provider_api_key_id,
exc,
)
if updated_count > 0:
logger.debug(
"异步同步 Codex 配额完成: queued_keys={}, updated_keys={}",
len(batch),
updated_count,
)
finally:
db.close()
_dispatcher_instance: CodexQuotaSyncDispatcher | None = None
def get_codex_quota_sync_dispatcher() -> CodexQuotaSyncDispatcher:
global _dispatcher_instance
if _dispatcher_instance is None:
_dispatcher_instance = CodexQuotaSyncDispatcher()
return _dispatcher_instance
async def init_codex_quota_sync_dispatcher() -> CodexQuotaSyncDispatcher:
dispatcher = get_codex_quota_sync_dispatcher()
await dispatcher.start()
return dispatcher
async def shutdown_codex_quota_sync_dispatcher() -> None:
global _dispatcher_instance
if _dispatcher_instance is None:
return
await _dispatcher_instance.stop()
_dispatcher_instance = None
def dispatch_codex_quota_sync_from_response_headers(
*,
provider_api_key_id: str | None,
response_headers: dict[str, Any] | None,
db: Session | None = None,
) -> None:
"""
投递 Codex 配额同步事件。
正常路径:
- 调度器已启动:异步去重后后台落库
回退路径:
- 调度器未启动且提供了 db退回同步执行避免数据丢失
"""
dispatcher = get_codex_quota_sync_dispatcher()
queued = dispatcher.enqueue(
provider_api_key_id=provider_api_key_id,
response_headers=response_headers,
)
if queued:
return
if db is not None:
sync_codex_quota_from_response_headers(
db=db,
provider_api_key_id=provider_api_key_id,
response_headers=response_headers,
)

View File

@@ -0,0 +1,156 @@
"""
Codex 配额实时同步(基于响应头)。
"""
from __future__ import annotations
import json
import time
from threading import Lock
from typing import Any
from sqlalchemy.orm import Session, joinedload
from src.core.logger import logger
from src.core.provider_types import ProviderType, normalize_provider_type
from src.models.database import ProviderAPIKey
from src.services.model.upstream_fetcher import merge_upstream_metadata
from src.services.provider_keys.codex_usage_parser import (
CodexUsageParseError,
parse_codex_usage_headers,
)
_COMPARE_IGNORE_FIELDS = frozenset(
{
"updated_at",
"primary_reset_seconds",
"secondary_reset_seconds",
"code_review_reset_seconds",
}
)
_CACHE_TTL_SECONDS = 30.0
_CACHE_MAX_ENTRIES = 4096
_header_fingerprint_cache: dict[str, tuple[str, float]] = {}
_cache_lock = Lock()
def _fingerprint_payload(payload: dict[str, Any]) -> str:
return json.dumps(payload, sort_keys=True, ensure_ascii=True, default=str)
def _build_compare_payload(data: dict[str, Any]) -> dict[str, Any]:
if not isinstance(data, dict):
return {}
return {k: v for k, v in data.items() if k not in _COMPARE_IGNORE_FIELDS}
def _get_cached_fingerprint(key_id: str, now_ts: float) -> str | None:
with _cache_lock:
cached = _header_fingerprint_cache.get(key_id)
if not cached:
return None
fp, expires_at = cached
if expires_at <= now_ts:
_header_fingerprint_cache.pop(key_id, None)
return None
return fp
def _set_cached_fingerprint(key_id: str, fingerprint: str, now_ts: float) -> None:
with _cache_lock:
_header_fingerprint_cache[key_id] = (fingerprint, now_ts + _CACHE_TTL_SECONDS)
_prune_cache_locked(now_ts)
def _prune_cache_locked(now_ts: float) -> None:
expired_keys = [
key for key, (_, expires_at) in _header_fingerprint_cache.items() if expires_at <= now_ts
]
for key in expired_keys:
_header_fingerprint_cache.pop(key, None)
overflow = len(_header_fingerprint_cache) - _CACHE_MAX_ENTRIES
if overflow <= 0:
return
keys_by_expiry = sorted(_header_fingerprint_cache.items(), key=lambda item: item[1][1])
for key, _ in keys_by_expiry[:overflow]:
_header_fingerprint_cache.pop(key, None)
def sync_codex_quota_from_response_headers(
*,
db: Session,
provider_api_key_id: str | None,
response_headers: dict[str, Any] | None,
) -> bool:
"""
从响应头同步 Codex 配额到 ProviderAPIKey.upstream_metadata。
返回值:
- True: 已产生数据库更新(由调用方统一 commit
- False: 无更新(无配额头/命中缓存/内容未变化/非 codex key
"""
if not provider_api_key_id or not isinstance(response_headers, dict):
return False
try:
parsed = parse_codex_usage_headers(response_headers)
except CodexUsageParseError as exc:
logger.warning(
"实时同步 Codex 配额头解析失败,已跳过: provider_api_key_id={}, error={}",
provider_api_key_id,
exc,
)
return False
if not parsed:
return False
now_ts = time.time()
incoming_compare = _build_compare_payload(parsed)
incoming_fingerprint = _fingerprint_payload(incoming_compare)
cached_fp = _get_cached_fingerprint(provider_api_key_id, now_ts)
if cached_fp == incoming_fingerprint:
return False
key = (
db.query(ProviderAPIKey)
.options(joinedload(ProviderAPIKey.provider))
.filter(ProviderAPIKey.id == provider_api_key_id)
.first()
)
if key is None:
_set_cached_fingerprint(provider_api_key_id, incoming_fingerprint, now_ts)
return False
provider_type = normalize_provider_type(
getattr(getattr(key, "provider", None), "provider_type", None)
)
if provider_type != ProviderType.CODEX:
_set_cached_fingerprint(provider_api_key_id, incoming_fingerprint, now_ts)
return False
current_metadata = key.upstream_metadata if isinstance(key.upstream_metadata, dict) else {}
current_codex = current_metadata.get("codex")
if not isinstance(current_codex, dict):
current_codex = {}
merged_codex = dict(current_codex)
merged_codex.update(parsed)
current_fingerprint = _fingerprint_payload(_build_compare_payload(current_codex))
merged_fingerprint = _fingerprint_payload(_build_compare_payload(merged_codex))
if current_fingerprint == merged_fingerprint:
_set_cached_fingerprint(provider_api_key_id, merged_fingerprint, now_ts)
return False
key.upstream_metadata = merge_upstream_metadata(
current_metadata,
{
"codex": merged_codex,
},
)
db.add(key)
_set_cached_fingerprint(provider_api_key_id, merged_fingerprint, now_ts)
return True

View File

@@ -5,6 +5,7 @@ Codex 配额响应解析器。
from __future__ import annotations
import time
from collections.abc import Mapping
from typing import Any
@@ -117,6 +118,84 @@ def _write_window(
)
def _normalize_plan_type(value: Any) -> str | None:
if not isinstance(value, str):
return None
normalized = value.strip().lower()
return normalized or None
def _is_blank_string(value: Any) -> bool:
return isinstance(value, str) and not value.strip()
def _coerce_optional_float(value: Any, field: str) -> float | None:
if value is None or _is_blank_string(value):
return None
return _coerce_float(value, field)
def _coerce_optional_int(value: Any, field: str) -> int | None:
if value is None or _is_blank_string(value):
return None
return _coerce_int(value, field)
def _coerce_optional_bool(value: Any, field: str) -> bool | None:
if value is None or _is_blank_string(value):
return None
return _coerce_bool(value, field)
def _normalize_header_map(headers: Mapping[str, Any]) -> dict[str, Any]:
normalized: dict[str, Any] = {}
for raw_key, raw_value in headers.items():
key = str(raw_key).strip().lower()
if not key:
continue
normalized[key] = raw_value
return normalized
def _read_header_window(
*,
headers: Mapping[str, Any],
used_percent_key: str,
reset_seconds_key: str,
reset_at_key: str,
window_minutes_key: str,
source_field: str,
) -> dict[str, Any]:
window: dict[str, Any] = {}
used_percent = _coerce_optional_float(
headers.get(used_percent_key),
f"{source_field}.used_percent",
)
if used_percent is not None:
window["used_percent"] = used_percent
reset_seconds = _coerce_optional_int(
headers.get(reset_seconds_key),
f"{source_field}.reset_after_seconds",
)
if reset_seconds is not None:
window["reset_after_seconds"] = reset_seconds
reset_at = _coerce_optional_int(headers.get(reset_at_key), f"{source_field}.reset_at")
if reset_at is not None:
window["reset_at"] = reset_at
window_minutes = _coerce_optional_int(
headers.get(window_minutes_key),
f"{source_field}.limit_window_minutes",
)
if window_minutes is not None:
window["limit_window_seconds"] = window_minutes * 60
return window
def parse_codex_wham_usage_response(data: dict[str, Any]) -> dict[str, Any] | None:
"""
解析 Codex wham/usage API 响应,提取限额信息
@@ -144,7 +223,7 @@ def parse_codex_wham_usage_response(data: dict[str, Any]) -> dict[str, Any] | No
if raw_plan_type is not None:
if not isinstance(raw_plan_type, str):
_raise_type_error("plan_type", "string", raw_plan_type)
normalized_plan_type = raw_plan_type.strip().lower()
normalized_plan_type = _normalize_plan_type(raw_plan_type)
if normalized_plan_type:
plan_type = normalized_plan_type
result["plan_type"] = normalized_plan_type
@@ -208,3 +287,124 @@ def parse_codex_wham_usage_response(data: dict[str, Any]) -> dict[str, Any] | No
result["updated_at"] = int(time.time())
return result if result else None
def parse_codex_usage_headers(headers: Mapping[str, Any] | None) -> dict[str, Any] | None:
"""
解析 Codex 反代响应头中的配额信息,提取账号配额(不依赖 code review 字段)。
Team/Plus/Enterprise:
- x-codex-primary-* : 5H 限额
- x-codex-secondary-* : 周限额
Free:
- x-codex-primary-* : 周限额
"""
if headers is None:
return None
if not isinstance(headers, Mapping):
_raise_type_error("headers", "object", headers)
if not headers:
return None
normalized_headers = _normalize_header_map(headers)
if not any(k.startswith("x-codex-") for k in normalized_headers):
return None
result: dict[str, Any] = {}
plan_type = _normalize_plan_type(normalized_headers.get("x-codex-plan-type"))
if plan_type:
result["plan_type"] = plan_type
primary_window = _read_header_window(
headers=normalized_headers,
used_percent_key="x-codex-primary-used-percent",
reset_seconds_key="x-codex-primary-reset-after-seconds",
reset_at_key="x-codex-primary-reset-at",
window_minutes_key="x-codex-primary-window-minutes",
source_field="headers.primary_window",
)
secondary_window = _read_header_window(
headers=normalized_headers,
used_percent_key="x-codex-secondary-used-percent",
reset_seconds_key="x-codex-secondary-reset-after-seconds",
reset_at_key="x-codex-secondary-reset-at",
window_minutes_key="x-codex-secondary-window-minutes",
source_field="headers.secondary_window",
)
# 兼容未来可能出现的 code review header当前反代可缺失
code_review_primary = _read_header_window(
headers=normalized_headers,
used_percent_key="x-codex-code-review-primary-used-percent",
reset_seconds_key="x-codex-code-review-primary-reset-after-seconds",
reset_at_key="x-codex-code-review-primary-reset-at",
window_minutes_key="x-codex-code-review-primary-window-minutes",
source_field="headers.code_review.primary_window",
)
# 与 wham/usage 解析保持一致:
# - metadata.primary_* 统一表示周限额
# - metadata.secondary_* 统一表示 5H 限额
use_paid_windows = bool(secondary_window) and plan_type != "free"
if use_paid_windows:
_write_window(
result,
source=secondary_window,
source_field="headers.secondary_window",
target_prefix="primary",
)
_write_window(
result,
source=primary_window,
source_field="headers.primary_window",
target_prefix="secondary",
)
else:
_write_window(
result,
source=primary_window,
source_field="headers.primary_window",
target_prefix="primary",
)
_write_window(
result,
source=code_review_primary,
source_field="headers.code_review.primary_window",
target_prefix="code_review",
)
# 当前窗口挤占占比(有值才记录)
primary_over_secondary_limit = _coerce_optional_float(
normalized_headers.get("x-codex-primary-over-secondary-limit-percent"),
"headers.primary_over_secondary_limit_percent",
)
if primary_over_secondary_limit is not None:
result["primary_over_secondary_limit_percent"] = primary_over_secondary_limit
has_credits = _coerce_optional_bool(
normalized_headers.get("x-codex-credits-has-credits"),
"headers.credits.has_credits",
)
if has_credits is not None:
result["has_credits"] = has_credits
credits_balance = _coerce_optional_float(
normalized_headers.get("x-codex-credits-balance"),
"headers.credits.balance",
)
if credits_balance is not None:
result["credits_balance"] = credits_balance
credits_unlimited = _coerce_optional_bool(
normalized_headers.get("x-codex-credits-unlimited"),
"headers.credits.unlimited",
)
if credits_unlimited is not None:
result["credits_unlimited"] = credits_unlimited
if result:
result["updated_at"] = int(time.time())
return result if result else None

View File

@@ -7,6 +7,9 @@ from sqlalchemy.orm import Session
from src.core.logger import logger
from src.models.database import ApiKey, Usage, User
from src.services.provider_keys.codex_quota_sync_dispatcher import (
dispatch_codex_quota_sync_from_response_headers,
)
from src.services.system.config import SystemConfigService
@@ -310,7 +313,14 @@ class UsageLifecycleMixin:
)
.values(**values)
)
return result.rowcount == 1
finalized = result.rowcount == 1
if finalized:
dispatch_codex_quota_sync_from_response_headers(
provider_api_key_id=provider_api_key_id,
response_headers=response_headers,
db=db,
)
return finalized
@classmethod
def update_settled_billing(

View File

@@ -8,6 +8,9 @@ from sqlalchemy.orm import Session
from src.core.logger import logger
from src.models.database import ApiKey, Provider, Usage, User, UserModelUsageCount
from src.services.provider_keys.codex_quota_sync_dispatcher import (
dispatch_codex_quota_sync_from_response_headers,
)
from src.services.usage._billing_integration import UsageBillingIntegrationMixin
from src.services.usage._recording_helpers import (
METADATA_KEEP_KEYS,
@@ -208,6 +211,12 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
usage.billing_status = "settled"
usage.finalized_at = datetime.now(timezone.utc)
dispatch_codex_quota_sync_from_response_headers(
provider_api_key_id=provider_api_key_id,
response_headers=response_headers,
db=db,
)
db.commit() # 立即提交事务,释放数据库锁
return usage
@@ -400,6 +409,12 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
usage.billing_status = "settled"
usage.finalized_at = datetime.now(timezone.utc)
dispatch_codex_quota_sync_from_response_headers(
provider_api_key_id=provider_api_key_id,
response_headers=response_headers,
db=db,
)
# 提交事务
try:
db.commit()
@@ -648,6 +663,12 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
usage.billing_status = "settled"
usage.finalized_at = datetime.now(timezone.utc)
dispatch_codex_quota_sync_from_response_headers(
provider_api_key_id=provider_api_key_id,
response_headers=response_headers,
db=db,
)
try:
db.commit()
except Exception as e:
@@ -764,6 +785,7 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
int
) # (user_id, model) -> count
provider_costs: dict[str, float] = defaultdict(float) # provider_id -> cost
quota_update_candidates: dict[str, dict[str, Any]] = {}
# 合并所有需要处理的记录(用于预取 user/api_key
all_records = records_to_insert + records_to_update
@@ -922,6 +944,15 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
apikey_stats[key_id]["cost"] += total_cost
apikey_stats[key_id]["is_standalone"] = api_key.is_standalone
provider_api_key_id = record.get("provider_api_key_id")
response_headers = record.get("response_headers")
if (
isinstance(provider_api_key_id, str)
and provider_api_key_id
and isinstance(response_headers, dict)
):
quota_update_candidates[provider_api_key_id] = response_headers
except Exception as e:
skipped_count += 1
logger.warning("批量记录中更新失败: {}, request_id={}", e, request_id)
@@ -974,6 +1005,15 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
apikey_stats[key_id]["cost"] += total_cost
apikey_stats[key_id]["is_standalone"] = api_key.is_standalone
provider_api_key_id = record.get("provider_api_key_id")
response_headers = record.get("response_headers")
if (
isinstance(provider_api_key_id, str)
and provider_api_key_id
and isinstance(response_headers, dict)
):
quota_update_candidates[provider_api_key_id] = response_headers
except Exception as e:
skipped_count += 1
logger.warning("批量记录中跳过无效记录: {}, request_id={}", e, request_id)
@@ -1090,6 +1130,14 @@ class UsageRecordingMixin(UsageBillingIntegrationMixin):
)
)
# 配额头实时同步:同一 key 仅取本批次最后一组响应头并执行一次对比更新。
for provider_api_key_id, response_headers in quota_update_candidates.items():
dispatch_codex_quota_sync_from_response_headers(
provider_api_key_id=provider_api_key_id,
response_headers=response_headers,
db=db,
)
# 单次提交所有更改
try:
db.commit()