mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
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:
18
src/main.py
18
src/main.py
@@ -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
|
||||
|
||||
234
src/services/provider_keys/codex_quota_sync_dispatcher.py
Normal file
234
src/services/provider_keys/codex_quota_sync_dispatcher.py
Normal 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,
|
||||
)
|
||||
156
src/services/provider_keys/codex_realtime_quota.py
Normal file
156
src/services/provider_keys/codex_realtime_quota.py
Normal 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
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user