mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
refactor: 移除 Python 后端源码，全面迁移至 Rust gateway 架构
- 删除全部 Python 源码 (src/) 及 Alembic 迁移脚本，归档至 _deprecated_py_src/ - 重构 Rust gateway ai_pipeline: 拆分 planner/finalize 模块，新增 contracts/adaptation 层 - 重组 handlers 模块为 admin/public/proxy/internal/shared 子模块结构 - 新增 executor 模块，引入 Rust 原生数据库迁移 (aether-data/migrations) - 简化 CI/Docker 构建流程，移除 base image 二级构建，统一为单一 app image - 移除 Python 相关基础设施文件 (entrypoint.sh, gunicorn_conf.py, Dockerfile.base)
This commit is contained in:
fawney19
126
_deprecated_py_src/services/provider_keys/__init__.py
Normal file
126
_deprecated_py_src/services/provider_keys/__init__.py
Normal file
@@ -0,0 +1,126 @@
|
||||
"""
|
||||
Provider Keys 领域服务模块。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.models.endpoint_models import (
|
||||
EndpointAPIKeyCreate,
|
||||
EndpointAPIKeyResponse,
|
||||
EndpointAPIKeyUpdate,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"batch_delete_endpoint_keys_response",
|
||||
"clear_oauth_invalid_response",
|
||||
"create_provider_key_response",
|
||||
"delete_endpoint_key_response",
|
||||
"update_endpoint_key_response",
|
||||
"reveal_endpoint_key_payload",
|
||||
"export_oauth_key_data",
|
||||
"get_keys_grouped_by_format",
|
||||
"list_provider_keys_responses",
|
||||
"refresh_provider_quota_for_provider",
|
||||
]
|
||||
|
||||
|
||||
def clear_oauth_invalid_response(db: Session, key_id: str) -> dict[str, str]:
|
||||
"""清除 OAuth 失效标记并返回统一响应（惰性导入实现）。"""
|
||||
from src.services.provider_keys.key_command_service import clear_oauth_invalid_response as _impl
|
||||
|
||||
return _impl(db=db, key_id=key_id)
|
||||
|
||||
|
||||
async def create_provider_key_response(
|
||||
db: Session,
|
||||
provider_id: str,
|
||||
key_data: EndpointAPIKeyCreate,
|
||||
) -> EndpointAPIKeyResponse:
|
||||
"""创建 Provider Key 并返回响应对象（惰性导入实现）。"""
|
||||
from src.services.provider_keys.key_command_service import create_provider_key_response as _impl
|
||||
|
||||
return await _impl(db=db, provider_id=provider_id, key_data=key_data)
|
||||
|
||||
|
||||
async def batch_delete_endpoint_keys_response(db: Session, key_ids: list[str]) -> dict:
|
||||
"""批量删除 Keys 并返回统一响应（惰性导入实现）。"""
|
||||
from src.services.provider_keys.key_command_service import (
|
||||
batch_delete_endpoint_keys_response as _impl,
|
||||
)
|
||||
|
||||
return await _impl(db=db, key_ids=key_ids)
|
||||
|
||||
|
||||
async def delete_endpoint_key_response(db: Session, key_id: str) -> dict[str, str]:
|
||||
"""删除 Key 并返回统一响应（惰性导入实现）。"""
|
||||
from src.services.provider_keys.key_command_service import delete_endpoint_key_response as _impl
|
||||
|
||||
return await _impl(db=db, key_id=key_id)
|
||||
|
||||
|
||||
async def update_endpoint_key_response(
|
||||
db: Session,
|
||||
key_id: str,
|
||||
key_data: EndpointAPIKeyUpdate,
|
||||
) -> EndpointAPIKeyResponse:
|
||||
"""更新 Key 并返回响应对象（惰性导入实现）。"""
|
||||
from src.services.provider_keys.key_command_service import update_endpoint_key_response as _impl
|
||||
|
||||
return await _impl(db=db, key_id=key_id, key_data=key_data)
|
||||
|
||||
|
||||
def reveal_endpoint_key_payload(db: Session, key_id: str) -> dict[str, Any]:
|
||||
"""获取完整的 API Key 或 Auth Config（惰性导入实现）。"""
|
||||
from src.services.provider_keys.key_query_service import reveal_endpoint_key_payload as _impl
|
||||
|
||||
return _impl(db=db, key_id=key_id)
|
||||
|
||||
|
||||
def export_oauth_key_data(db: Session, key_id: str) -> dict[str, Any]:
|
||||
"""导出 OAuth Key 凭据（惰性导入实现）。"""
|
||||
from src.services.provider_keys.key_query_service import export_oauth_key_data as _impl
|
||||
|
||||
return _impl(db=db, key_id=key_id)
|
||||
|
||||
|
||||
def get_keys_grouped_by_format(db: Session) -> dict:
|
||||
"""按 API 格式分组查询所有 Key（惰性导入实现）。"""
|
||||
from src.services.provider_keys.key_query_service import get_keys_grouped_by_format as _impl
|
||||
|
||||
return _impl(db=db)
|
||||
|
||||
|
||||
def list_provider_keys_responses(
|
||||
db: Session,
|
||||
provider_id: str,
|
||||
skip: int,
|
||||
limit: int,
|
||||
) -> list[EndpointAPIKeyResponse]:
|
||||
"""查询 Provider 下的 Key 列表（惰性导入实现）。"""
|
||||
from src.services.provider_keys.key_query_service import list_provider_keys_responses as _impl
|
||||
|
||||
return _impl(db=db, provider_id=provider_id, skip=skip, limit=limit)
|
||||
|
||||
|
||||
async def refresh_provider_quota_for_provider(
|
||||
db: Session,
|
||||
provider_id: str,
|
||||
codex_wham_usage_url: str,
|
||||
key_ids: list[str] | None = None,
|
||||
) -> dict:
|
||||
"""刷新 Provider 限额信息（惰性导入实现）。"""
|
||||
from src.services.provider_keys.key_quota_service import (
|
||||
refresh_provider_quota_for_provider as _impl,
|
||||
)
|
||||
|
||||
return await _impl(
|
||||
db=db,
|
||||
provider_id=provider_id,
|
||||
codex_wham_usage_url=codex_wham_usage_url,
|
||||
key_ids=key_ids,
|
||||
)
|
||||
21
_deprecated_py_src/services/provider_keys/auth_type.py
Normal file
21
_deprecated_py_src/services/provider_keys/auth_type.py
Normal file
@@ -0,0 +1,21 @@
|
||||
"""
|
||||
Provider Key 认证类型相关规则。
|
||||
"""
|
||||
|
||||
# 数据库中所有属于 OAuth 的 auth_type 值（含历史别名）。
|
||||
# 新增 OAuth 类型时只需在此追加，SQL 过滤和 Python 判断均引用此常量。
|
||||
OAUTH_AUTH_TYPES: tuple[str, ...] = ("oauth", "kiro")
|
||||
|
||||
|
||||
def normalize_auth_type(raw: str) -> str:
|
||||
"""将数据库中的 auth_type 归一化为逻辑类型。
|
||||
|
||||
- ``"kiro"`` -> ``"oauth"`` (Kiro 使用 OAuth 流程)
|
||||
- ``"vertex_ai"`` -> ``"service_account"`` (旧的 Vertex AI auth_type 已重命名)
|
||||
"""
|
||||
t = str(raw or "api_key").strip() or "api_key"
|
||||
if t == "kiro": # TODO: 迁移稳定后清理，同步清理各处 in ("...", "kiro") 兼容检查
|
||||
return "oauth"
|
||||
if t == "vertex_ai": # TODO: 迁移稳定后清理，同步清理各处 in ("...", "vertex_ai") 兼容检查
|
||||
return "service_account"
|
||||
return t
|
||||
488
_deprecated_py_src/services/provider_keys/batch_delete_task.py
Normal file
488
_deprecated_py_src/services/provider_keys/batch_delete_task.py
Normal file
@@ -0,0 +1,488 @@
|
||||
"""Pool Key 批量删除异步任务。
|
||||
|
||||
接口立即返回 task_id，后台执行删除，前端轮询进度。
|
||||
任务状态存储在 Redis 中，支持多 worker 进程共享。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Callable
|
||||
from concurrent.futures import Future
|
||||
|
||||
import redis.asyncio as aioredis
|
||||
from sqlalchemy import delete as sa_delete
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.clients.redis_client import get_redis_client
|
||||
from src.core.logger import logger
|
||||
from src.services.provider_keys.key_side_effects import cleanup_key_references
|
||||
|
||||
# 任务状态
|
||||
STATUS_PENDING = "pending"
|
||||
STATUS_RUNNING = "running"
|
||||
STATUS_COMPLETED = "completed"
|
||||
STATUS_FAILED = "failed"
|
||||
|
||||
# 任务完成后保留时长（秒），也用作 Redis key 的 TTL
|
||||
_TASK_RETAIN_SECONDS = 600
|
||||
|
||||
# 批量删除时每个独立事务处理的 key 数量
|
||||
_CLEANUP_BATCH_SIZE = 50
|
||||
|
||||
# 单个批次的数据库 statement 超时（秒）
|
||||
_BATCH_STATEMENT_TIMEOUT_S = 30
|
||||
|
||||
# 单个批次的数据库锁等待超时（秒）
|
||||
_BATCH_LOCK_TIMEOUT_S = 5
|
||||
|
||||
# 整个任务的最大执行时间（秒）
|
||||
_TASK_TIMEOUT_S = 600
|
||||
|
||||
# 发生超时/锁等待时降批到的最小 Key 数
|
||||
_MIN_RETRY_BATCH_SIZE = 1
|
||||
|
||||
# Redis key 前缀
|
||||
_REDIS_KEY_PREFIX = "batch_delete_task"
|
||||
|
||||
# 持有后台 asyncio.Task 引用，防止 GC 回收（进程内部，不需要跨进程共享）
|
||||
_running_tasks: set[asyncio.Task[None]] = set()
|
||||
|
||||
|
||||
def _task_key(task_id: str) -> str:
|
||||
return f"{_REDIS_KEY_PREFIX}:{task_id}"
|
||||
|
||||
|
||||
|
||||
def _apply_statement_timeouts(db: Session) -> None:
|
||||
db.execute(text(f"SET LOCAL statement_timeout = '{_BATCH_STATEMENT_TIMEOUT_S * 1000}'"))
|
||||
db.execute(text(f"SET LOCAL lock_timeout = '{_BATCH_LOCK_TIMEOUT_S * 1000}'"))
|
||||
|
||||
|
||||
def _is_retryable_batch_error(exc: Exception) -> bool:
|
||||
messages = [str(exc)]
|
||||
orig = getattr(exc, "orig", None)
|
||||
if orig is not None:
|
||||
messages.append(str(orig))
|
||||
text_blob = " ".join(messages).lower()
|
||||
return any(
|
||||
marker in text_blob
|
||||
for marker in (
|
||||
"querycanceled",
|
||||
"statement timeout",
|
||||
"canceling statement due to statement timeout",
|
||||
"lock timeout",
|
||||
"canceling statement due to lock timeout",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class BatchDeleteTaskInfo:
|
||||
"""任务状态数据对象（从 Redis 反序列化）。"""
|
||||
|
||||
__slots__ = ("task_id", "provider_id", "status", "total", "deleted", "message")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
task_id: str,
|
||||
provider_id: str,
|
||||
status: str = STATUS_PENDING,
|
||||
total: int = 0,
|
||||
deleted: int = 0,
|
||||
message: str = "",
|
||||
):
|
||||
self.task_id = task_id
|
||||
self.provider_id = provider_id
|
||||
self.status = status
|
||||
self.total = total
|
||||
self.deleted = deleted
|
||||
self.message = message
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"task_id": self.task_id,
|
||||
"provider_id": self.provider_id,
|
||||
"status": self.status,
|
||||
"total": self.total,
|
||||
"deleted": self.deleted,
|
||||
"message": self.message,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict) -> BatchDeleteTaskInfo:
|
||||
return cls(
|
||||
task_id=data["task_id"],
|
||||
provider_id=data["provider_id"],
|
||||
status=data.get("status", STATUS_PENDING),
|
||||
total=int(data.get("total", 0)),
|
||||
deleted=int(data.get("deleted", 0)),
|
||||
message=data.get("message", ""),
|
||||
)
|
||||
|
||||
|
||||
async def _save_task(
|
||||
task: BatchDeleteTaskInfo,
|
||||
ttl: int = _TASK_RETAIN_SECONDS,
|
||||
r: aioredis.Redis | None = None,
|
||||
) -> None:
|
||||
"""将任务状态写入 Redis。"""
|
||||
if r is None:
|
||||
r = await get_redis_client(require_redis=False)
|
||||
if not r:
|
||||
return
|
||||
try:
|
||||
await r.setex(_task_key(task.task_id), ttl, json.dumps(task.to_dict()))
|
||||
except Exception as e:
|
||||
logger.warning("Failed to save batch delete task to Redis: {}", e)
|
||||
|
||||
|
||||
async def _load_task(task_id: str, r: aioredis.Redis | None = None) -> BatchDeleteTaskInfo | None:
|
||||
"""从 Redis 加载任务状态。"""
|
||||
if r is None:
|
||||
r = await get_redis_client(require_redis=False)
|
||||
if not r:
|
||||
return None
|
||||
try:
|
||||
data = await r.get(_task_key(task_id))
|
||||
if data is None:
|
||||
return None
|
||||
return BatchDeleteTaskInfo.from_dict(json.loads(data))
|
||||
except Exception as e:
|
||||
logger.warning("Failed to load batch delete task from Redis: {}", e)
|
||||
return None
|
||||
|
||||
|
||||
async def _update_task_field(
|
||||
task_id: str, r: aioredis.Redis | None = None, **fields: object
|
||||
) -> None:
|
||||
"""局部更新任务状态字段（read-modify-write）。"""
|
||||
if r is None:
|
||||
r = await get_redis_client(require_redis=False)
|
||||
if not r:
|
||||
return
|
||||
task = await _load_task(task_id, r=r)
|
||||
if task is None:
|
||||
return
|
||||
for k, v in fields.items():
|
||||
setattr(task, k, v)
|
||||
# 已完成/失败的任务只保留 _TASK_RETAIN_SECONDS
|
||||
if task.status in (STATUS_COMPLETED, STATUS_FAILED):
|
||||
await _save_task(task, ttl=_TASK_RETAIN_SECONDS, r=r)
|
||||
else:
|
||||
# 运行中的任务使用更长的 TTL，防止超长任务过期
|
||||
await _save_task(task, ttl=_TASK_RETAIN_SECONDS * 2, r=r)
|
||||
|
||||
|
||||
async def submit_batch_delete(provider_id: str, key_ids: list[str]) -> str:
|
||||
"""提交批量删除任务，返回 task_id。
|
||||
|
||||
Raises:
|
||||
RuntimeError: Redis 不可用时无法追踪任务状态。
|
||||
"""
|
||||
r = await get_redis_client(require_redis=False)
|
||||
if not r:
|
||||
raise RuntimeError("Redis is required for batch delete tasks but is not available")
|
||||
|
||||
task_id = uuid.uuid4().hex[:16]
|
||||
task = BatchDeleteTaskInfo(
|
||||
task_id=task_id,
|
||||
provider_id=provider_id,
|
||||
total=len(key_ids),
|
||||
)
|
||||
await _save_task(task, ttl=_TASK_RETAIN_SECONDS * 2, r=r)
|
||||
|
||||
bg = asyncio.create_task(
|
||||
_run_batch_delete(task_id, provider_id, key_ids),
|
||||
name=f"batch-delete-{task_id}",
|
||||
)
|
||||
_running_tasks.add(bg)
|
||||
bg.add_done_callback(_running_tasks.discard)
|
||||
return task_id
|
||||
|
||||
|
||||
async def get_batch_delete_task(task_id: str) -> BatchDeleteTaskInfo | None:
|
||||
return await _load_task(task_id)
|
||||
|
||||
|
||||
def _delete_key_batch(
|
||||
db: Session,
|
||||
provider_id: str,
|
||||
batch: list[str],
|
||||
) -> int:
|
||||
from src.models.database import ProviderAPIKey
|
||||
|
||||
phase = "cleanup_key_references"
|
||||
|
||||
def _set_phase(stage_name: str, _batch_size: int) -> None:
|
||||
nonlocal phase
|
||||
phase = stage_name
|
||||
|
||||
try:
|
||||
_apply_statement_timeouts(db)
|
||||
cleanup_key_references(
|
||||
db,
|
||||
batch,
|
||||
batch_size=len(batch),
|
||||
stage_callback=_set_phase,
|
||||
)
|
||||
phase = "provider_api_keys"
|
||||
result = db.execute(
|
||||
sa_delete(ProviderAPIKey).where(
|
||||
ProviderAPIKey.provider_id == provider_id,
|
||||
ProviderAPIKey.id.in_(batch),
|
||||
)
|
||||
)
|
||||
rowcount = getattr(result, "rowcount", 0) or 0
|
||||
db.commit()
|
||||
return int(rowcount)
|
||||
except Exception as exc:
|
||||
setattr(exc, "_aether_batch_phase", phase)
|
||||
raise
|
||||
|
||||
|
||||
def _delete_key_batch_with_retry(
|
||||
db: Session,
|
||||
provider_id: str,
|
||||
batch: list[str],
|
||||
*,
|
||||
batch_idx: int,
|
||||
total_batches: int,
|
||||
start_offset: int,
|
||||
attempt: int = 1,
|
||||
) -> int:
|
||||
try:
|
||||
return _delete_key_batch(db, provider_id, batch)
|
||||
except Exception as exc:
|
||||
try:
|
||||
db.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
is_retryable = _is_retryable_batch_error(exc)
|
||||
phase = getattr(exc, "_aether_batch_phase", "unknown")
|
||||
can_split = len(batch) > _MIN_RETRY_BATCH_SIZE
|
||||
if is_retryable and can_split:
|
||||
split_at = max(len(batch) // 2, _MIN_RETRY_BATCH_SIZE)
|
||||
left = batch[:split_at]
|
||||
right = batch[split_at:]
|
||||
logger.warning(
|
||||
"[BATCH_DELETE] batch {}/{} retrying after timeout/lock (keys {}-{}, size={}, attempt={}, phase={}): split into {} + {}",
|
||||
batch_idx,
|
||||
total_batches,
|
||||
start_offset,
|
||||
start_offset + len(batch),
|
||||
len(batch),
|
||||
attempt,
|
||||
phase,
|
||||
len(left),
|
||||
len(right),
|
||||
)
|
||||
deleted = _delete_key_batch_with_retry(
|
||||
db,
|
||||
provider_id,
|
||||
left,
|
||||
batch_idx=batch_idx,
|
||||
total_batches=total_batches,
|
||||
start_offset=start_offset,
|
||||
attempt=attempt + 1,
|
||||
)
|
||||
if right:
|
||||
deleted += _delete_key_batch_with_retry(
|
||||
db,
|
||||
provider_id,
|
||||
right,
|
||||
batch_idx=batch_idx,
|
||||
total_batches=total_batches,
|
||||
start_offset=start_offset + len(left),
|
||||
attempt=attempt + 1,
|
||||
)
|
||||
return deleted
|
||||
|
||||
logger.warning(
|
||||
"[BATCH_DELETE] batch {}/{} failed (keys {}-{} size={} attempt={} phase={} retryable={}): {}",
|
||||
batch_idx,
|
||||
total_batches,
|
||||
start_offset,
|
||||
start_offset + len(batch),
|
||||
len(batch),
|
||||
attempt,
|
||||
phase,
|
||||
is_retryable,
|
||||
exc,
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
def _sync_delete(
|
||||
provider_id: str,
|
||||
key_ids: list[str],
|
||||
progress_callback: Callable[[int], None] | None = None,
|
||||
) -> int:
|
||||
"""在线程中执行的同步删除逻辑，避免阻塞事件循环。
|
||||
|
||||
按小批次（_CLEANUP_BATCH_SIZE）先显式处理关联表，再删除 key。
|
||||
每个批次独立事务，单批失败跳过并继续。
|
||||
"""
|
||||
from src.database import create_session
|
||||
|
||||
db = create_session()
|
||||
try:
|
||||
affected = 0
|
||||
total_batches = (len(key_ids) + _CLEANUP_BATCH_SIZE - 1) // _CLEANUP_BATCH_SIZE
|
||||
batch_idx = 0
|
||||
task_start = time.monotonic()
|
||||
for i in range(0, len(key_ids), _CLEANUP_BATCH_SIZE):
|
||||
# 整体超时保护
|
||||
if time.monotonic() - task_start > _TASK_TIMEOUT_S:
|
||||
logger.warning(
|
||||
"[BATCH_DELETE] task timeout after {}s, deleted {}/{}",
|
||||
_TASK_TIMEOUT_S,
|
||||
affected,
|
||||
len(key_ids),
|
||||
)
|
||||
break
|
||||
|
||||
batch = key_ids[i : i + _CLEANUP_BATCH_SIZE]
|
||||
batch_idx += 1
|
||||
affected += _delete_key_batch_with_retry(
|
||||
db,
|
||||
provider_id,
|
||||
batch,
|
||||
batch_idx=batch_idx,
|
||||
total_batches=total_batches,
|
||||
start_offset=i,
|
||||
)
|
||||
# 每个批次都上报一次进度
|
||||
if progress_callback is not None:
|
||||
progress_callback(affected)
|
||||
logger.info(
|
||||
"[BATCH_DELETE] sync delete finished: deleted={}/{} batches={} elapsed={:.1f}s",
|
||||
affected,
|
||||
len(key_ids),
|
||||
total_batches,
|
||||
time.monotonic() - task_start,
|
||||
)
|
||||
return affected
|
||||
finally:
|
||||
try:
|
||||
db.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
async def _run_batch_delete(
|
||||
task_id: str,
|
||||
provider_id: str,
|
||||
key_ids: list[str],
|
||||
) -> None:
|
||||
r = await get_redis_client(require_redis=False)
|
||||
await _update_task_field(task_id, r=r, status=STATUS_RUNNING)
|
||||
logger.info(
|
||||
"[BATCH_DELETE_TASK] started task={} provider={} total={}",
|
||||
task_id,
|
||||
provider_id[:8],
|
||||
len(key_ids),
|
||||
)
|
||||
|
||||
# 从工作线程安全地触发 Redis 进度更新
|
||||
loop = asyncio.get_running_loop()
|
||||
progress_futures: list[Future[object]] = []
|
||||
|
||||
def on_progress(current: int) -> None:
|
||||
try:
|
||||
future = asyncio.run_coroutine_threadsafe(
|
||||
_update_task_field(task_id, r=r, deleted=current), loop
|
||||
)
|
||||
progress_futures.append(future)
|
||||
except RuntimeError:
|
||||
pass
|
||||
|
||||
try:
|
||||
affected = await asyncio.wait_for(
|
||||
asyncio.to_thread(_sync_delete, provider_id, key_ids, on_progress),
|
||||
timeout=_TASK_TIMEOUT_S + 30, # 留一点余量给 _sync_delete 内部超时
|
||||
)
|
||||
|
||||
if progress_futures:
|
||||
results = await asyncio.gather(
|
||||
*(asyncio.wrap_future(f) for f in progress_futures),
|
||||
return_exceptions=True,
|
||||
)
|
||||
for exc in results:
|
||||
if isinstance(exc, Exception):
|
||||
logger.debug(
|
||||
"[BATCH_DELETE_TASK] progress update failed task={}: {}",
|
||||
task_id,
|
||||
exc,
|
||||
)
|
||||
|
||||
# 副作用（缓存失效等）是异步操作，在事件循环中执行
|
||||
if affected > 0:
|
||||
try:
|
||||
from src.database import get_db_context
|
||||
from src.services.provider_keys.key_side_effects import (
|
||||
run_delete_key_side_effects,
|
||||
)
|
||||
|
||||
with get_db_context() as db:
|
||||
await run_delete_key_side_effects(
|
||||
db=db,
|
||||
provider_id=provider_id,
|
||||
deleted_key_allowed_models=None,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error("batch delete side effects failed: {}", exc)
|
||||
|
||||
await _update_task_field(
|
||||
task_id,
|
||||
r=r,
|
||||
status=STATUS_COMPLETED,
|
||||
deleted=affected,
|
||||
message=f"{affected} keys deleted",
|
||||
)
|
||||
logger.info(
|
||||
"[BATCH_DELETE_TASK] completed task={} provider={} total={} deleted={}",
|
||||
task_id,
|
||||
provider_id[:8],
|
||||
len(key_ids),
|
||||
affected,
|
||||
)
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
msg = f"task timeout after {_TASK_TIMEOUT_S + 30}s"
|
||||
await _update_task_field(task_id, r=r, status=STATUS_FAILED, message=msg)
|
||||
logger.error(
|
||||
"[BATCH_DELETE_TASK] {} task={} provider={}",
|
||||
msg,
|
||||
task_id,
|
||||
provider_id[:8],
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
await _update_task_field(
|
||||
task_id,
|
||||
r=r,
|
||||
status=STATUS_FAILED,
|
||||
message="task cancelled (shutdown)",
|
||||
)
|
||||
logger.warning(
|
||||
"[BATCH_DELETE_TASK] cancelled task={} provider={}",
|
||||
task_id,
|
||||
provider_id[:8],
|
||||
)
|
||||
except Exception as exc:
|
||||
await _update_task_field(
|
||||
task_id,
|
||||
r=r,
|
||||
status=STATUS_FAILED,
|
||||
message=str(exc),
|
||||
)
|
||||
logger.error(
|
||||
"[BATCH_DELETE_TASK] failed task={} provider={}: {}",
|
||||
task_id,
|
||||
provider_id[:8],
|
||||
exc,
|
||||
)
|
||||
@@ -0,0 +1,358 @@
|
||||
"""
|
||||
Codex 配额实时同步调度器（异步去重版）。
|
||||
|
||||
目标：
|
||||
- 请求主路径只投递事件，不阻塞在解析/查询/提交上
|
||||
- 同一 provider_api_key_id 在短窗口内仅保留最后一份响应头
|
||||
- 后台批量 flush 到数据库，降低请求路径抖动
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
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
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class FlushResult:
|
||||
queued_count: int
|
||||
updated_count: int
|
||||
retry_batch: dict[str, dict[str, Any]]
|
||||
|
||||
|
||||
class CodexQuotaSyncDispatcher:
|
||||
"""Codex 配额同步异步调度器。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
flush_interval_seconds: float = 0.5,
|
||||
*,
|
||||
max_backoff_seconds: float = 8.0,
|
||||
error_log_interval_seconds: float = 30.0,
|
||||
) -> None:
|
||||
self.flush_interval_seconds = max(float(flush_interval_seconds), 0.001)
|
||||
self.max_backoff_seconds = max(float(max_backoff_seconds), self.flush_interval_seconds)
|
||||
self.error_log_interval_seconds = max(float(error_log_interval_seconds), 0.0)
|
||||
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._current_flush_delay_seconds = self.flush_interval_seconds
|
||||
self._last_flush_error_log_at: float | 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._current_flush_delay_seconds = self.flush_interval_seconds
|
||||
self._last_flush_error_log_at = None
|
||||
self._running = True
|
||||
logger.info(
|
||||
"Codex 配额异步同步器已启动，flush_interval={}s, max_backoff={}s",
|
||||
self.flush_interval_seconds,
|
||||
self.max_backoff_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._current_flush_delay_seconds)
|
||||
batch = self._drain_pending()
|
||||
if batch:
|
||||
try:
|
||||
result = await asyncio.to_thread(self._flush_batch_sync, batch)
|
||||
except Exception as exc:
|
||||
self._merge_back_pending(batch)
|
||||
self._increase_backoff()
|
||||
self._log_flush_failure_with_rate_limit(
|
||||
queued_count=len(batch),
|
||||
retry_count=len(batch),
|
||||
error=exc,
|
||||
)
|
||||
else:
|
||||
if result.updated_count > 0:
|
||||
logger.debug(
|
||||
"异步同步 Codex 配额完成: queued_keys={}, updated_keys={}",
|
||||
result.queued_count,
|
||||
result.updated_count,
|
||||
)
|
||||
if result.retry_batch:
|
||||
self._merge_back_pending(result.retry_batch)
|
||||
self._increase_backoff()
|
||||
self._log_flush_failure_with_rate_limit(
|
||||
queued_count=result.queued_count,
|
||||
retry_count=len(result.retry_batch),
|
||||
)
|
||||
else:
|
||||
self._reset_backoff()
|
||||
with self._pending_lock:
|
||||
if not self._pending:
|
||||
event.clear()
|
||||
except asyncio.CancelledError:
|
||||
batch = self._drain_pending()
|
||||
if batch:
|
||||
try:
|
||||
result = await asyncio.to_thread(self._flush_batch_sync, batch)
|
||||
if result.retry_batch:
|
||||
logger.warning(
|
||||
"Codex 配额异步同步器停止时仍有未同步事件: retry_keys={}",
|
||||
len(result.retry_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 _reset_backoff(self) -> None:
|
||||
self._current_flush_delay_seconds = self.flush_interval_seconds
|
||||
|
||||
def _increase_backoff(self) -> None:
|
||||
self._current_flush_delay_seconds = min(
|
||||
self.max_backoff_seconds,
|
||||
max(self.flush_interval_seconds, self._current_flush_delay_seconds * 2),
|
||||
)
|
||||
|
||||
def _log_flush_failure_with_rate_limit(
|
||||
self,
|
||||
*,
|
||||
queued_count: int,
|
||||
retry_count: int,
|
||||
error: Exception | None = None,
|
||||
) -> None:
|
||||
now = time.monotonic()
|
||||
if (
|
||||
self._last_flush_error_log_at is not None
|
||||
and now - self._last_flush_error_log_at < self.error_log_interval_seconds
|
||||
):
|
||||
return
|
||||
self._last_flush_error_log_at = now
|
||||
if error is not None:
|
||||
logger.warning(
|
||||
"Codex 配额异步同步器 flush 失败，将重试: queued_keys={}, retry_keys={}, backoff={}s, error={}",
|
||||
queued_count,
|
||||
retry_count,
|
||||
round(self._current_flush_delay_seconds, 3),
|
||||
error,
|
||||
)
|
||||
return
|
||||
logger.warning(
|
||||
"Codex 配额异步同步器 flush 部分失败，将重试: queued_keys={}, retry_keys={}, backoff={}s",
|
||||
queued_count,
|
||||
retry_count,
|
||||
round(self._current_flush_delay_seconds, 3),
|
||||
)
|
||||
|
||||
def _flush_batch_fallback(
|
||||
self,
|
||||
entries: list[tuple[str, dict[str, Any]]],
|
||||
) -> tuple[int, dict[str, dict[str, Any]]]:
|
||||
updated_count = 0
|
||||
retry_batch: dict[str, dict[str, Any]] = {}
|
||||
for provider_api_key_id, response_headers in entries:
|
||||
db: Session = create_session()
|
||||
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:
|
||||
db.rollback()
|
||||
retry_batch[provider_api_key_id] = response_headers
|
||||
finally:
|
||||
db.close()
|
||||
return updated_count, retry_batch
|
||||
|
||||
def _flush_batch_sync(self, batch: dict[str, dict[str, Any]]) -> FlushResult:
|
||||
if not batch:
|
||||
return FlushResult(
|
||||
queued_count=0,
|
||||
updated_count=0,
|
||||
retry_batch={},
|
||||
)
|
||||
|
||||
db: Session = create_session()
|
||||
updated_entries: list[tuple[str, dict[str, Any]]] = []
|
||||
retry_batch: dict[str, dict[str, Any]] = {}
|
||||
try:
|
||||
for provider_api_key_id, response_headers in batch.items():
|
||||
try:
|
||||
with db.begin_nested():
|
||||
updated = sync_codex_quota_from_response_headers(
|
||||
db=db,
|
||||
provider_api_key_id=provider_api_key_id,
|
||||
response_headers=response_headers,
|
||||
)
|
||||
if updated:
|
||||
updated_entries.append((provider_api_key_id, response_headers))
|
||||
except Exception:
|
||||
retry_batch[provider_api_key_id] = response_headers
|
||||
|
||||
updated_count = 0
|
||||
if updated_entries:
|
||||
try:
|
||||
db.commit()
|
||||
updated_count = len(updated_entries)
|
||||
except Exception:
|
||||
db.rollback()
|
||||
fallback_updated, fallback_retry_batch = self._flush_batch_fallback(
|
||||
updated_entries
|
||||
)
|
||||
updated_count = fallback_updated
|
||||
retry_batch.update(fallback_retry_batch)
|
||||
|
||||
return FlushResult(
|
||||
queued_count=len(batch),
|
||||
updated_count=updated_count,
|
||||
retry_batch=retry_batch,
|
||||
)
|
||||
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,
|
||||
)
|
||||
@@ -0,0 +1,155 @@
|
||||
"""
|
||||
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",
|
||||
}
|
||||
)
|
||||
_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
|
||||
384
_deprecated_py_src/services/provider_keys/codex_usage_parser.py
Normal file
384
_deprecated_py_src/services/provider_keys/codex_usage_parser.py
Normal file
@@ -0,0 +1,384 @@
|
||||
"""
|
||||
Codex 配额响应解析器。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from typing import Any
|
||||
|
||||
|
||||
class CodexUsageParseError(ValueError):
|
||||
"""Codex 配额响应结构异常。"""
|
||||
|
||||
|
||||
def _raise_type_error(field: str, expected: str, value: Any) -> None:
|
||||
raise CodexUsageParseError(
|
||||
f"{field} 类型错误，期望 {expected}，实际 {type(value).__name__}: {value!r}"
|
||||
)
|
||||
|
||||
|
||||
def _as_dict(value: Any, field: str) -> dict[str, Any]:
|
||||
if value is None:
|
||||
return {}
|
||||
if not isinstance(value, dict):
|
||||
_raise_type_error(field, "object", value)
|
||||
return value
|
||||
|
||||
|
||||
def _coerce_float(value: Any, field: str) -> float:
|
||||
if isinstance(value, bool):
|
||||
_raise_type_error(field, "number", value)
|
||||
if isinstance(value, (int, float)):
|
||||
return float(value)
|
||||
if isinstance(value, str):
|
||||
raw = value.strip()
|
||||
if not raw:
|
||||
_raise_type_error(field, "number", value)
|
||||
try:
|
||||
return float(raw)
|
||||
except ValueError as exc: # pragma: no cover - 仅防御
|
||||
raise CodexUsageParseError(f"{field} 不是合法数字: {value!r}") from exc
|
||||
_raise_type_error(field, "number", value)
|
||||
|
||||
|
||||
def _coerce_int(value: Any, field: str) -> int:
|
||||
if isinstance(value, bool):
|
||||
_raise_type_error(field, "integer", value)
|
||||
if isinstance(value, int):
|
||||
return value
|
||||
if isinstance(value, float):
|
||||
if not value.is_integer():
|
||||
raise CodexUsageParseError(f"{field} 必须是整数，实际为小数: {value!r}")
|
||||
return int(value)
|
||||
if isinstance(value, str):
|
||||
raw = value.strip()
|
||||
if not raw:
|
||||
_raise_type_error(field, "integer", value)
|
||||
try:
|
||||
if "." in raw or "e" in raw.lower():
|
||||
parsed = float(raw)
|
||||
if not parsed.is_integer():
|
||||
raise CodexUsageParseError(f"{field} 必须是整数，实际为小数: {value!r}")
|
||||
return int(parsed)
|
||||
return int(raw)
|
||||
except ValueError as exc: # pragma: no cover - 仅防御
|
||||
raise CodexUsageParseError(f"{field} 不是合法整数: {value!r}") from exc
|
||||
_raise_type_error(field, "integer", value)
|
||||
|
||||
|
||||
def _coerce_bool(value: Any, field: str) -> bool:
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
if isinstance(value, int):
|
||||
if value in (0, 1):
|
||||
return bool(value)
|
||||
raise CodexUsageParseError(f"{field} 仅支持 0/1 整数，实际为: {value!r}")
|
||||
if isinstance(value, str):
|
||||
normalized = value.strip().lower()
|
||||
if normalized in {"true", "1", "yes"}:
|
||||
return True
|
||||
if normalized in {"false", "0", "no"}:
|
||||
return False
|
||||
raise CodexUsageParseError(f"{field} 不是合法布尔值: {value!r}")
|
||||
_raise_type_error(field, "boolean", value)
|
||||
|
||||
|
||||
def _write_window(
|
||||
result: dict[str, Any],
|
||||
*,
|
||||
source: dict[str, Any],
|
||||
source_field: str,
|
||||
target_prefix: str,
|
||||
) -> None:
|
||||
if not source:
|
||||
return
|
||||
|
||||
used_percent = source.get("used_percent")
|
||||
if used_percent is not None:
|
||||
result[f"{target_prefix}_used_percent"] = _coerce_float(
|
||||
used_percent, f"{source_field}.used_percent"
|
||||
)
|
||||
|
||||
reset_seconds = source.get("reset_after_seconds")
|
||||
if reset_seconds is not None:
|
||||
result[f"{target_prefix}_reset_seconds"] = _coerce_int(
|
||||
reset_seconds, f"{source_field}.reset_after_seconds"
|
||||
)
|
||||
|
||||
reset_at = source.get("reset_at")
|
||||
if reset_at is not None:
|
||||
result[f"{target_prefix}_reset_at"] = _coerce_int(reset_at, f"{source_field}.reset_at")
|
||||
|
||||
limit_window_seconds = source.get("limit_window_seconds")
|
||||
if limit_window_seconds is not None:
|
||||
result[f"{target_prefix}_window_minutes"] = (
|
||||
_coerce_int(limit_window_seconds, f"{source_field}.limit_window_seconds") // 60
|
||||
)
|
||||
|
||||
|
||||
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 响应，提取限额信息
|
||||
|
||||
Free 账号:
|
||||
- rate_limit.primary_window: 周限额
|
||||
|
||||
Team/Plus/Enterprise 账号:
|
||||
- rate_limit.primary_window: 5H 限额
|
||||
- rate_limit.secondary_window: 周限额
|
||||
"""
|
||||
if data is None:
|
||||
return None
|
||||
if not isinstance(data, dict):
|
||||
_raise_type_error("root", "object", data)
|
||||
if not data:
|
||||
return None
|
||||
|
||||
result: dict[str, Any] = {}
|
||||
|
||||
plan_type: str | None = None
|
||||
raw_plan_type = data.get("plan_type")
|
||||
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 = _normalize_plan_type(raw_plan_type)
|
||||
if normalized_plan_type:
|
||||
plan_type = normalized_plan_type
|
||||
result["plan_type"] = normalized_plan_type
|
||||
|
||||
# 解析 rate_limit
|
||||
rate_limit = _as_dict(data.get("rate_limit"), "rate_limit")
|
||||
primary_window = _as_dict(rate_limit.get("primary_window"), "rate_limit.primary_window")
|
||||
secondary_window = _as_dict(rate_limit.get("secondary_window"), "rate_limit.secondary_window")
|
||||
|
||||
# 根据账号类型解析限额
|
||||
# Free 账号: primary_window 是周限额，无 secondary_window
|
||||
# Team/Plus/Enterprise: primary_window 是 5H 限额，secondary_window 是周限额
|
||||
use_paid_windows = bool(secondary_window) and plan_type != "free"
|
||||
if use_paid_windows:
|
||||
# 周限额 (secondary_window)
|
||||
_write_window(
|
||||
result,
|
||||
source=secondary_window,
|
||||
source_field="rate_limit.secondary_window",
|
||||
target_prefix="primary",
|
||||
)
|
||||
# 5H 限额 (primary_window)
|
||||
_write_window(
|
||||
result,
|
||||
source=primary_window,
|
||||
source_field="rate_limit.primary_window",
|
||||
target_prefix="secondary",
|
||||
)
|
||||
else:
|
||||
# Free / 或 secondary_window 缺失时，primary_window 视为主窗口
|
||||
_write_window(
|
||||
result,
|
||||
source=primary_window,
|
||||
source_field="rate_limit.primary_window",
|
||||
target_prefix="primary",
|
||||
)
|
||||
|
||||
# 解析 credits
|
||||
credits = _as_dict(data.get("credits"), "credits")
|
||||
has_credits = credits.get("has_credits")
|
||||
if has_credits is not None:
|
||||
result["has_credits"] = _coerce_bool(has_credits, "credits.has_credits")
|
||||
|
||||
credits_balance = _coerce_optional_float(credits.get("balance"), "credits.balance")
|
||||
if credits_balance is not None:
|
||||
result["credits_balance"] = credits_balance
|
||||
|
||||
credits_unlimited = _coerce_optional_bool(credits.get("unlimited"), "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
|
||||
|
||||
|
||||
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",
|
||||
)
|
||||
# 与 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",
|
||||
)
|
||||
|
||||
# 当前窗口挤占占比（有值才记录）
|
||||
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
|
||||
100
_deprecated_py_src/services/provider_keys/duplicate_check.py
Normal file
100
_deprecated_py_src/services/provider_keys/duplicate_check.py
Normal file
@@ -0,0 +1,100 @@
|
||||
"""
|
||||
Provider Key 重复校验规则。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.crypto import crypto_service
|
||||
from src.core.exceptions import InvalidRequestException
|
||||
from src.models.database import ProviderAPIKey
|
||||
|
||||
|
||||
def check_duplicate_key(
|
||||
db: Session,
|
||||
provider_id: str,
|
||||
auth_type: str,
|
||||
new_api_key: str | None = None,
|
||||
new_auth_config: dict | None = None,
|
||||
exclude_key_id: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
检查密钥是否与其他现有密钥重复
|
||||
|
||||
对于不同的认证类型，使用不同的比较方式：
|
||||
- api_key: 比较 API Key 的哈希值
|
||||
- service_account: 比较 Service Account 的 client_email
|
||||
|
||||
Args:
|
||||
db: 数据库会话
|
||||
provider_id: Provider ID
|
||||
auth_type: 认证类型 (api_key, service_account, oauth)
|
||||
new_api_key: 新的 API Key（用于 api_key 类型）
|
||||
new_auth_config: 新的认证配置（用于 service_account 类型）
|
||||
exclude_key_id: 要排除的 Key ID（用于更新场景）
|
||||
"""
|
||||
if auth_type == "api_key" and new_api_key:
|
||||
# 跳过占位符
|
||||
if new_api_key == "__placeholder__":
|
||||
return
|
||||
|
||||
# 仅查询同 auth_type 的 Keys，减少不必要的解密操作
|
||||
query = db.query(ProviderAPIKey).filter(
|
||||
ProviderAPIKey.provider_id == provider_id,
|
||||
ProviderAPIKey.auth_type == "api_key",
|
||||
)
|
||||
if exclude_key_id:
|
||||
query = query.filter(ProviderAPIKey.id != exclude_key_id)
|
||||
|
||||
new_key_hash = crypto_service.hash_api_key(new_api_key)
|
||||
for existing_key in query:
|
||||
try:
|
||||
decrypted_key = crypto_service.decrypt(existing_key.api_key, silent=True)
|
||||
if decrypted_key == "__placeholder__":
|
||||
continue
|
||||
existing_hash = crypto_service.hash_api_key(decrypted_key)
|
||||
if new_key_hash == existing_hash:
|
||||
raise InvalidRequestException(
|
||||
f"该 API Key 已存在于当前 Provider 中（名称: {existing_key.name}）"
|
||||
)
|
||||
except InvalidRequestException:
|
||||
raise
|
||||
except Exception:
|
||||
# 解密失败时跳过该 Key
|
||||
continue
|
||||
|
||||
elif auth_type in ("service_account", "vertex_ai") and new_auth_config:
|
||||
new_client_email = (
|
||||
new_auth_config.get("client_email") if isinstance(new_auth_config, dict) else None
|
||||
)
|
||||
if not new_client_email:
|
||||
return
|
||||
|
||||
# 仅查询同 auth_type 且有 auth_config 的 Keys
|
||||
query = db.query(ProviderAPIKey).filter(
|
||||
ProviderAPIKey.provider_id == provider_id,
|
||||
ProviderAPIKey.auth_type.in_(["service_account", "vertex_ai"]),
|
||||
ProviderAPIKey.auth_config.isnot(None),
|
||||
)
|
||||
if exclude_key_id:
|
||||
query = query.filter(ProviderAPIKey.id != exclude_key_id)
|
||||
|
||||
for existing_key in query:
|
||||
try:
|
||||
decrypted_config = json.loads(
|
||||
crypto_service.decrypt(existing_key.auth_config, silent=True)
|
||||
)
|
||||
existing_email = decrypted_config.get("client_email")
|
||||
if existing_email and existing_email == new_client_email:
|
||||
raise InvalidRequestException(
|
||||
f"该 Service Account ({new_client_email}) 已存在于当前 Provider 中"
|
||||
f"（名称: {existing_key.name}）"
|
||||
)
|
||||
except InvalidRequestException:
|
||||
raise
|
||||
except Exception:
|
||||
# 解密失败时跳过该 Key
|
||||
continue
|
||||
596
_deprecated_py_src/services/provider_keys/key_command_service.py
Normal file
596
_deprecated_py_src/services/provider_keys/key_command_service.py
Normal file
@@ -0,0 +1,596 @@
|
||||
"""
|
||||
Provider Key 写操作命令服务。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from fastapi.concurrency import run_in_threadpool
|
||||
from sqlalchemy import delete as sa_delete
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.crypto import crypto_service
|
||||
from src.core.exceptions import InvalidRequestException, NotFoundException
|
||||
from src.core.logger import logger
|
||||
from src.core.provider_types import ProviderType
|
||||
from src.database import get_db_context
|
||||
from src.models.database import (
|
||||
Provider,
|
||||
ProviderAPIKey,
|
||||
)
|
||||
from src.models.endpoint_models import (
|
||||
EndpointAPIKeyCreate,
|
||||
EndpointAPIKeyResponse,
|
||||
EndpointAPIKeyUpdate,
|
||||
)
|
||||
from src.services.provider.fingerprint import generate_fingerprint, normalize_fingerprint
|
||||
from src.services.provider_keys.auth_type import normalize_auth_type
|
||||
from src.services.provider_keys.duplicate_check import check_duplicate_key
|
||||
from src.services.provider_keys.key_side_effects import (
|
||||
cleanup_key_references,
|
||||
run_create_key_side_effects,
|
||||
run_delete_key_side_effects,
|
||||
run_update_key_side_effects,
|
||||
)
|
||||
from src.services.provider_keys.response_builder import build_key_response
|
||||
|
||||
|
||||
def _validate_vertex_api_formats(
|
||||
provider_type: str | None,
|
||||
auth_type: str,
|
||||
api_formats: list[str] | None,
|
||||
) -> None:
|
||||
"""校验 Vertex Provider 的 key.api_formats 与 auth_type 是否匹配。"""
|
||||
if str(provider_type or "").strip().lower() != ProviderType.VERTEX_AI.value:
|
||||
return
|
||||
|
||||
formats = [
|
||||
str(fmt or "").strip().lower() for fmt in (api_formats or []) if str(fmt or "").strip()
|
||||
]
|
||||
if not formats:
|
||||
return
|
||||
|
||||
if auth_type == "api_key":
|
||||
allowed = {"gemini:chat"}
|
||||
elif auth_type in {"service_account", "vertex_ai"}:
|
||||
allowed = {"claude:chat", "gemini:chat"}
|
||||
else:
|
||||
return
|
||||
|
||||
invalid = sorted({fmt for fmt in formats if fmt not in allowed})
|
||||
if invalid:
|
||||
allowed_text = ", ".join(sorted(allowed))
|
||||
invalid_text = ", ".join(invalid)
|
||||
raise InvalidRequestException(
|
||||
f"Vertex {auth_type} 不支持以下 API 格式: {invalid_text}；允许: {allowed_text}"
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _UpdateKeyPreparation:
|
||||
"""更新 Key 前置准备结果。"""
|
||||
|
||||
update_data: dict[str, Any]
|
||||
auto_fetch_enabled_before: bool
|
||||
auto_fetch_enabled_after: bool
|
||||
allowed_models_before: set[str]
|
||||
include_patterns_before: list[str] | None
|
||||
exclude_patterns_before: list[str] | None
|
||||
|
||||
|
||||
@dataclass
|
||||
class _DeleteKeyResult:
|
||||
"""删除 Key 的执行结果。"""
|
||||
|
||||
provider_id: str | None
|
||||
deleted_key_allowed_models: list[str] | None
|
||||
|
||||
|
||||
def _update_endpoint_key_core_sync(
|
||||
key_id: str,
|
||||
key_data: EndpointAPIKeyUpdate,
|
||||
) -> _UpdateKeyPreparation:
|
||||
with get_db_context() as db:
|
||||
key = db.query(ProviderAPIKey).filter(ProviderAPIKey.id == key_id).first()
|
||||
if not key:
|
||||
raise NotFoundException(f"Key {key_id} 不存在")
|
||||
|
||||
prepared = _prepare_update_key_payload(
|
||||
db=db,
|
||||
key=key,
|
||||
key_id=key_id,
|
||||
key_data=key_data,
|
||||
)
|
||||
|
||||
for field, value in prepared.update_data.items():
|
||||
setattr(key, field, value)
|
||||
key.updated_at = datetime.now(timezone.utc)
|
||||
|
||||
return prepared
|
||||
|
||||
|
||||
def _create_provider_key_core_sync(
|
||||
provider_id: str,
|
||||
key_data: EndpointAPIKeyCreate,
|
||||
) -> str:
|
||||
with get_db_context() as db:
|
||||
provider = db.query(Provider).filter(Provider.id == provider_id).first()
|
||||
if not provider:
|
||||
raise NotFoundException(f"Provider {provider_id} 不存在")
|
||||
|
||||
if not key_data.api_formats:
|
||||
raise InvalidRequestException("api_formats 为必填字段")
|
||||
|
||||
auth_type, new_key = _prepare_create_key_payload(
|
||||
db=db,
|
||||
provider_id=provider_id,
|
||||
key_data=key_data,
|
||||
)
|
||||
|
||||
_validate_vertex_api_formats(
|
||||
getattr(provider, "provider_type", None),
|
||||
auth_type,
|
||||
key_data.api_formats,
|
||||
)
|
||||
|
||||
db.add(new_key)
|
||||
db.flush()
|
||||
return str(new_key.id)
|
||||
|
||||
|
||||
def _delete_endpoint_key_core_sync(key_id: str) -> _DeleteKeyResult:
|
||||
with get_db_context() as db:
|
||||
return _delete_endpoint_key(db, key_id)
|
||||
|
||||
|
||||
def _batch_delete_endpoint_keys_core_sync(key_ids: list[str]) -> dict[str, Any]:
|
||||
with get_db_context() as db:
|
||||
keys = db.query(ProviderAPIKey).filter(ProviderAPIKey.id.in_(key_ids)).all()
|
||||
found_ids = {key.id for key in keys}
|
||||
not_found_ids = [kid for kid in key_ids if kid not in found_ids]
|
||||
|
||||
failed: list[dict[str, str]] = [{"id": kid, "error": "not found"} for kid in not_found_ids]
|
||||
affected_provider_ids = {key.provider_id for key in keys if key.provider_id}
|
||||
|
||||
success_count = 0
|
||||
try:
|
||||
found_id_list = list(found_ids)
|
||||
cleanup_key_references(db, found_id_list)
|
||||
db.execute(sa_delete(ProviderAPIKey).where(ProviderAPIKey.id.in_(found_id_list)))
|
||||
db.commit()
|
||||
success_count = len(found_ids)
|
||||
except Exception as exc:
|
||||
db.rollback()
|
||||
logger.error("批量删除 Key 提交失败: {}", exc)
|
||||
failed.extend({"id": kid, "error": str(exc)} for kid in found_ids)
|
||||
return {
|
||||
"success_count": 0,
|
||||
"failed": failed,
|
||||
"affected_provider_ids": set(),
|
||||
}
|
||||
|
||||
return {
|
||||
"success_count": success_count,
|
||||
"failed": failed,
|
||||
"affected_provider_ids": affected_provider_ids,
|
||||
}
|
||||
|
||||
|
||||
def _run_async_with_fallback(coro: Any) -> None:
|
||||
"""在同步上下文中执行异步任务（有事件循环则调度，无则阻塞执行）。"""
|
||||
try:
|
||||
asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
asyncio.run(coro)
|
||||
return
|
||||
|
||||
from src.utils.async_utils import safe_create_task
|
||||
|
||||
safe_create_task(coro)
|
||||
|
||||
|
||||
async def _invalidate_cache_after_clear_oauth_invalid(key_id: str) -> None:
|
||||
"""清除 OAuth 失效标记后同步失效相关缓存。"""
|
||||
from src.services.cache.model_list_cache import invalidate_models_list_cache
|
||||
from src.services.cache.provider_cache import ProviderCacheService
|
||||
|
||||
await ProviderCacheService.invalidate_provider_api_key_cache(key_id)
|
||||
await invalidate_models_list_cache()
|
||||
|
||||
|
||||
def _clear_oauth_invalid_marker(db: Session, key_id: str) -> dict[str, str]:
|
||||
"""清除 Key 的 OAuth 失效标记。"""
|
||||
key = db.query(ProviderAPIKey).filter(ProviderAPIKey.id == key_id).first()
|
||||
if not key:
|
||||
raise NotFoundException(f"Key {key_id} 不存在")
|
||||
|
||||
if not key.oauth_invalid_at:
|
||||
return {"message": "该 Key 当前无失效标记，无需清除"}
|
||||
|
||||
old_reason = key.oauth_invalid_reason
|
||||
key.oauth_invalid_at = None
|
||||
key.oauth_invalid_reason = None
|
||||
db.commit()
|
||||
_run_async_with_fallback(_invalidate_cache_after_clear_oauth_invalid(key_id))
|
||||
|
||||
logger.info("[OK] 手动清除 Key {}... 的 OAuth 失效标记 (原因: {})", key_id[:8], old_reason)
|
||||
return {"message": "已清除 OAuth 失效标记"}
|
||||
|
||||
|
||||
def clear_oauth_invalid_response(db: Session, key_id: str) -> dict[str, str]:
|
||||
"""清除 OAuth 失效标记并返回统一响应。"""
|
||||
return _clear_oauth_invalid_marker(db=db, key_id=key_id)
|
||||
|
||||
|
||||
def _prepare_update_key_payload(
|
||||
db: Session,
|
||||
key: ProviderAPIKey,
|
||||
key_id: str,
|
||||
key_data: EndpointAPIKeyUpdate,
|
||||
) -> _UpdateKeyPreparation:
|
||||
"""准备更新 Key 的数据，并执行规则校验。"""
|
||||
# 检查是否开启了 auto_fetch_models（用于后续立即获取模型）
|
||||
auto_fetch_enabled_before = key.auto_fetch_models
|
||||
auto_fetch_enabled_after = (
|
||||
key_data.auto_fetch_models
|
||||
if "auto_fetch_models" in key_data.model_fields_set
|
||||
else auto_fetch_enabled_before
|
||||
)
|
||||
|
||||
# 记录 allowed_models 变化前的值
|
||||
allowed_models_before = set(key.allowed_models or [])
|
||||
|
||||
# 记录过滤规则变化前的值（用于检测是否需要重新应用过滤）
|
||||
include_patterns_before = key.model_include_patterns
|
||||
exclude_patterns_before = key.model_exclude_patterns
|
||||
|
||||
update_data = key_data.model_dump(exclude_unset=True)
|
||||
# 显式传 null 等价于“不更新 auth_type”，避免写入 NULL 触发数据库约束错误。
|
||||
if update_data.get("auth_type") is None:
|
||||
update_data.pop("auth_type", None)
|
||||
if "api_key" in update_data and isinstance(update_data["api_key"], str):
|
||||
update_data["api_key"] = update_data["api_key"].strip()
|
||||
|
||||
# 验证 auth_type
|
||||
current_auth_type = normalize_auth_type(getattr(key, "auth_type", "api_key"))
|
||||
target_auth_type = normalize_auth_type(update_data.get("auth_type", current_auth_type))
|
||||
is_auth_type_switch = "auth_type" in update_data and target_auth_type != current_auth_type
|
||||
api_key_in_payload = "api_key" in update_data
|
||||
api_key_value = update_data.get("api_key")
|
||||
|
||||
if api_key_in_payload and api_key_value == "":
|
||||
raise InvalidRequestException("api_key 不能为空")
|
||||
|
||||
# auth_type 校验 + 字段归一化
|
||||
if target_auth_type == "api_key":
|
||||
if is_auth_type_switch and (not api_key_value or api_key_value == "__placeholder__"):
|
||||
raise InvalidRequestException("切换到 API Key 认证模式时，必须提供新的 API Key")
|
||||
if api_key_in_payload and (api_key_value is None or api_key_value == "__placeholder__"):
|
||||
raise InvalidRequestException("API Key 认证模式下 api_key 不能为空")
|
||||
# 切换回 API Key：清理非本模式配置
|
||||
update_data["auth_config"] = None
|
||||
elif target_auth_type == "service_account":
|
||||
if is_auth_type_switch and not update_data.get("auth_config"):
|
||||
raise InvalidRequestException(
|
||||
"切换到 Service Account 认证模式时，必须提供 Service Account JSON"
|
||||
)
|
||||
# Service Account 不允许手工写入 api_key，仅保留占位符
|
||||
if api_key_in_payload and api_key_value not in {None, "__placeholder__"}:
|
||||
raise InvalidRequestException("Service Account 认证模式下不允许直接填写 api_key")
|
||||
if is_auth_type_switch or api_key_in_payload:
|
||||
update_data["api_key"] = "__placeholder__"
|
||||
elif target_auth_type == "oauth":
|
||||
# OAuth 的 token 不允许在 key 更新接口里手工写入
|
||||
if api_key_in_payload and api_key_value not in {None, "__placeholder__"}:
|
||||
raise InvalidRequestException("OAuth 认证模式下不允许直接填写 api_key")
|
||||
if is_auth_type_switch:
|
||||
update_data["api_key"] = "__placeholder__"
|
||||
# 从非 OAuth 切换到 OAuth 时，清理旧认证配置（如 Vertex SA 凭证）。
|
||||
update_data["auth_config"] = None
|
||||
elif api_key_in_payload:
|
||||
# 避免把 null 写入 DB 或意外覆盖现有 OAuth token。
|
||||
update_data.pop("api_key", None)
|
||||
|
||||
# 检查密钥是否与其他现有密钥重复（排除当前正在更新的密钥）
|
||||
check_duplicate_key(
|
||||
db=db,
|
||||
provider_id=key.provider_id,
|
||||
auth_type=target_auth_type,
|
||||
new_api_key=update_data.get("api_key"),
|
||||
new_auth_config=update_data.get("auth_config"),
|
||||
exclude_key_id=key_id,
|
||||
)
|
||||
|
||||
# Vertex Provider: 仅在 auth_type/api_formats 变更时校验组合，
|
||||
# 避免历史旧数据在无关编辑时被强制阻断。
|
||||
if "auth_type" in update_data or "api_formats" in update_data:
|
||||
provider = getattr(key, "provider", None)
|
||||
effective_api_formats = update_data.get("api_formats", key.api_formats)
|
||||
_validate_vertex_api_formats(
|
||||
getattr(provider, "provider_type", None),
|
||||
target_auth_type,
|
||||
effective_api_formats,
|
||||
)
|
||||
|
||||
if "api_key" in update_data:
|
||||
api_key_raw = update_data["api_key"]
|
||||
if api_key_raw is None:
|
||||
# 防御式处理：避免将 NULL 写入 NOT NULL 字段导致 500。
|
||||
update_data.pop("api_key", None)
|
||||
else:
|
||||
update_data["api_key"] = crypto_service.encrypt(api_key_raw)
|
||||
|
||||
# 加密 auth_config（包含敏感凭证）；即便是 {} 也必须加密存储。
|
||||
if "auth_config" in update_data:
|
||||
auth_config_raw = update_data["auth_config"]
|
||||
if auth_config_raw is None:
|
||||
pass
|
||||
elif isinstance(auth_config_raw, dict):
|
||||
update_data["auth_config"] = crypto_service.encrypt(json.dumps(auth_config_raw))
|
||||
else:
|
||||
raise InvalidRequestException("auth_config 必须是 JSON 对象")
|
||||
|
||||
# 特殊处理 rpm_limit：需要区分"未提供"和"显式设置为 null"
|
||||
if "rpm_limit" in key_data.model_fields_set:
|
||||
update_data["rpm_limit"] = key_data.rpm_limit
|
||||
if key_data.rpm_limit is None:
|
||||
update_data["learned_rpm_limit"] = None
|
||||
logger.info("Key {} 切换为自适应 RPM 模式", key_id)
|
||||
|
||||
# 统一处理 allowed_models：空列表 -> None（表示不限制）
|
||||
if "allowed_models" in update_data:
|
||||
am = update_data["allowed_models"]
|
||||
if isinstance(am, list) and len(am) == 0:
|
||||
update_data["allowed_models"] = None
|
||||
|
||||
# 统一处理 locked_models：空列表 -> None
|
||||
if "locked_models" in update_data:
|
||||
lm = update_data["locked_models"]
|
||||
if isinstance(lm, list) and len(lm) == 0:
|
||||
update_data["locked_models"] = None
|
||||
|
||||
# 处理模型过滤规则：空字符串 -> None
|
||||
if "model_include_patterns" in update_data:
|
||||
patterns = update_data["model_include_patterns"]
|
||||
if isinstance(patterns, list) and len(patterns) == 0:
|
||||
update_data["model_include_patterns"] = None
|
||||
|
||||
if "model_exclude_patterns" in update_data:
|
||||
patterns = update_data["model_exclude_patterns"]
|
||||
if isinstance(patterns, list) and len(patterns) == 0:
|
||||
update_data["model_exclude_patterns"] = None
|
||||
|
||||
# 处理 proxy：将 ProxyConfig 转换为 dict 存储，null 清除代理
|
||||
if "proxy" in key_data.model_fields_set:
|
||||
if key_data.proxy is None:
|
||||
update_data["proxy"] = None
|
||||
else:
|
||||
update_data["proxy"] = key_data.proxy.model_dump(exclude_none=True)
|
||||
|
||||
if "fingerprint" in key_data.model_fields_set:
|
||||
if key_data.fingerprint is None:
|
||||
update_data["fingerprint"] = None
|
||||
else:
|
||||
update_data["fingerprint"] = normalize_fingerprint(key_data.fingerprint, key_id)
|
||||
|
||||
return _UpdateKeyPreparation(
|
||||
update_data=update_data,
|
||||
auto_fetch_enabled_before=auto_fetch_enabled_before,
|
||||
auto_fetch_enabled_after=auto_fetch_enabled_after,
|
||||
allowed_models_before=allowed_models_before,
|
||||
include_patterns_before=include_patterns_before,
|
||||
exclude_patterns_before=exclude_patterns_before,
|
||||
)
|
||||
|
||||
|
||||
def _prepare_create_key_payload(
|
||||
db: Session,
|
||||
provider_id: str,
|
||||
key_data: EndpointAPIKeyCreate,
|
||||
) -> tuple[str, ProviderAPIKey]:
|
||||
"""准备创建 Key 的认证类型校验、重复校验与实体构造。"""
|
||||
auth_type = key_data.auth_type or "api_key"
|
||||
if auth_type == "api_key":
|
||||
if not key_data.api_key:
|
||||
raise InvalidRequestException("API Key 认证模式下 api_key 为必填字段")
|
||||
elif auth_type == "service_account":
|
||||
if not key_data.auth_config:
|
||||
raise InvalidRequestException("Service Account 认证模式下 auth_config 为必填字段")
|
||||
elif auth_type == "oauth":
|
||||
# OAuth key 的 token 通过 provider-oauth 授权流程写入（此处不允许手填）
|
||||
if key_data.api_key:
|
||||
raise InvalidRequestException("OAuth 认证模式下不允许直接填写 api_key")
|
||||
|
||||
# 检查密钥是否已存在（防止重复添加）
|
||||
check_duplicate_key(
|
||||
db=db,
|
||||
provider_id=provider_id,
|
||||
auth_type=auth_type,
|
||||
new_api_key=key_data.api_key,
|
||||
new_auth_config=key_data.auth_config,
|
||||
)
|
||||
|
||||
# 加密 API Key（如果有）
|
||||
encrypted_key = (
|
||||
crypto_service.encrypt(key_data.api_key)
|
||||
if key_data.api_key
|
||||
else crypto_service.encrypt("__placeholder__") # 占位符，保持 NOT NULL 约束
|
||||
)
|
||||
# OAuth 类型 key 初始写入占位符（token 由 provider-oauth 流程写入）
|
||||
if auth_type == "oauth":
|
||||
encrypted_key = crypto_service.encrypt("__placeholder__")
|
||||
now = datetime.now(timezone.utc)
|
||||
|
||||
# 加密 auth_config（包含敏感的 Service Account 凭证）
|
||||
encrypted_auth_config = None
|
||||
if key_data.auth_config:
|
||||
encrypted_auth_config = crypto_service.encrypt(json.dumps(key_data.auth_config))
|
||||
|
||||
new_key_id = str(uuid.uuid4())
|
||||
|
||||
new_key = ProviderAPIKey(
|
||||
id=new_key_id,
|
||||
provider_id=provider_id,
|
||||
api_formats=key_data.api_formats,
|
||||
auth_type=auth_type,
|
||||
api_key=encrypted_key,
|
||||
auth_config=encrypted_auth_config,
|
||||
name=key_data.name,
|
||||
note=key_data.note,
|
||||
rate_multipliers=key_data.rate_multipliers,
|
||||
internal_priority=key_data.internal_priority,
|
||||
rpm_limit=key_data.rpm_limit,
|
||||
allowed_models=key_data.allowed_models if key_data.allowed_models else None,
|
||||
capabilities=key_data.capabilities if key_data.capabilities else None,
|
||||
cache_ttl_minutes=key_data.cache_ttl_minutes,
|
||||
max_probe_interval_minutes=key_data.max_probe_interval_minutes,
|
||||
auto_fetch_models=key_data.auto_fetch_models,
|
||||
locked_models=key_data.locked_models if key_data.locked_models else None,
|
||||
model_include_patterns=(
|
||||
key_data.model_include_patterns if key_data.model_include_patterns else None
|
||||
),
|
||||
model_exclude_patterns=(
|
||||
key_data.model_exclude_patterns if key_data.model_exclude_patterns else None
|
||||
),
|
||||
fingerprint=generate_fingerprint(seed=new_key_id),
|
||||
request_count=0,
|
||||
success_count=0,
|
||||
error_count=0,
|
||||
total_response_time_ms=0,
|
||||
health_by_format={}, # 按格式存储健康度
|
||||
circuit_breaker_by_format={}, # 按格式存储熔断器状态
|
||||
is_active=True,
|
||||
last_used_at=None,
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
)
|
||||
return auth_type, new_key
|
||||
|
||||
|
||||
async def update_endpoint_key_response(
|
||||
db: Session,
|
||||
key_id: str,
|
||||
key_data: EndpointAPIKeyUpdate,
|
||||
) -> EndpointAPIKeyResponse:
|
||||
"""更新 Key 并返回响应对象。"""
|
||||
prepared = await run_in_threadpool(_update_endpoint_key_core_sync, key_id, key_data)
|
||||
|
||||
db.expire_all()
|
||||
key = db.query(ProviderAPIKey).filter(ProviderAPIKey.id == key_id).first()
|
||||
if not key:
|
||||
raise NotFoundException(f"Key {key_id} 不存在")
|
||||
|
||||
await run_update_key_side_effects(
|
||||
db=db,
|
||||
key=key,
|
||||
key_id=key_id,
|
||||
auto_fetch_enabled_before=prepared.auto_fetch_enabled_before,
|
||||
auto_fetch_enabled_after=prepared.auto_fetch_enabled_after,
|
||||
include_patterns_before=prepared.include_patterns_before,
|
||||
exclude_patterns_before=prepared.exclude_patterns_before,
|
||||
allowed_models_before=prepared.allowed_models_before,
|
||||
)
|
||||
|
||||
logger.info("[OK] 更新 Key: ID={}, Updates={}", key_id, list(prepared.update_data.keys()))
|
||||
return build_key_response(key)
|
||||
|
||||
|
||||
async def create_provider_key_response(
|
||||
db: Session,
|
||||
provider_id: str,
|
||||
key_data: EndpointAPIKeyCreate,
|
||||
) -> EndpointAPIKeyResponse:
|
||||
"""创建 Provider Key 并返回响应对象。"""
|
||||
key_id = await run_in_threadpool(_create_provider_key_core_sync, provider_id, key_data)
|
||||
|
||||
db.expire_all()
|
||||
new_key = db.query(ProviderAPIKey).filter(ProviderAPIKey.id == key_id).first()
|
||||
if not new_key:
|
||||
raise NotFoundException(f"Key {key_id} 不存在")
|
||||
|
||||
key_tail = (key_data.api_key or "")[-4:]
|
||||
logger.info(
|
||||
"[OK] 添加 Key: Provider={}, Formats={}, Key=***{}, ID={}",
|
||||
provider_id,
|
||||
key_data.api_formats,
|
||||
key_tail,
|
||||
new_key.id,
|
||||
)
|
||||
|
||||
await run_create_key_side_effects(db=db, provider_id=provider_id, key=new_key)
|
||||
return build_key_response(new_key, api_key_plain=key_data.api_key)
|
||||
|
||||
|
||||
def _delete_endpoint_key(db: Session, key_id: str) -> _DeleteKeyResult:
|
||||
"""删除指定 Key 并返回删除上下文。"""
|
||||
key = db.query(ProviderAPIKey).filter(ProviderAPIKey.id == key_id).first()
|
||||
if not key:
|
||||
raise NotFoundException(f"Key {key_id} 不存在")
|
||||
|
||||
provider_id = key.provider_id
|
||||
deleted_key_allowed_models = key.allowed_models # 保存被删除 Key 的 allowed_models
|
||||
try:
|
||||
cleanup_key_references(db, [key_id])
|
||||
db.delete(key)
|
||||
db.commit()
|
||||
except Exception as exc:
|
||||
db.rollback()
|
||||
logger.error(f"删除 Key 失败: ID={key_id}, Error={exc}")
|
||||
raise
|
||||
|
||||
return _DeleteKeyResult(
|
||||
provider_id=provider_id,
|
||||
deleted_key_allowed_models=deleted_key_allowed_models,
|
||||
)
|
||||
|
||||
|
||||
async def delete_endpoint_key_response(db: Session, key_id: str) -> dict[str, str]:
|
||||
"""删除 Key，执行副作用并返回统一响应。"""
|
||||
delete_result = await run_in_threadpool(_delete_endpoint_key_core_sync, key_id)
|
||||
|
||||
await run_delete_key_side_effects(
|
||||
db=db,
|
||||
provider_id=delete_result.provider_id,
|
||||
deleted_key_allowed_models=delete_result.deleted_key_allowed_models,
|
||||
)
|
||||
logger.warning("[DELETE] 删除 Key: ID={}, Provider={}", key_id, delete_result.provider_id)
|
||||
return {"message": f"Key {key_id} 已删除"}
|
||||
|
||||
|
||||
async def batch_delete_endpoint_keys_response(db: Session, key_ids: list[str]) -> dict[str, Any]:
|
||||
"""批量删除 Keys，按 provider_id 聚合后仅执行一次副作用。"""
|
||||
if not key_ids:
|
||||
return {"success_count": 0, "failed_count": 0, "failed": []}
|
||||
|
||||
result = await run_in_threadpool(_batch_delete_endpoint_keys_core_sync, key_ids)
|
||||
affected_provider_ids = result["affected_provider_ids"]
|
||||
failed = result["failed"]
|
||||
success_count = result["success_count"]
|
||||
|
||||
for provider_id in affected_provider_ids:
|
||||
try:
|
||||
await run_delete_key_side_effects(
|
||||
db=db,
|
||||
provider_id=provider_id,
|
||||
deleted_key_allowed_models=None,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error("批量删除副作用执行失败: provider_id={}, Error={}", provider_id, exc)
|
||||
|
||||
logger.warning(
|
||||
"[BATCH_DELETE] 批量删除 Keys: success={}, failed={}, providers={}",
|
||||
success_count,
|
||||
len(failed),
|
||||
len(affected_provider_ids),
|
||||
)
|
||||
return {
|
||||
"success_count": success_count,
|
||||
"failed_count": len(failed),
|
||||
"failed": failed,
|
||||
}
|
||||
325
_deprecated_py_src/services/provider_keys/key_query_service.py
Normal file
325
_deprecated_py_src/services/provider_keys/key_query_service.py
Normal file
@@ -0,0 +1,325 @@
|
||||
"""
|
||||
Provider Key 查询服务。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.api_format.signature import normalize_signature_key
|
||||
from src.core.crypto import crypto_service
|
||||
from src.core.exceptions import InvalidRequestException, NotFoundException
|
||||
from src.core.key_capabilities import get_capability
|
||||
from src.core.logger import logger
|
||||
from src.models.database import Provider, ProviderAPIKey, ProviderEndpoint
|
||||
from src.models.endpoint_models import EndpointAPIKeyResponse
|
||||
from src.services.provider_keys.auth_type import normalize_auth_type
|
||||
from src.services.provider_keys.response_builder import build_key_response
|
||||
|
||||
_LEGACY_API_FORMAT_MAP: dict[str, str] = {
|
||||
"CLAUDE": "claude:chat",
|
||||
"CLAUDE_CLI": "claude:cli",
|
||||
"OPENAI": "openai:chat",
|
||||
"OPENAI_CLI": "openai:cli",
|
||||
"OPENAI_COMPACT": "openai:compact",
|
||||
"OPENAI_VIDEO": "openai:video",
|
||||
"GEMINI": "gemini:chat",
|
||||
"GEMINI_CLI": "gemini:cli",
|
||||
"GEMINI_VIDEO": "gemini:video",
|
||||
}
|
||||
|
||||
|
||||
def _normalize_api_format_key(raw_format: Any) -> str | None:
|
||||
"""Normalize api_format to canonical signature key; keeps legacy import compatibility."""
|
||||
text = str(raw_format or "").strip()
|
||||
if not text:
|
||||
return None
|
||||
|
||||
try:
|
||||
return normalize_signature_key(text)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
legacy = text.upper().replace("-", "_")
|
||||
return _LEGACY_API_FORMAT_MAP.get(legacy)
|
||||
|
||||
|
||||
def _normalize_format_dict(raw_dict: Any) -> dict[str, Any]:
|
||||
"""Normalize dict keys from any format aliases to canonical api_format."""
|
||||
if not isinstance(raw_dict, dict):
|
||||
return {}
|
||||
|
||||
normalized: dict[str, Any] = {}
|
||||
for raw_key, value in raw_dict.items():
|
||||
format_key = _normalize_api_format_key(raw_key)
|
||||
if not format_key or format_key in normalized:
|
||||
continue
|
||||
normalized[format_key] = value
|
||||
return normalized
|
||||
|
||||
|
||||
def get_keys_grouped_by_format(db: Session) -> dict:
|
||||
"""查询所有 Key，并按 API 格式分组返回。"""
|
||||
# Key 属于 Provider：按 key.api_formats 分组展示
|
||||
# 包含所有 Key（含停用的 Key 和停用的 Provider），前端可显示停用标签和快捷开关
|
||||
keys = (
|
||||
db.query(ProviderAPIKey, Provider)
|
||||
.join(Provider, ProviderAPIKey.provider_id == Provider.id)
|
||||
.order_by(
|
||||
ProviderAPIKey.internal_priority.asc(),
|
||||
)
|
||||
.all()
|
||||
)
|
||||
|
||||
provider_ids = {str(provider.id) for _key, provider in keys}
|
||||
endpoints = (
|
||||
db.query(
|
||||
ProviderEndpoint.provider_id,
|
||||
ProviderEndpoint.api_format,
|
||||
ProviderEndpoint.base_url,
|
||||
)
|
||||
.filter(
|
||||
ProviderEndpoint.provider_id.in_(provider_ids),
|
||||
ProviderEndpoint.is_active.is_(True),
|
||||
)
|
||||
.all()
|
||||
)
|
||||
endpoint_base_url_map: dict[tuple[str, str], str] = {}
|
||||
for provider_id, api_format, base_url in endpoints:
|
||||
fmt = api_format.value if hasattr(api_format, "value") else str(api_format)
|
||||
normalized_fmt = _normalize_api_format_key(fmt)
|
||||
if not normalized_fmt:
|
||||
continue
|
||||
endpoint_base_url_map[(str(provider_id), normalized_fmt)] = base_url
|
||||
|
||||
grouped: dict[str, list[dict]] = {}
|
||||
for key, provider in keys:
|
||||
raw_api_formats = key.api_formats or []
|
||||
api_formats: list[str] = []
|
||||
seen_formats: set[str] = set()
|
||||
for raw_format in raw_api_formats:
|
||||
normalized_format = _normalize_api_format_key(raw_format)
|
||||
if not normalized_format or normalized_format in seen_formats:
|
||||
continue
|
||||
seen_formats.add(normalized_format)
|
||||
api_formats.append(normalized_format)
|
||||
|
||||
if not api_formats:
|
||||
continue # 跳过没有 API 格式的 Key
|
||||
|
||||
auth_type = normalize_auth_type(getattr(key, "auth_type", "api_key"))
|
||||
if auth_type in ("service_account", "vertex_ai"):
|
||||
masked_key = "[Service Account]"
|
||||
elif auth_type == "oauth":
|
||||
masked_key = "[OAuth Token]"
|
||||
else:
|
||||
try:
|
||||
decrypted_key = crypto_service.decrypt(key.api_key)
|
||||
masked_key = f"{decrypted_key[:8]}***{decrypted_key[-4:]}"
|
||||
except Exception as e:
|
||||
logger.error(f"解密 Key 失败: key_id={key.id}, error={e}")
|
||||
masked_key = "***ERROR***"
|
||||
|
||||
# 计算健康度指标
|
||||
success_rate = key.success_count / key.request_count if key.request_count > 0 else None
|
||||
avg_response_time_ms = (
|
||||
round(key.total_response_time_ms / key.success_count, 2)
|
||||
if key.success_count > 0
|
||||
else None
|
||||
)
|
||||
|
||||
# 将 capabilities dict 转换为启用的能力简短名称列表
|
||||
caps_list = []
|
||||
if key.capabilities:
|
||||
for cap_name, enabled in key.capabilities.items():
|
||||
if enabled:
|
||||
cap_def = get_capability(cap_name)
|
||||
caps_list.append(cap_def.short_name if cap_def else cap_name)
|
||||
|
||||
# 构建 Key 信息（基础数据）
|
||||
normalized_rate_multipliers = _normalize_format_dict(key.rate_multipliers)
|
||||
normalized_priority_by_format = _normalize_format_dict(key.global_priority_by_format)
|
||||
key_info = {
|
||||
"id": key.id,
|
||||
"provider_id": str(provider.id),
|
||||
"name": key.name,
|
||||
"auth_type": auth_type,
|
||||
"api_key_masked": masked_key,
|
||||
"internal_priority": key.internal_priority,
|
||||
"global_priority_by_format": normalized_priority_by_format,
|
||||
"rate_multipliers": normalized_rate_multipliers or None,
|
||||
"is_active": key.is_active,
|
||||
"provider_active": provider.is_active,
|
||||
"provider_name": provider.name,
|
||||
"api_formats": api_formats,
|
||||
"capabilities": caps_list,
|
||||
"success_rate": success_rate,
|
||||
"avg_response_time_ms": avg_response_time_ms,
|
||||
"request_count": key.request_count,
|
||||
}
|
||||
|
||||
# 将 Key 添加到每个支持的格式分组中，并附加格式特定的数据
|
||||
health_by_format = _normalize_format_dict(key.health_by_format)
|
||||
circuit_by_format = _normalize_format_dict(key.circuit_breaker_by_format)
|
||||
priority_by_format: dict[str, int] = {}
|
||||
for k, v in normalized_priority_by_format.items():
|
||||
try:
|
||||
priority_by_format[k] = int(v)
|
||||
except Exception:
|
||||
continue
|
||||
provider_id = str(provider.id)
|
||||
for api_format in api_formats:
|
||||
if api_format not in grouped:
|
||||
grouped[api_format] = []
|
||||
# 为每个格式创建副本，设置当前格式
|
||||
format_key_info = key_info.copy()
|
||||
format_key_info["api_format"] = api_format
|
||||
format_key_info["endpoint_base_url"] = endpoint_base_url_map.get(
|
||||
(provider_id, api_format)
|
||||
)
|
||||
# 添加格式特定的优先级
|
||||
format_key_info["format_priority"] = priority_by_format.get(api_format)
|
||||
# 添加格式特定的健康度数据
|
||||
format_health = health_by_format.get(api_format, {})
|
||||
format_circuit = circuit_by_format.get(api_format, {})
|
||||
format_key_info["health_score"] = float(format_health.get("health_score") or 1.0)
|
||||
format_key_info["circuit_breaker_open"] = bool(format_circuit.get("open", False))
|
||||
grouped[api_format].append(format_key_info)
|
||||
|
||||
# 直接返回分组对象，供前端使用
|
||||
return grouped
|
||||
|
||||
|
||||
def list_provider_keys_responses(
|
||||
db: Session,
|
||||
provider_id: str,
|
||||
skip: int,
|
||||
limit: int,
|
||||
) -> list[EndpointAPIKeyResponse]:
|
||||
"""查询 Provider 下的 Key 列表并构建响应。"""
|
||||
provider = db.query(Provider).filter(Provider.id == provider_id).first()
|
||||
if not provider:
|
||||
raise NotFoundException(f"Provider {provider_id} 不存在")
|
||||
provider_type = (
|
||||
str(
|
||||
getattr(provider, "provider_type", None) or getattr(provider, "type", None) or ""
|
||||
).strip()
|
||||
or None
|
||||
)
|
||||
|
||||
keys = (
|
||||
db.query(ProviderAPIKey)
|
||||
.filter(ProviderAPIKey.provider_id == provider_id)
|
||||
.order_by(ProviderAPIKey.internal_priority.asc(), ProviderAPIKey.created_at.asc())
|
||||
.offset(skip)
|
||||
.limit(limit)
|
||||
.all()
|
||||
)
|
||||
return [build_key_response(key, provider_type=provider_type) for key in keys]
|
||||
|
||||
|
||||
def reveal_endpoint_key_payload(
|
||||
db: Session,
|
||||
key_id: str,
|
||||
) -> dict[str, Any]:
|
||||
"""获取完整的 API Key 或 Auth Config（用于查看和复制）。"""
|
||||
key = db.query(ProviderAPIKey).filter(ProviderAPIKey.id == key_id).first()
|
||||
if not key:
|
||||
raise NotFoundException(f"Key {key_id} 不存在")
|
||||
|
||||
auth_type = normalize_auth_type(getattr(key, "auth_type", "api_key"))
|
||||
|
||||
# Service Account 类型返回 auth_config（需要解密）
|
||||
if auth_type in ("service_account", "vertex_ai"):
|
||||
encrypted_auth_config = getattr(key, "auth_config", None)
|
||||
if encrypted_auth_config:
|
||||
try:
|
||||
decrypted_config = crypto_service.decrypt(encrypted_auth_config)
|
||||
auth_config = json.loads(decrypted_config)
|
||||
logger.info(f"[REVEAL] 查看 Auth Config: ID={key_id}, Name={key.name}")
|
||||
return {"auth_type": auth_type, "auth_config": auth_config}
|
||||
except Exception as e:
|
||||
logger.error(f"解密 Auth Config 失败: ID={key_id}, Error={e}")
|
||||
raise InvalidRequestException(
|
||||
"无法解密认证配置，可能是加密密钥已更改。请重新添加该密钥。"
|
||||
)
|
||||
|
||||
# 兼容：auth_config 为空时尝试从 api_key 解密（仅对迁移前的旧数据有效）
|
||||
try:
|
||||
decrypted_key = crypto_service.decrypt(key.api_key)
|
||||
if decrypted_key == "__placeholder__":
|
||||
logger.error(f"Service Account Key 缺少 auth_config: ID={key_id}")
|
||||
raise InvalidRequestException("认证配置丢失，请重新添加该密钥。")
|
||||
logger.info(f"[REVEAL] 查看完整 Key (legacy SA): ID={key_id}, Name={key.name}")
|
||||
return {"auth_type": auth_type, "auth_config": decrypted_key}
|
||||
except InvalidRequestException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"解密 Key 失败: ID={key_id}, Error={e}")
|
||||
raise InvalidRequestException(
|
||||
"无法解密认证配置，可能是加密密钥已更改。请重新添加该密钥。"
|
||||
)
|
||||
|
||||
# OAuth 类型：返回 access_token（导出走 /export 端点）
|
||||
if auth_type == "oauth":
|
||||
try:
|
||||
decrypted_key = crypto_service.decrypt(key.api_key)
|
||||
except Exception as e:
|
||||
logger.error(f"解密 Key 失败: ID={key_id}, Error={e}")
|
||||
raise InvalidRequestException(
|
||||
"无法解密 API Key，可能是加密密钥已更改。请重新添加该密钥。"
|
||||
)
|
||||
logger.info(f"[REVEAL] 查看 OAuth Key: ID={key_id}, Name={key.name}")
|
||||
return {"auth_type": "oauth", "api_key": decrypted_key}
|
||||
|
||||
# API Key 类型返回 api_key
|
||||
try:
|
||||
decrypted_key = crypto_service.decrypt(key.api_key)
|
||||
except Exception as e:
|
||||
logger.error(f"解密 Key 失败: ID={key_id}, Error={e}")
|
||||
raise InvalidRequestException("无法解密 API Key，可能是加密密钥已更改。请重新添加该密钥。")
|
||||
|
||||
logger.info(f"[REVEAL] 查看完整 Key: ID={key_id}, Name={key.name}")
|
||||
return {"auth_type": "api_key", "api_key": decrypted_key}
|
||||
|
||||
|
||||
def export_oauth_key_data(
|
||||
db: Session,
|
||||
key_id: str,
|
||||
) -> dict[str, Any]:
|
||||
"""导出 OAuth Key 凭据。"""
|
||||
from src.services.provider.export import build_export_data
|
||||
|
||||
key = db.query(ProviderAPIKey).filter(ProviderAPIKey.id == key_id).first()
|
||||
if not key:
|
||||
raise NotFoundException(f"Key {key_id} 不存在")
|
||||
|
||||
auth_type = normalize_auth_type(getattr(key, "auth_type", "api_key"))
|
||||
if auth_type != "oauth":
|
||||
raise InvalidRequestException("仅 OAuth 类型的 Key 支持导出")
|
||||
|
||||
encrypted_auth_config = getattr(key, "auth_config", None)
|
||||
if not encrypted_auth_config:
|
||||
raise InvalidRequestException("缺少认证配置，无法导出")
|
||||
|
||||
try:
|
||||
auth_config: dict[str, Any] = json.loads(crypto_service.decrypt(encrypted_auth_config))
|
||||
except Exception:
|
||||
raise InvalidRequestException("无法解密认证配置")
|
||||
|
||||
if not auth_config.get("refresh_token"):
|
||||
raise InvalidRequestException("缺少 refresh_token，无法导出")
|
||||
|
||||
provider_type = str(auth_config.get("provider_type") or "").strip()
|
||||
upstream = getattr(key, "upstream_metadata", None)
|
||||
|
||||
export_data = build_export_data(provider_type, auth_config, upstream)
|
||||
export_data["name"] = key.name or ""
|
||||
export_data["exported_at"] = datetime.now(timezone.utc).isoformat()
|
||||
|
||||
logger.info("[EXPORT] Key {}... 导出成功", key_id[:8])
|
||||
return export_data
|
||||
301
_deprecated_py_src/services/provider_keys/key_quota_service.py
Normal file
301
_deprecated_py_src/services/provider_keys/key_quota_service.py
Normal file
@@ -0,0 +1,301 @@
|
||||
"""Provider Key 配额刷新编排服务。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import or_ as db_or
|
||||
from sqlalchemy.orm import Session, defer
|
||||
|
||||
from src.core.exceptions import InvalidRequestException, NotFoundException
|
||||
from src.core.logger import logger
|
||||
from src.core.provider_types import ProviderType, normalize_provider_type
|
||||
from src.models.database import Provider, ProviderAPIKey, ProviderEndpoint
|
||||
from src.services.model.upstream_fetcher import merge_upstream_metadata
|
||||
from src.services.provider.pool import redis_ops as pool_redis
|
||||
from src.services.provider.pool.account_state import (
|
||||
resolve_pool_account_state,
|
||||
should_auto_remove_account_state,
|
||||
)
|
||||
from src.services.provider.pool.config import parse_pool_config
|
||||
from src.services.provider_keys.key_side_effects import run_delete_key_side_effects
|
||||
from src.services.provider_keys.quota_refresh import (
|
||||
refresh_antigravity_key_quota,
|
||||
refresh_codex_key_quota,
|
||||
refresh_kiro_key_quota,
|
||||
)
|
||||
|
||||
QuotaRefreshHandler = Callable[..., Awaitable[dict]]
|
||||
|
||||
|
||||
CODEX_WHAM_USAGE_URL = "https://chatgpt.com/backend-api/wham/usage"
|
||||
|
||||
_QUOTA_REFRESH_HANDLERS: dict[str, QuotaRefreshHandler] = {
|
||||
ProviderType.CODEX: refresh_codex_key_quota,
|
||||
ProviderType.ANTIGRAVITY: refresh_antigravity_key_quota,
|
||||
ProviderType.KIRO: refresh_kiro_key_quota,
|
||||
}
|
||||
|
||||
QUOTA_REFRESH_PROVIDER_TYPES: frozenset[str] = frozenset(_QUOTA_REFRESH_HANDLERS.keys())
|
||||
|
||||
|
||||
def _normalize_api_format(api_format: Any) -> str:
|
||||
"""规范化 api_format，兼容大小写和首尾空白。"""
|
||||
if not isinstance(api_format, str):
|
||||
return ""
|
||||
return api_format.strip().lower()
|
||||
|
||||
|
||||
def _select_refresh_endpoint(provider: Provider, provider_type: str) -> ProviderEndpoint | None:
|
||||
"""为配额刷新选择端点。"""
|
||||
if provider_type == ProviderType.CODEX:
|
||||
for ep in provider.endpoints:
|
||||
if _normalize_api_format(ep.api_format) == "openai:cli" and ep.is_active:
|
||||
return ep
|
||||
raise InvalidRequestException("找不到有效的 openai:cli 端点")
|
||||
|
||||
if provider_type == ProviderType.ANTIGRAVITY:
|
||||
# Prefer the new signature, but keep backward-compat with existing DB rows.
|
||||
for sig in ("gemini:chat", "gemini:cli"):
|
||||
for ep in provider.endpoints:
|
||||
if _normalize_api_format(ep.api_format) == sig and ep.is_active:
|
||||
return ep
|
||||
raise InvalidRequestException("找不到有效的 gemini:chat/gemini:cli 端点")
|
||||
|
||||
# Kiro 不需要端点检查，直接使用 auth_config
|
||||
return None
|
||||
|
||||
|
||||
def _resolve_quota_refresh_handler(provider_type: str) -> QuotaRefreshHandler:
|
||||
"""按 provider 类型返回刷新策略。"""
|
||||
handler = _QUOTA_REFRESH_HANDLERS.get(provider_type)
|
||||
if handler is not None:
|
||||
return handler
|
||||
raise InvalidRequestException("仅支持 Codex / Antigravity / Kiro 类型的 Provider 刷新限额")
|
||||
|
||||
|
||||
async def refresh_provider_quota_for_provider(
|
||||
db: Session,
|
||||
provider_id: str,
|
||||
codex_wham_usage_url: str,
|
||||
key_ids: list[str] | None = None,
|
||||
) -> dict:
|
||||
"""刷新指定 Provider 的限额信息（默认所有活跃 Key，可按 key_ids 限定）。"""
|
||||
provider = db.query(Provider).filter(Provider.id == provider_id).first()
|
||||
if not provider:
|
||||
raise NotFoundException(f"Provider {provider_id} 不存在")
|
||||
|
||||
provider_type = normalize_provider_type(getattr(provider, "provider_type", ""))
|
||||
if provider_type not in QUOTA_REFRESH_PROVIDER_TYPES:
|
||||
raise InvalidRequestException("仅支持 Codex / Antigravity / Kiro 类型的 Provider 刷新限额")
|
||||
pool_cfg = parse_pool_config(getattr(provider, "config", None))
|
||||
auto_remove_abnormal_keys = bool(pool_cfg and pool_cfg.auto_remove_banned_keys)
|
||||
|
||||
selected_key_ids: list[str] | None = None
|
||||
if key_ids is not None:
|
||||
deduped: list[str] = []
|
||||
seen: set[str] = set()
|
||||
for raw in key_ids:
|
||||
value = str(raw).strip()
|
||||
if not value or value in seen:
|
||||
continue
|
||||
seen.add(value)
|
||||
deduped.append(value)
|
||||
selected_key_ids = deduped
|
||||
|
||||
keys_query = (
|
||||
db.query(ProviderAPIKey)
|
||||
.options(
|
||||
defer(ProviderAPIKey.health_by_format),
|
||||
defer(ProviderAPIKey.circuit_breaker_by_format),
|
||||
defer(ProviderAPIKey.adjustment_history),
|
||||
defer(ProviderAPIKey.utilization_samples),
|
||||
)
|
||||
.filter(
|
||||
ProviderAPIKey.provider_id == provider_id,
|
||||
)
|
||||
)
|
||||
if selected_key_ids is None:
|
||||
# 全量刷新：活跃 key + 被系统自动标记 ACCOUNT_BLOCK 的 key。
|
||||
# 后者使 ACCOUNT_BLOCK 标记的 key 也参与刷新，账号恢复后可自动解除。
|
||||
# 注意：不能用宽泛的 oauth_invalid_reason IS NOT NULL，否则会纳入
|
||||
# 用户手动停用 (is_active=False) 但恰好也有 reason 的历史 key。
|
||||
keys_query = keys_query.filter(
|
||||
db_or(
|
||||
ProviderAPIKey.is_active.is_(True),
|
||||
ProviderAPIKey.oauth_invalid_reason.startswith("[ACCOUNT_BLOCK]"),
|
||||
)
|
||||
)
|
||||
else:
|
||||
if not selected_key_ids:
|
||||
return {
|
||||
"success": 0,
|
||||
"failed": 0,
|
||||
"total": 0,
|
||||
"results": [],
|
||||
"message": "未提供可刷新的 Key",
|
||||
"auto_removed": 0,
|
||||
}
|
||||
keys_query = keys_query.filter(ProviderAPIKey.id.in_(selected_key_ids))
|
||||
|
||||
keys = keys_query.all()
|
||||
if not keys:
|
||||
return {
|
||||
"success": 0,
|
||||
"failed": 0,
|
||||
"total": 0,
|
||||
"results": [],
|
||||
"message": "没有可刷新的 Key",
|
||||
"auto_removed": 0,
|
||||
}
|
||||
|
||||
endpoint = _select_refresh_endpoint(provider, provider_type)
|
||||
handler = _resolve_quota_refresh_handler(provider_type)
|
||||
|
||||
results: list[dict[str, Any]] = []
|
||||
success_count = 0
|
||||
failed_count = 0
|
||||
metadata_updates: dict[str, dict] = {} # key_id -> metadata
|
||||
state_updates: dict[str, dict[str, Any]] = {} # key_id -> model field updates
|
||||
|
||||
async def refresh_single_key(key: ProviderAPIKey) -> dict:
|
||||
try:
|
||||
return await handler(
|
||||
db=db,
|
||||
provider=provider,
|
||||
key=key,
|
||||
endpoint=endpoint,
|
||||
codex_wham_usage_url=codex_wham_usage_url,
|
||||
metadata_updates=metadata_updates,
|
||||
state_updates=state_updates,
|
||||
)
|
||||
except Exception as e:
|
||||
error_msg = str(e) or type(e).__name__
|
||||
logger.error("刷新 Key {} 限额失败: {}", key.id, error_msg)
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": error_msg,
|
||||
}
|
||||
|
||||
# 分批执行，每批最多 5 个并发
|
||||
batch_size = 5
|
||||
for i in range(0, len(keys), batch_size):
|
||||
batch = keys[i : i + batch_size]
|
||||
batch_tasks = [refresh_single_key(key) for key in batch]
|
||||
batch_results = await asyncio.gather(*batch_tasks)
|
||||
results.extend(batch_results)
|
||||
|
||||
# 统计本批次结果
|
||||
for result in batch_results:
|
||||
if result["status"] == "success":
|
||||
success_count += 1
|
||||
else:
|
||||
failed_count += 1
|
||||
|
||||
# 统一更新数据库（避免在并发任务中操作 session）
|
||||
auto_removed_contexts: list[tuple[str, str | None, list[str] | None]] = []
|
||||
result_index_by_key_id: dict[str, dict[str, Any]] = {}
|
||||
for result in results:
|
||||
rid = str(result.get("key_id", "")).strip()
|
||||
if rid:
|
||||
result_index_by_key_id[rid] = result
|
||||
|
||||
if metadata_updates or state_updates or auto_remove_abnormal_keys:
|
||||
for key in keys:
|
||||
key_dirty = False
|
||||
if key.id in metadata_updates:
|
||||
updates = metadata_updates[key.id]
|
||||
if isinstance(updates, dict):
|
||||
key.upstream_metadata = merge_upstream_metadata(key.upstream_metadata, updates)
|
||||
key_dirty = True
|
||||
if key.id in state_updates:
|
||||
updates = state_updates[key.id]
|
||||
if isinstance(updates, dict):
|
||||
for field_name, field_value in updates.items():
|
||||
setattr(key, field_name, field_value)
|
||||
key_dirty = True
|
||||
|
||||
if auto_remove_abnormal_keys:
|
||||
account_state = resolve_pool_account_state(
|
||||
provider_type=provider_type,
|
||||
upstream_metadata=getattr(key, "upstream_metadata", None),
|
||||
oauth_invalid_reason=getattr(key, "oauth_invalid_reason", None),
|
||||
)
|
||||
if should_auto_remove_account_state(account_state):
|
||||
key_id = str(getattr(key, "id", "") or "")
|
||||
auto_removed_contexts.append(
|
||||
(
|
||||
key_id,
|
||||
(
|
||||
str(getattr(key, "provider_id", "") or "")
|
||||
if getattr(key, "provider_id", None)
|
||||
else None
|
||||
),
|
||||
getattr(key, "allowed_models", None),
|
||||
)
|
||||
)
|
||||
if key_id and key_id in result_index_by_key_id:
|
||||
result_index_by_key_id[key_id]["auto_removed"] = True
|
||||
db.delete(key)
|
||||
continue
|
||||
|
||||
if key_dirty:
|
||||
db.add(key)
|
||||
|
||||
db.commit()
|
||||
|
||||
if auto_removed_contexts:
|
||||
cleanup_coros = []
|
||||
for key_id, pid, _allowed_models in auto_removed_contexts:
|
||||
if not key_id or not pid:
|
||||
continue
|
||||
cleanup_coros.append(pool_redis.clear_cooldown(pid, key_id))
|
||||
cleanup_coros.append(pool_redis.clear_cost(pid, key_id))
|
||||
if cleanup_coros:
|
||||
await asyncio.gather(*cleanup_coros, return_exceptions=True)
|
||||
for _key_id, pid, allowed_models in auto_removed_contexts:
|
||||
await run_delete_key_side_effects(
|
||||
db=db,
|
||||
provider_id=pid,
|
||||
deleted_key_allowed_models=allowed_models,
|
||||
)
|
||||
logger.warning(
|
||||
"[QUOTA_REFRESH] Provider {}: auto removed {} abnormal key(s): {}",
|
||||
provider_id,
|
||||
len(auto_removed_contexts),
|
||||
[ctx[0][:8] for ctx in auto_removed_contexts if ctx[0]],
|
||||
)
|
||||
|
||||
failed_details = [
|
||||
f"{r.get('key_name', r.get('key_id', '?'))}: {r.get('message', 'unknown')}"
|
||||
for r in results
|
||||
if r["status"] != "success"
|
||||
]
|
||||
if failed_details:
|
||||
logger.info(
|
||||
"[QUOTA_REFRESH] Provider {}: 成功 {}/{}, 失败 {} [{}]",
|
||||
provider_id,
|
||||
success_count,
|
||||
len(keys),
|
||||
failed_count,
|
||||
"; ".join(failed_details),
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
"[QUOTA_REFRESH] Provider {}: 成功 {}/{}",
|
||||
provider_id,
|
||||
success_count,
|
||||
len(keys),
|
||||
)
|
||||
|
||||
return {
|
||||
"success": success_count,
|
||||
"failed": failed_count,
|
||||
"total": len(keys),
|
||||
"results": results,
|
||||
"auto_removed": len(auto_removed_contexts),
|
||||
}
|
||||
222
_deprecated_py_src/services/provider_keys/key_side_effects.py
Normal file
222
_deprecated_py_src/services/provider_keys/key_side_effects.py
Normal file
@@ -0,0 +1,222 @@
|
||||
"""
|
||||
Provider Key 写操作后的副作用处理。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
from sqlalchemy import delete as sa_delete
|
||||
from sqlalchemy import update as sa_update
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.models.database import (
|
||||
GeminiFileMapping,
|
||||
ProviderAPIKey,
|
||||
Usage,
|
||||
VideoTask,
|
||||
)
|
||||
from src.services.cache.model_list_cache import invalidate_models_list_cache
|
||||
from src.services.cache.provider_cache import ProviderCacheService
|
||||
|
||||
_SQLITE_BATCH_SIZE = 900
|
||||
_DEFAULT_BATCH_SIZE = 2000
|
||||
|
||||
_CLEANUP_STAGES = (
|
||||
(
|
||||
"gemini_file_mappings",
|
||||
lambda batch: sa_delete(GeminiFileMapping).where(GeminiFileMapping.key_id.in_(batch)),
|
||||
),
|
||||
(
|
||||
"usage",
|
||||
lambda batch: sa_update(Usage)
|
||||
.where(Usage.provider_api_key_id.in_(batch))
|
||||
.values(provider_api_key_id=None),
|
||||
),
|
||||
(
|
||||
"video_tasks",
|
||||
lambda batch: sa_update(VideoTask).where(VideoTask.key_id.in_(batch)).values(key_id=None),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def cleanup_key_references(
|
||||
db: Session,
|
||||
key_ids: list[str],
|
||||
*,
|
||||
batch_size: int | None = None,
|
||||
stage_callback: Callable[[str, int], None] | None = None,
|
||||
) -> None:
|
||||
"""在删除 ProviderAPIKey 前，先显式处理关联表引用，降低级联删除/置空成本。
|
||||
|
||||
- gemini_file_mappings: 直接删除
|
||||
- usage / video_tasks: 先置空外键，保留快照与历史记录
|
||||
|
||||
PostgreSQL 下直接按 key_id/provider_api_key_id 批量处理；
|
||||
SQLite 下按 key_id 批次拆分，避免超出变量数限制。
|
||||
"""
|
||||
if not key_ids:
|
||||
return
|
||||
effective_batch_size = batch_size if batch_size is not None else _resolve_batch_size(db)
|
||||
for batch in iter_key_batches(key_ids, effective_batch_size):
|
||||
for stage_name, statement_factory in _CLEANUP_STAGES:
|
||||
if stage_callback is not None:
|
||||
stage_callback(stage_name, len(batch))
|
||||
db.execute(statement_factory(batch))
|
||||
|
||||
|
||||
def _resolve_batch_size(db: Session) -> int:
|
||||
try:
|
||||
bind = db.get_bind()
|
||||
dialect_name = str(getattr(getattr(bind, "dialect", None), "name", "") or "").lower()
|
||||
except Exception:
|
||||
dialect_name = ""
|
||||
if dialect_name == "sqlite":
|
||||
return _SQLITE_BATCH_SIZE
|
||||
return _DEFAULT_BATCH_SIZE
|
||||
|
||||
|
||||
def iter_key_batches(items: list[str], batch_size: int) -> list[list[str]]:
|
||||
"""将 key_ids 列表按 batch_size 拆分为子列表。"""
|
||||
if not items:
|
||||
return []
|
||||
if batch_size <= 0:
|
||||
return [list(items)]
|
||||
return [list(items[i : i + batch_size]) for i in range(0, len(items), batch_size)]
|
||||
|
||||
|
||||
async def run_update_key_side_effects(
|
||||
db: Session,
|
||||
key: ProviderAPIKey,
|
||||
key_id: str,
|
||||
auto_fetch_enabled_before: bool,
|
||||
auto_fetch_enabled_after: bool,
|
||||
include_patterns_before: list[str] | None,
|
||||
exclude_patterns_before: list[str] | None,
|
||||
allowed_models_before: set[str],
|
||||
) -> None:
|
||||
"""执行更新 Key 后的副作用。"""
|
||||
if not auto_fetch_enabled_before and auto_fetch_enabled_after:
|
||||
# 刚刚开启了 auto_fetch_models，同步执行模型获取
|
||||
logger.info("[AUTO_FETCH] Key {} 开启自动获取模型，同步执行模型获取", key_id)
|
||||
try:
|
||||
from src.services.model.fetch_scheduler import get_model_fetch_scheduler
|
||||
|
||||
scheduler = get_model_fetch_scheduler()
|
||||
# 同步等待模型获取完成，确保前端刷新时能看到最新数据
|
||||
await scheduler._fetch_models_for_key_by_id(key_id)
|
||||
# fetch_scheduler 可能在独立 session 更新 allowed_models，需刷新当前对象避免后续比较使用旧值。
|
||||
db.refresh(key)
|
||||
except Exception as e:
|
||||
logger.error(f"触发模型获取失败: {e}")
|
||||
# 不抛出异常，避免影响 Key 更新操作
|
||||
elif auto_fetch_enabled_before and not auto_fetch_enabled_after:
|
||||
# 关闭了 auto_fetch_models，只保留锁定的模型，清除自动获取的模型
|
||||
locked = key.locked_models or []
|
||||
if locked:
|
||||
key.allowed_models = locked
|
||||
logger.info(
|
||||
"[AUTO_FETCH] Key {} 关闭自动获取模型，保留 {} 个锁定模型",
|
||||
key_id,
|
||||
len(locked),
|
||||
)
|
||||
else:
|
||||
key.allowed_models = None
|
||||
logger.info(
|
||||
"[AUTO_FETCH] Key {} 关闭自动获取模型，无锁定模型，清空 allowed_models",
|
||||
key_id,
|
||||
)
|
||||
db.commit()
|
||||
db.refresh(key)
|
||||
elif auto_fetch_enabled_after:
|
||||
# auto_fetch_models 保持开启状态，检查过滤规则是否变更
|
||||
include_patterns_after = key.model_include_patterns
|
||||
exclude_patterns_after = key.model_exclude_patterns
|
||||
patterns_changed = (
|
||||
include_patterns_before != include_patterns_after
|
||||
or exclude_patterns_before != exclude_patterns_after
|
||||
)
|
||||
if patterns_changed:
|
||||
# 过滤规则变更，重新应用过滤（使用缓存的上游模型数据）
|
||||
logger.info("[AUTO_FETCH] Key {} 过滤规则变更，重新应用过滤", key_id)
|
||||
try:
|
||||
from src.services.model.fetch_scheduler import get_model_fetch_scheduler
|
||||
|
||||
scheduler = get_model_fetch_scheduler()
|
||||
await scheduler._fetch_models_for_key_by_id(key_id)
|
||||
# 重新应用过滤后，刷新当前对象以读取最新 allowed_models。
|
||||
db.refresh(key)
|
||||
except Exception as e:
|
||||
logger.error(f"重新应用过滤规则失败: {e}")
|
||||
|
||||
# 任何字段更新都清除缓存，确保缓存一致性
|
||||
# 包括 is_active、allowed_models、capabilities 等影响权限和行为的字段
|
||||
await ProviderCacheService.invalidate_provider_api_key_cache(key_id)
|
||||
|
||||
# 检查 allowed_models 是否有变化，触发缓存失效和自动关联
|
||||
allowed_models_after = set(key.allowed_models or [])
|
||||
if allowed_models_before != allowed_models_after and key.provider_id:
|
||||
from src.services.model.global_model import on_key_allowed_models_changed
|
||||
|
||||
await on_key_allowed_models_changed(
|
||||
db=db,
|
||||
provider_id=key.provider_id,
|
||||
allowed_models=list(key.allowed_models or []),
|
||||
)
|
||||
else:
|
||||
# allowed_models 未变化时，仍需清除 /v1/models 缓存（is_active、api_formats 变更会影响模型可用性）
|
||||
await invalidate_models_list_cache()
|
||||
|
||||
|
||||
async def run_create_key_side_effects(
|
||||
db: Session,
|
||||
provider_id: str,
|
||||
key: ProviderAPIKey,
|
||||
) -> None:
|
||||
"""执行创建 Key 后的副作用。"""
|
||||
# 如果开启了 auto_fetch_models，同步执行模型获取
|
||||
if key.auto_fetch_models:
|
||||
logger.info("[AUTO_FETCH] 新 Key {} 开启自动获取模型，同步执行模型获取", key.id)
|
||||
try:
|
||||
from src.services.model.fetch_scheduler import get_model_fetch_scheduler
|
||||
|
||||
scheduler = get_model_fetch_scheduler()
|
||||
# 同步等待模型获取完成，确保前端刷新时能看到最新数据
|
||||
await scheduler._fetch_models_for_key_by_id(key.id)
|
||||
except Exception as e:
|
||||
logger.error(f"触发模型获取失败: {e}")
|
||||
# 不抛出异常，避免影响 Key 创建操作
|
||||
|
||||
# 如果创建时指定了 allowed_models，触发自动关联检查（内部会清除 /v1/models 缓存）
|
||||
if key.allowed_models:
|
||||
from src.services.model.global_model import on_key_allowed_models_changed
|
||||
|
||||
await on_key_allowed_models_changed(
|
||||
db=db,
|
||||
provider_id=provider_id,
|
||||
allowed_models=list(key.allowed_models),
|
||||
)
|
||||
else:
|
||||
# 没有 allowed_models 时，仍需清除 /v1/models 缓存
|
||||
await invalidate_models_list_cache()
|
||||
|
||||
|
||||
async def run_delete_key_side_effects(
|
||||
db: Session,
|
||||
provider_id: str | None,
|
||||
deleted_key_allowed_models: list[str] | None,
|
||||
) -> None:
|
||||
"""执行删除 Key 后的副作用。"""
|
||||
_ = deleted_key_allowed_models
|
||||
if provider_id:
|
||||
from src.services.model.global_model import on_key_allowed_models_changed
|
||||
|
||||
await on_key_allowed_models_changed(
|
||||
db=db,
|
||||
provider_id=provider_id,
|
||||
skip_disassociate=True,
|
||||
)
|
||||
else:
|
||||
# 无 provider_id 时仅清除缓存
|
||||
await invalidate_models_list_cache()
|
||||
@@ -0,0 +1,373 @@
|
||||
"""
|
||||
号池额度主动探测调度器。
|
||||
|
||||
行为：
|
||||
- 当 provider.pool_advanced.probing_enabled=true 时启用
|
||||
- Key 以固定间隔主动触发额度刷新，用于检查 OAuth / 额度状态
|
||||
- 实际请求使用不会跳过定期探测；探测节流仅由刷新时间与主动探测时间控制
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.orm import load_only
|
||||
|
||||
from src.clients.redis_client import get_redis_client
|
||||
from src.core.logger import logger
|
||||
from src.core.provider_types import ProviderType, normalize_provider_type
|
||||
from src.database import create_session
|
||||
from src.models.database import Provider, ProviderAPIKey
|
||||
from src.services.provider.pool.config import parse_pool_config
|
||||
from src.services.provider_keys.key_quota_service import (
|
||||
CODEX_WHAM_USAGE_URL,
|
||||
QUOTA_REFRESH_PROVIDER_TYPES,
|
||||
refresh_provider_quota_for_provider,
|
||||
)
|
||||
from src.services.system.scheduler import get_scheduler
|
||||
|
||||
_REDIS_PREFIX = "ap:quota_probe:last"
|
||||
_DEFAULT_INTERVAL_MINUTES = 10
|
||||
_DEFAULT_SCAN_INTERVAL_SECONDS = 60
|
||||
_DEFAULT_MAX_KEYS_PER_PROVIDER = 50
|
||||
_MAX_INTERVAL_MINUTES = 1440
|
||||
|
||||
|
||||
def _probe_stamp_key(provider_id: str, key_id: str) -> str:
|
||||
return f"{_REDIS_PREFIX}:{provider_id}:{key_id}"
|
||||
|
||||
|
||||
def _to_float(value: Any) -> float | None:
|
||||
if isinstance(value, bool):
|
||||
return None
|
||||
if isinstance(value, (int, float)):
|
||||
return float(value)
|
||||
if isinstance(value, str):
|
||||
text = value.strip()
|
||||
if not text:
|
||||
return None
|
||||
try:
|
||||
return float(text)
|
||||
except ValueError:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _extract_quota_updated_at(provider_type: str, upstream_metadata: Any) -> int | None:
|
||||
if not isinstance(upstream_metadata, dict):
|
||||
return None
|
||||
|
||||
normalized = normalize_provider_type(provider_type)
|
||||
if normalized == ProviderType.CODEX.value:
|
||||
bucket = upstream_metadata.get("codex")
|
||||
elif normalized == ProviderType.KIRO.value:
|
||||
bucket = upstream_metadata.get("kiro")
|
||||
elif normalized == ProviderType.ANTIGRAVITY.value:
|
||||
bucket = upstream_metadata.get("antigravity")
|
||||
else:
|
||||
return None
|
||||
|
||||
if not isinstance(bucket, dict):
|
||||
return None
|
||||
|
||||
updated_at = _to_float(bucket.get("updated_at"))
|
||||
if updated_at is None or updated_at <= 0:
|
||||
return None
|
||||
|
||||
# 兼容毫秒时间戳
|
||||
if updated_at > 1_000_000_000_000:
|
||||
updated_at /= 1000
|
||||
return int(updated_at)
|
||||
|
||||
|
||||
def _parse_probe_stamp(raw_value: Any) -> int | None:
|
||||
parsed = _to_float(raw_value)
|
||||
if parsed is None or parsed <= 0:
|
||||
return None
|
||||
return int(parsed)
|
||||
|
||||
|
||||
def _normalize_probe_interval_minutes(raw_value: Any) -> int:
|
||||
parsed = _to_float(raw_value)
|
||||
if parsed is None:
|
||||
return _DEFAULT_INTERVAL_MINUTES
|
||||
return max(1, min(int(parsed), _MAX_INTERVAL_MINUTES))
|
||||
|
||||
|
||||
def _select_probe_key_ids(
|
||||
*,
|
||||
keys: list[ProviderAPIKey],
|
||||
provider_type: str,
|
||||
now_ts: int,
|
||||
interval_seconds: int,
|
||||
last_probe_timestamps: dict[str, int],
|
||||
limit: int,
|
||||
) -> list[str]:
|
||||
stale: list[tuple[int, str]] = []
|
||||
for key in keys:
|
||||
key_id = str(getattr(key, "id", "") or "")
|
||||
if not key_id:
|
||||
continue
|
||||
quota_updated_ts = _extract_quota_updated_at(
|
||||
provider_type,
|
||||
getattr(key, "upstream_metadata", None),
|
||||
)
|
||||
last_probe_ts = last_probe_timestamps.get(key_id)
|
||||
anchor_ts = max(quota_updated_ts or 0, last_probe_ts or 0)
|
||||
if anchor_ts <= 0 or (now_ts - anchor_ts) >= interval_seconds:
|
||||
stale.append((anchor_ts, key_id))
|
||||
|
||||
# anchor 越小说明越久未被探测/使用，优先探测
|
||||
stale.sort(key=lambda item: item[0])
|
||||
if limit > 0:
|
||||
stale = stale[:limit]
|
||||
return [key_id for _, key_id in stale]
|
||||
|
||||
|
||||
class PoolQuotaProbeScheduler:
|
||||
"""按号池高级配置执行额度主动探测。"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
scan_interval_raw = os.getenv(
|
||||
"POOL_QUOTA_PROBE_SCAN_INTERVAL_SECONDS",
|
||||
str(_DEFAULT_SCAN_INTERVAL_SECONDS),
|
||||
)
|
||||
max_keys_raw = os.getenv(
|
||||
"POOL_QUOTA_PROBE_MAX_KEYS_PER_PROVIDER",
|
||||
str(_DEFAULT_MAX_KEYS_PER_PROVIDER),
|
||||
)
|
||||
self.scan_interval_seconds = max(
|
||||
15, int(_to_float(scan_interval_raw) or _DEFAULT_SCAN_INTERVAL_SECONDS)
|
||||
)
|
||||
self.max_keys_per_provider = max(
|
||||
0, int(_to_float(max_keys_raw) or _DEFAULT_MAX_KEYS_PER_PROVIDER)
|
||||
)
|
||||
self.running = False
|
||||
|
||||
async def start(self) -> Any:
|
||||
if self.running:
|
||||
logger.warning("PoolQuotaProbeScheduler already running")
|
||||
return
|
||||
self.running = True
|
||||
logger.info(
|
||||
"PoolQuotaProbeScheduler started: scan={}s, max_keys_per_provider={}",
|
||||
self.scan_interval_seconds,
|
||||
self.max_keys_per_provider,
|
||||
)
|
||||
|
||||
scheduler = get_scheduler()
|
||||
scheduler.add_interval_job(
|
||||
self._scheduled_probe_check,
|
||||
seconds=self.scan_interval_seconds,
|
||||
job_id="pool_quota_probe_check",
|
||||
name="号池额度主动探测检查",
|
||||
)
|
||||
# 不在启动时立即探测：大号池场景下会阻塞启动并占用大量内存。
|
||||
# 首次探测由定时调度器在 scan_interval_seconds 后自动触发。
|
||||
|
||||
async def stop(self) -> Any:
|
||||
if not self.running:
|
||||
return
|
||||
self.running = False
|
||||
scheduler = get_scheduler()
|
||||
scheduler.remove_job("pool_quota_probe_check")
|
||||
logger.info("PoolQuotaProbeScheduler stopped")
|
||||
|
||||
async def _scheduled_probe_check(self) -> None:
|
||||
if not self.running:
|
||||
return
|
||||
await self._run_probe_cycle()
|
||||
|
||||
async def _load_probe_timestamps(
|
||||
self,
|
||||
*,
|
||||
redis_client: Any,
|
||||
provider_id: str,
|
||||
key_ids: list[str],
|
||||
) -> dict[str, int]:
|
||||
if redis_client is None or not key_ids:
|
||||
return {}
|
||||
redis_keys = [_probe_stamp_key(provider_id, key_id) for key_id in key_ids]
|
||||
try:
|
||||
values = await redis_client.mget(redis_keys)
|
||||
except Exception as exc:
|
||||
logger.debug("PoolQuotaProbeScheduler mget probe stamps failed: {}", exc)
|
||||
return {}
|
||||
|
||||
mapping: dict[str, int] = {}
|
||||
for key_id, raw in zip(key_ids, values, strict=False):
|
||||
parsed = _parse_probe_stamp(raw)
|
||||
if parsed is not None:
|
||||
mapping[key_id] = parsed
|
||||
return mapping
|
||||
|
||||
async def _mark_probe_timestamps(
|
||||
self,
|
||||
*,
|
||||
redis_client: Any,
|
||||
provider_id: str,
|
||||
key_ids: list[str],
|
||||
now_ts: int,
|
||||
interval_seconds: int,
|
||||
) -> None:
|
||||
if redis_client is None or not key_ids:
|
||||
return
|
||||
ttl_seconds = max(interval_seconds * 2, 120)
|
||||
try:
|
||||
pipe = redis_client.pipeline(transaction=False)
|
||||
value = str(now_ts)
|
||||
for key_id in key_ids:
|
||||
pipe.set(_probe_stamp_key(provider_id, key_id), value, ex=ttl_seconds)
|
||||
await pipe.execute()
|
||||
except Exception as exc:
|
||||
logger.debug("PoolQuotaProbeScheduler set probe stamps failed: {}", exc)
|
||||
|
||||
async def _run_probe_cycle(self) -> None:
|
||||
now_ts = int(time.time())
|
||||
redis_client = await get_redis_client(require_redis=False)
|
||||
|
||||
# 第一阶段：查出符合条件的 provider 列表（轻量查询）
|
||||
eligible_providers: list[tuple[str, str, int]] = [] # (id, type, interval_seconds)
|
||||
db = create_session()
|
||||
try:
|
||||
providers = db.query(Provider).filter(Provider.is_active == True).all() # noqa: E712
|
||||
for provider in providers:
|
||||
provider_id = str(getattr(provider, "id", "") or "")
|
||||
provider_type = normalize_provider_type(getattr(provider, "provider_type", ""))
|
||||
if not provider_id or provider_type not in QUOTA_REFRESH_PROVIDER_TYPES:
|
||||
continue
|
||||
|
||||
pool_cfg = parse_pool_config(getattr(provider, "config", None))
|
||||
if pool_cfg is None or not pool_cfg.probing_enabled:
|
||||
continue
|
||||
|
||||
interval_minutes = _normalize_probe_interval_minutes(
|
||||
pool_cfg.probing_interval_minutes
|
||||
)
|
||||
eligible_providers.append((provider_id, provider_type, interval_minutes * 60))
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
if not eligible_providers:
|
||||
return
|
||||
|
||||
# 第二阶段：逐个 provider 查询 key 并筛选探测目标
|
||||
# 避免一次性加载所有 provider 的全部 key 到内存
|
||||
for provider_id, provider_type, interval_seconds in eligible_providers:
|
||||
if not self.running:
|
||||
break
|
||||
|
||||
probe_key_ids = await self._select_keys_for_provider(
|
||||
provider_id=provider_id,
|
||||
provider_type=provider_type,
|
||||
interval_seconds=interval_seconds,
|
||||
now_ts=now_ts,
|
||||
redis_client=redis_client,
|
||||
)
|
||||
if not probe_key_ids:
|
||||
continue
|
||||
|
||||
# 先写探测节流时间戳，避免异常时高频重入
|
||||
await self._mark_probe_timestamps(
|
||||
redis_client=redis_client,
|
||||
provider_id=provider_id,
|
||||
key_ids=probe_key_ids,
|
||||
now_ts=now_ts,
|
||||
interval_seconds=interval_seconds,
|
||||
)
|
||||
|
||||
probe_db = create_session()
|
||||
try:
|
||||
result = await refresh_provider_quota_for_provider(
|
||||
db=probe_db,
|
||||
provider_id=provider_id,
|
||||
codex_wham_usage_url=CODEX_WHAM_USAGE_URL,
|
||||
key_ids=probe_key_ids,
|
||||
)
|
||||
logger.info(
|
||||
"[POOL_PROBE] Provider {} ({}) 静默探测完成: selected={}, success={}, failed={}",
|
||||
provider_id[:8],
|
||||
provider_type,
|
||||
len(probe_key_ids),
|
||||
int(result.get("success") or 0),
|
||||
int(result.get("failed") or 0),
|
||||
)
|
||||
except Exception as exc:
|
||||
try:
|
||||
probe_db.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
logger.warning(
|
||||
"[POOL_PROBE] Provider {} ({}) 静默探测失败: {}",
|
||||
provider_id[:8],
|
||||
provider_type,
|
||||
exc,
|
||||
)
|
||||
finally:
|
||||
probe_db.close()
|
||||
|
||||
async def _select_keys_for_provider(
|
||||
self,
|
||||
*,
|
||||
provider_id: str,
|
||||
provider_type: str,
|
||||
interval_seconds: int,
|
||||
now_ts: int,
|
||||
redis_client: Any,
|
||||
) -> list[str]:
|
||||
"""为单个 provider 筛选需要探测的 key，使用独立短生命周期 session。"""
|
||||
db = create_session()
|
||||
try:
|
||||
keys = (
|
||||
db.query(ProviderAPIKey)
|
||||
.options(
|
||||
load_only(
|
||||
ProviderAPIKey.id,
|
||||
ProviderAPIKey.provider_id,
|
||||
ProviderAPIKey.upstream_metadata,
|
||||
)
|
||||
)
|
||||
.filter(
|
||||
ProviderAPIKey.provider_id == provider_id,
|
||||
ProviderAPIKey.is_active == True, # noqa: E712
|
||||
)
|
||||
.all()
|
||||
)
|
||||
if not keys:
|
||||
return []
|
||||
|
||||
key_ids = [str(key.id) for key in keys if getattr(key, "id", None)]
|
||||
probe_stamps = await self._load_probe_timestamps(
|
||||
redis_client=redis_client,
|
||||
provider_id=provider_id,
|
||||
key_ids=key_ids,
|
||||
)
|
||||
return _select_probe_key_ids(
|
||||
keys=keys,
|
||||
provider_type=provider_type,
|
||||
now_ts=now_ts,
|
||||
interval_seconds=interval_seconds,
|
||||
last_probe_timestamps=probe_stamps,
|
||||
limit=self.max_keys_per_provider,
|
||||
)
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
_pool_quota_probe_scheduler: PoolQuotaProbeScheduler | None = None
|
||||
|
||||
|
||||
def get_pool_quota_probe_scheduler() -> PoolQuotaProbeScheduler:
|
||||
global _pool_quota_probe_scheduler
|
||||
if _pool_quota_probe_scheduler is None:
|
||||
_pool_quota_probe_scheduler = PoolQuotaProbeScheduler()
|
||||
return _pool_quota_probe_scheduler
|
||||
|
||||
|
||||
__all__ = [
|
||||
"PoolQuotaProbeScheduler",
|
||||
"get_pool_quota_probe_scheduler",
|
||||
"_select_probe_key_ids",
|
||||
]
|
||||
35
_deprecated_py_src/services/provider_keys/quota_cooldown.py
Normal file
35
_deprecated_py_src/services/provider_keys/quota_cooldown.py
Normal file
@@ -0,0 +1,35 @@
|
||||
"""配额冷却判定工具。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.services.scheduling.quota_skipper import is_key_quota_exhausted
|
||||
|
||||
|
||||
def resolve_effective_cooldown_reason(
|
||||
*,
|
||||
provider_type: str | None,
|
||||
key: Any,
|
||||
redis_reason: str | None,
|
||||
) -> str | None:
|
||||
"""返回 Key 的有效冷却原因。
|
||||
|
||||
规则：
|
||||
- Redis 冷却存在时，优先返回 Redis 原因（429/403/quota_exhausted 等）。
|
||||
- Redis 冷却不存在时，回退到 upstream_metadata 配额判断：
|
||||
若账号级配额耗尽（Codex/Kiro），返回 ``quota_exhausted``。
|
||||
"""
|
||||
if redis_reason:
|
||||
return redis_reason
|
||||
|
||||
try:
|
||||
exhausted, _ = is_key_quota_exhausted(provider_type, key, model_name="")
|
||||
except Exception:
|
||||
logger.opt(exception=True).debug(
|
||||
"quota_cooldown: is_key_quota_exhausted failed for key={}",
|
||||
getattr(key, "id", "?"),
|
||||
)
|
||||
return None
|
||||
return "quota_exhausted" if exhausted else None
|
||||
580
_deprecated_py_src/services/provider_keys/quota_reader.py
Normal file
580
_deprecated_py_src/services/provider_keys/quota_reader.py
Normal file
@@ -0,0 +1,580 @@
|
||||
"""Unified quota readers for provider key upstream metadata."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import time
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from src.core.provider_types import ProviderType, normalize_provider_type
|
||||
|
||||
|
||||
def _to_float(value: Any) -> float | None:
|
||||
try:
|
||||
parsed = float(value)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
if math.isnan(parsed) or math.isinf(parsed):
|
||||
return None
|
||||
return parsed
|
||||
|
||||
|
||||
def _normalize_plan(value: Any) -> str | None:
|
||||
if not isinstance(value, str):
|
||||
return None
|
||||
normalized = value.strip().lower()
|
||||
return normalized or None
|
||||
|
||||
|
||||
def _is_truthy_flag(value: Any) -> bool:
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
if isinstance(value, (int, float)):
|
||||
return value != 0
|
||||
if isinstance(value, str):
|
||||
normalized = value.strip().lower()
|
||||
return normalized in {"1", "true", "yes", "y"}
|
||||
return False
|
||||
|
||||
|
||||
def _extract_reason(source: dict[str, Any], *fields: str) -> str | None:
|
||||
for field in fields:
|
||||
value = source.get(field)
|
||||
if not isinstance(value, str):
|
||||
continue
|
||||
text = value.strip()
|
||||
if text:
|
||||
return text
|
||||
return None
|
||||
|
||||
|
||||
def _is_workspace_deactivated_reason(reason: str | None) -> bool:
|
||||
if not reason:
|
||||
return False
|
||||
lowered = reason.strip().lower()
|
||||
if not lowered:
|
||||
return False
|
||||
return "deactivated_workspace" in lowered
|
||||
|
||||
|
||||
def _pct_is_exhausted(value: Any) -> bool:
|
||||
pct = _to_float(value)
|
||||
if pct is None:
|
||||
return False
|
||||
return pct >= 100.0 - 1e-6
|
||||
|
||||
|
||||
def _format_percent(value: float) -> str:
|
||||
clamped = max(0.0, min(value, 100.0))
|
||||
return f"{clamped:.1f}%"
|
||||
|
||||
|
||||
def _format_quota_value(value: float) -> str:
|
||||
rounded = round(value)
|
||||
if abs(value - rounded) < 1e-6:
|
||||
return str(rounded)
|
||||
return f"{value:.1f}"
|
||||
|
||||
|
||||
def _has_quota_consumption(used_percent_raw: Any) -> bool:
|
||||
used = _to_float(used_percent_raw)
|
||||
if used is None:
|
||||
return False
|
||||
clamped_used = max(0.0, min(used, 100.0))
|
||||
return clamped_used > 1e-6
|
||||
|
||||
|
||||
def _format_reset_after(seconds_raw: Any) -> str | None:
|
||||
seconds = _to_float(seconds_raw)
|
||||
if seconds is None:
|
||||
return None
|
||||
|
||||
total_seconds = int(seconds)
|
||||
if total_seconds <= 0:
|
||||
return "已重置"
|
||||
|
||||
days = total_seconds // 86400
|
||||
hours = (total_seconds % 86400) // 3600
|
||||
minutes = (total_seconds % 3600) // 60
|
||||
|
||||
if days > 0:
|
||||
return f"{days}天{hours}小时后重置"
|
||||
if hours > 0:
|
||||
return f"{hours}小时{minutes}分钟后重置"
|
||||
if minutes > 0:
|
||||
return f"{minutes}分钟后重置"
|
||||
return "即将重置"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class QuotaExhaustedResult:
|
||||
exhausted: bool
|
||||
reason: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AccountBlockResult:
|
||||
blocked: bool
|
||||
code: str | None = None
|
||||
label: str | None = None
|
||||
reason: str | None = None
|
||||
|
||||
|
||||
class PoolQuotaReader(ABC):
|
||||
"""Read-only view over one provider namespace in upstream_metadata."""
|
||||
|
||||
namespace: str | None = None
|
||||
|
||||
def __init__(self, data: dict[str, Any] | None) -> None:
|
||||
self._data: dict[str, Any] = data if isinstance(data, dict) else {}
|
||||
|
||||
@abstractmethod
|
||||
def is_exhausted(self, model_name: str | None = None) -> QuotaExhaustedResult:
|
||||
"""Return whether this key/model should be skipped for quota exhaustion."""
|
||||
|
||||
@abstractmethod
|
||||
def usage_ratio(self) -> float | None:
|
||||
"""Return usage ratio within [0, 1], when available."""
|
||||
|
||||
@abstractmethod
|
||||
def plan_type(self) -> str | None:
|
||||
"""Return normalized plan type, when available."""
|
||||
|
||||
@abstractmethod
|
||||
def reset_seconds(self) -> float | None:
|
||||
"""Return seconds until next reset, when available."""
|
||||
|
||||
@abstractmethod
|
||||
def account_block(self) -> AccountBlockResult:
|
||||
"""Return account-level block state derived from metadata."""
|
||||
|
||||
@abstractmethod
|
||||
def display_summary(self) -> str | None:
|
||||
"""Return admin-facing quota summary string."""
|
||||
|
||||
def updated_at(self) -> int | None:
|
||||
updated_at = _to_float(self._data.get("updated_at"))
|
||||
if updated_at is None or updated_at <= 0:
|
||||
return None
|
||||
if updated_at > 1_000_000_000_000:
|
||||
updated_at /= 1000
|
||||
return int(updated_at)
|
||||
|
||||
|
||||
class NullQuotaReader(PoolQuotaReader):
|
||||
def is_exhausted(self, model_name: str | None = None) -> QuotaExhaustedResult:
|
||||
_ = model_name
|
||||
return QuotaExhaustedResult(exhausted=False)
|
||||
|
||||
def usage_ratio(self) -> float | None:
|
||||
return None
|
||||
|
||||
def plan_type(self) -> str | None:
|
||||
return None
|
||||
|
||||
def reset_seconds(self) -> float | None:
|
||||
return None
|
||||
|
||||
def account_block(self) -> AccountBlockResult:
|
||||
return AccountBlockResult(blocked=False)
|
||||
|
||||
def display_summary(self) -> str | None:
|
||||
return None
|
||||
|
||||
|
||||
class CodexQuotaReader(PoolQuotaReader):
|
||||
namespace = "codex"
|
||||
|
||||
def is_exhausted(self, model_name: str | None = None) -> QuotaExhaustedResult:
|
||||
_ = model_name
|
||||
exhausted_parts: list[str] = []
|
||||
if _pct_is_exhausted(self._data.get("primary_used_percent")):
|
||||
exhausted_parts.append("周限额剩余 0%")
|
||||
if _pct_is_exhausted(self._data.get("secondary_used_percent")):
|
||||
exhausted_parts.append("5H 限额剩余 0%")
|
||||
if exhausted_parts:
|
||||
return QuotaExhaustedResult(True, "Codex " + "，".join(exhausted_parts))
|
||||
return QuotaExhaustedResult(False)
|
||||
|
||||
def usage_ratio(self) -> float | None:
|
||||
values: list[float] = []
|
||||
for field in ("primary_used_percent", "secondary_used_percent"):
|
||||
parsed = _to_float(self._data.get(field))
|
||||
if parsed is None:
|
||||
continue
|
||||
values.append(max(0.0, min(parsed, 100.0)) / 100.0)
|
||||
if not values:
|
||||
return None
|
||||
return sum(values) / len(values)
|
||||
|
||||
def plan_type(self) -> str | None:
|
||||
return _normalize_plan(self._data.get("plan_type"))
|
||||
|
||||
def reset_seconds(self) -> float | None:
|
||||
candidates: list[float] = []
|
||||
for field in ("secondary_reset_seconds", "primary_reset_seconds"):
|
||||
parsed = _to_float(self._data.get(field))
|
||||
if parsed is None or parsed < 0:
|
||||
continue
|
||||
candidates.append(parsed)
|
||||
if not candidates:
|
||||
return None
|
||||
return min(candidates)
|
||||
|
||||
def account_block(self) -> AccountBlockResult:
|
||||
if not _is_truthy_flag(self._data.get("account_disabled")):
|
||||
return AccountBlockResult(blocked=False)
|
||||
reason = _extract_reason(self._data, "forbidden_reason", "ban_reason", "reason", "message")
|
||||
if _is_workspace_deactivated_reason(reason):
|
||||
return AccountBlockResult(
|
||||
blocked=True,
|
||||
code="workspace_deactivated",
|
||||
label="工作区停用",
|
||||
reason=reason or "工作区已停用",
|
||||
)
|
||||
return AccountBlockResult(
|
||||
blocked=True,
|
||||
code="account_forbidden",
|
||||
label="访问受限",
|
||||
reason=reason or "账号访问受限",
|
||||
)
|
||||
|
||||
def display_summary(self) -> str | None:
|
||||
parts: list[str] = []
|
||||
|
||||
primary_used = _to_float(self._data.get("primary_used_percent"))
|
||||
if primary_used is not None:
|
||||
part = f"周剩余 {_format_percent(100.0 - primary_used)}"
|
||||
reset_text = (
|
||||
_format_reset_after(self._data.get("primary_reset_seconds"))
|
||||
if _has_quota_consumption(primary_used)
|
||||
else None
|
||||
)
|
||||
if reset_text:
|
||||
part = f"{part} ({reset_text})"
|
||||
parts.append(part)
|
||||
|
||||
secondary_used = _to_float(self._data.get("secondary_used_percent"))
|
||||
if secondary_used is not None:
|
||||
part = f"5H剩余 {_format_percent(100.0 - secondary_used)}"
|
||||
reset_text = (
|
||||
_format_reset_after(self._data.get("secondary_reset_seconds"))
|
||||
if _has_quota_consumption(secondary_used)
|
||||
else None
|
||||
)
|
||||
if reset_text:
|
||||
part = f"{part} ({reset_text})"
|
||||
parts.append(part)
|
||||
|
||||
if parts:
|
||||
return " | ".join(parts)
|
||||
|
||||
has_credits = self._data.get("has_credits")
|
||||
credits_balance = _to_float(self._data.get("credits_balance"))
|
||||
if has_credits is True and credits_balance is not None:
|
||||
return f"积分 {credits_balance:.2f}"
|
||||
if has_credits is True:
|
||||
return "有积分"
|
||||
return None
|
||||
|
||||
|
||||
class KiroQuotaReader(PoolQuotaReader):
|
||||
namespace = "kiro"
|
||||
|
||||
def is_exhausted(self, model_name: str | None = None) -> QuotaExhaustedResult:
|
||||
_ = model_name
|
||||
remaining = _to_float(self._data.get("remaining"))
|
||||
if remaining is not None and remaining <= 0.0:
|
||||
return QuotaExhaustedResult(True, "Kiro 账号配额剩余 0")
|
||||
return QuotaExhaustedResult(False)
|
||||
|
||||
def usage_ratio(self) -> float | None:
|
||||
parsed = _to_float(self._data.get("usage_percentage"))
|
||||
if parsed is None:
|
||||
return None
|
||||
return max(0.0, min(parsed, 100.0)) / 100.0
|
||||
|
||||
def plan_type(self) -> str | None:
|
||||
subscription_title = _normalize_plan(self._data.get("subscription_title"))
|
||||
if not subscription_title:
|
||||
return None
|
||||
if "team" in subscription_title:
|
||||
return "team"
|
||||
if "free" in subscription_title:
|
||||
return "free"
|
||||
if "pro" in subscription_title:
|
||||
return "pro"
|
||||
if "plus" in subscription_title:
|
||||
return "plus"
|
||||
return subscription_title
|
||||
|
||||
def reset_seconds(self) -> float | None:
|
||||
next_reset_at = _to_float(self._data.get("next_reset_at"))
|
||||
if next_reset_at is None or next_reset_at <= 0:
|
||||
return None
|
||||
return max(0.0, next_reset_at - time.time())
|
||||
|
||||
def account_block(self) -> AccountBlockResult:
|
||||
if not _is_truthy_flag(self._data.get("is_banned")):
|
||||
return AccountBlockResult(blocked=False)
|
||||
reason = _extract_reason(self._data, "ban_reason", "reason", "message")
|
||||
return AccountBlockResult(
|
||||
blocked=True,
|
||||
code="account_banned",
|
||||
label="账号封禁",
|
||||
reason=reason or "Kiro 账号已封禁",
|
||||
)
|
||||
|
||||
def display_summary(self) -> str | None:
|
||||
if self._data.get("is_banned") is True:
|
||||
return "账号已封禁"
|
||||
|
||||
usage_percentage = _to_float(self._data.get("usage_percentage"))
|
||||
if usage_percentage is not None:
|
||||
remaining = 100.0 - usage_percentage
|
||||
current_usage = _to_float(self._data.get("current_usage"))
|
||||
usage_limit = _to_float(self._data.get("usage_limit"))
|
||||
if current_usage is not None and usage_limit is not None and usage_limit > 0:
|
||||
return (
|
||||
f"剩余 {_format_percent(remaining)} "
|
||||
f"({_format_quota_value(current_usage)}/{_format_quota_value(usage_limit)})"
|
||||
)
|
||||
return f"剩余 {_format_percent(remaining)}"
|
||||
|
||||
remaining = _to_float(self._data.get("remaining"))
|
||||
usage_limit = _to_float(self._data.get("usage_limit"))
|
||||
if remaining is not None and usage_limit is not None and usage_limit > 0:
|
||||
return f"剩余 {_format_quota_value(remaining)}/{_format_quota_value(usage_limit)}"
|
||||
return None
|
||||
|
||||
|
||||
class AntigravityQuotaReader(PoolQuotaReader):
|
||||
namespace = "antigravity"
|
||||
|
||||
def _quota_by_model(self) -> dict[str, Any]:
|
||||
quota_by_model = self._data.get("quota_by_model")
|
||||
if not isinstance(quota_by_model, dict):
|
||||
return {}
|
||||
return quota_by_model
|
||||
|
||||
def _used_percent(self, model_info: dict[str, Any]) -> float | None:
|
||||
used_percent = _to_float(model_info.get("used_percent"))
|
||||
if used_percent is not None:
|
||||
return max(0.0, min(used_percent, 100.0))
|
||||
remaining_fraction = _to_float(model_info.get("remaining_fraction"))
|
||||
if remaining_fraction is None:
|
||||
return None
|
||||
return max(0.0, min((1.0 - remaining_fraction) * 100.0, 100.0))
|
||||
|
||||
def is_exhausted(self, model_name: str | None = None) -> QuotaExhaustedResult:
|
||||
if not model_name:
|
||||
return QuotaExhaustedResult(False)
|
||||
model_quota = self._quota_by_model().get(model_name)
|
||||
if not isinstance(model_quota, dict):
|
||||
return QuotaExhaustedResult(False)
|
||||
|
||||
remaining_fraction = _to_float(model_quota.get("remaining_fraction"))
|
||||
if remaining_fraction is not None and remaining_fraction <= 0.0:
|
||||
return QuotaExhaustedResult(True, f"Antigravity 模型 {model_name} 配额剩余 0%")
|
||||
if _pct_is_exhausted(model_quota.get("used_percent")):
|
||||
return QuotaExhaustedResult(True, f"Antigravity 模型 {model_name} 配额剩余 0%")
|
||||
return QuotaExhaustedResult(False)
|
||||
|
||||
def usage_ratio(self) -> float | None:
|
||||
usage_values: list[float] = []
|
||||
for model_info in self._quota_by_model().values():
|
||||
if not isinstance(model_info, dict):
|
||||
continue
|
||||
used_percent = self._used_percent(model_info)
|
||||
if used_percent is None:
|
||||
continue
|
||||
usage_values.append(used_percent / 100.0)
|
||||
if not usage_values:
|
||||
return None
|
||||
return sum(usage_values) / len(usage_values)
|
||||
|
||||
def plan_type(self) -> str | None:
|
||||
return None
|
||||
|
||||
def reset_seconds(self) -> float | None:
|
||||
return None
|
||||
|
||||
def account_block(self) -> AccountBlockResult:
|
||||
if not _is_truthy_flag(self._data.get("is_forbidden")):
|
||||
return AccountBlockResult(blocked=False)
|
||||
reason = _extract_reason(self._data, "forbidden_reason", "reason", "message")
|
||||
return AccountBlockResult(
|
||||
blocked=True,
|
||||
code="account_forbidden",
|
||||
label="访问受限",
|
||||
reason=reason or "Antigravity 账户访问受限",
|
||||
)
|
||||
|
||||
def display_summary(self) -> str | None:
|
||||
if self._data.get("is_forbidden") is True:
|
||||
return "访问受限"
|
||||
|
||||
remaining_list: list[float] = []
|
||||
for raw_info in self._quota_by_model().values():
|
||||
if not isinstance(raw_info, dict):
|
||||
continue
|
||||
used_percent = self._used_percent(raw_info)
|
||||
if used_percent is None:
|
||||
continue
|
||||
remaining_list.append(max(0.0, min(100.0 - used_percent, 100.0)))
|
||||
|
||||
if not remaining_list:
|
||||
return None
|
||||
|
||||
min_remaining = min(remaining_list)
|
||||
if len(remaining_list) == 1:
|
||||
return f"剩余 {_format_percent(min_remaining)}"
|
||||
return f"最低剩余 {_format_percent(min_remaining)} ({len(remaining_list)} 模型)"
|
||||
|
||||
|
||||
class GeminiCliQuotaReader(PoolQuotaReader):
|
||||
namespace = "gemini_cli"
|
||||
|
||||
def _quota_by_model(self) -> dict[str, Any]:
|
||||
quota_by_model = self._data.get("quota_by_model")
|
||||
if not isinstance(quota_by_model, dict):
|
||||
return {}
|
||||
return quota_by_model
|
||||
|
||||
def _reset_at(self, model_info: dict[str, Any]) -> int | None:
|
||||
reset_at = _to_float(model_info.get("reset_at"))
|
||||
if reset_at is None or reset_at <= 0:
|
||||
return None
|
||||
if reset_at > 1_000_000_000_000:
|
||||
reset_at /= 1000
|
||||
return int(reset_at)
|
||||
|
||||
def _is_model_exhausted(self, model_info: dict[str, Any]) -> bool:
|
||||
if _is_truthy_flag(model_info.get("is_exhausted")):
|
||||
return True
|
||||
remaining_fraction = _to_float(model_info.get("remaining_fraction"))
|
||||
if remaining_fraction is not None and remaining_fraction <= 0.0:
|
||||
return True
|
||||
return _pct_is_exhausted(model_info.get("used_percent"))
|
||||
|
||||
def _active_exhausted_models(self) -> list[tuple[str, dict[str, Any], int | None]]:
|
||||
now = int(time.time())
|
||||
active: list[tuple[str, dict[str, Any], int | None]] = []
|
||||
for model_name, raw_info in self._quota_by_model().items():
|
||||
if not isinstance(raw_info, dict):
|
||||
continue
|
||||
if not self._is_model_exhausted(raw_info):
|
||||
continue
|
||||
reset_at = self._reset_at(raw_info)
|
||||
if reset_at is not None and reset_at <= now:
|
||||
continue
|
||||
active.append((str(model_name), raw_info, reset_at))
|
||||
return active
|
||||
|
||||
def is_exhausted(self, model_name: str | None = None) -> QuotaExhaustedResult:
|
||||
if not model_name:
|
||||
return QuotaExhaustedResult(False)
|
||||
model_quota = self._quota_by_model().get(model_name)
|
||||
if not isinstance(model_quota, dict) or not self._is_model_exhausted(model_quota):
|
||||
return QuotaExhaustedResult(False)
|
||||
|
||||
reset_at = self._reset_at(model_quota)
|
||||
if reset_at is not None:
|
||||
now = int(time.time())
|
||||
if reset_at <= now:
|
||||
return QuotaExhaustedResult(False)
|
||||
reset_text = _format_reset_after(reset_at - now)
|
||||
if reset_text:
|
||||
return QuotaExhaustedResult(
|
||||
True, f"Gemini CLI 模型 {model_name} 冷却中（{reset_text}）"
|
||||
)
|
||||
return QuotaExhaustedResult(True, f"Gemini CLI 模型 {model_name} 配额已耗尽")
|
||||
|
||||
def usage_ratio(self) -> float | None:
|
||||
active = self._active_exhausted_models()
|
||||
if not active:
|
||||
return None
|
||||
return 1.0
|
||||
|
||||
def plan_type(self) -> str | None:
|
||||
return _normalize_plan(self._data.get("plan_type")) or _normalize_plan(
|
||||
self._data.get("tier")
|
||||
)
|
||||
|
||||
def reset_seconds(self) -> float | None:
|
||||
now = int(time.time())
|
||||
reset_values = [
|
||||
reset_at - now
|
||||
for _, _, reset_at in self._active_exhausted_models()
|
||||
if reset_at is not None and reset_at > now
|
||||
]
|
||||
if not reset_values:
|
||||
return None
|
||||
return float(min(reset_values))
|
||||
|
||||
def account_block(self) -> AccountBlockResult:
|
||||
return AccountBlockResult(blocked=False)
|
||||
|
||||
def display_summary(self) -> str | None:
|
||||
active = self._active_exhausted_models()
|
||||
if not active:
|
||||
return None
|
||||
|
||||
active_sorted = sorted(
|
||||
active,
|
||||
key=lambda item: item[2] if item[2] is not None else 2**31 - 1,
|
||||
)
|
||||
first_model, _, first_reset_at = active_sorted[0]
|
||||
if len(active_sorted) == 1:
|
||||
if first_reset_at is not None:
|
||||
reset_text = _format_reset_after(first_reset_at - int(time.time()))
|
||||
if reset_text:
|
||||
return f"{first_model} 冷却中 ({reset_text})"
|
||||
return f"{first_model} 冷却中"
|
||||
|
||||
if first_reset_at is not None:
|
||||
reset_text = _format_reset_after(first_reset_at - int(time.time()))
|
||||
if reset_text:
|
||||
return f"{len(active_sorted)} 个模型冷却中（最早 {reset_text}）"
|
||||
return f"{len(active_sorted)} 个模型冷却中"
|
||||
|
||||
|
||||
_READER_CLASSES: dict[str, type[PoolQuotaReader]] = {
|
||||
ProviderType.CODEX: CodexQuotaReader,
|
||||
ProviderType.GEMINI_CLI: GeminiCliQuotaReader,
|
||||
ProviderType.KIRO: KiroQuotaReader,
|
||||
ProviderType.ANTIGRAVITY: AntigravityQuotaReader,
|
||||
}
|
||||
|
||||
|
||||
def get_quota_reader(provider_type: str | None, upstream_metadata: Any) -> PoolQuotaReader:
|
||||
"""Return a quota reader for one provider namespace in upstream_metadata."""
|
||||
|
||||
normalized_type = normalize_provider_type(provider_type)
|
||||
reader_cls = _READER_CLASSES.get(normalized_type)
|
||||
if reader_cls is None or not isinstance(upstream_metadata, dict):
|
||||
return NullQuotaReader(None)
|
||||
|
||||
namespace = reader_cls.namespace
|
||||
if not namespace:
|
||||
return NullQuotaReader(None)
|
||||
|
||||
data = upstream_metadata.get(namespace)
|
||||
if not isinstance(data, dict):
|
||||
return NullQuotaReader(None)
|
||||
return reader_cls(data)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AccountBlockResult",
|
||||
"AntigravityQuotaReader",
|
||||
"CodexQuotaReader",
|
||||
"GeminiCliQuotaReader",
|
||||
"KiroQuotaReader",
|
||||
"NullQuotaReader",
|
||||
"PoolQuotaReader",
|
||||
"QuotaExhaustedResult",
|
||||
"get_quota_reader",
|
||||
]
|
||||
@@ -0,0 +1,15 @@
|
||||
"""
|
||||
Provider Key 配额刷新策略模块。
|
||||
"""
|
||||
|
||||
from src.services.provider_keys.quota_refresh.antigravity_refresher import (
|
||||
refresh_antigravity_key_quota,
|
||||
)
|
||||
from src.services.provider_keys.quota_refresh.codex_refresher import refresh_codex_key_quota
|
||||
from src.services.provider_keys.quota_refresh.kiro_refresher import refresh_kiro_key_quota
|
||||
|
||||
__all__ = [
|
||||
"refresh_codex_key_quota",
|
||||
"refresh_antigravity_key_quota",
|
||||
"refresh_kiro_key_quota",
|
||||
]
|
||||
@@ -0,0 +1,25 @@
|
||||
"""Shared helpers for quota refresh strategies."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from src.models.database import ProviderAPIKey
|
||||
from src.services.provider.pool.account_state import OAUTH_REFRESH_FAILED_PREFIX
|
||||
|
||||
|
||||
def build_success_state_update(key: ProviderAPIKey) -> dict[str, Any]:
|
||||
"""配额刷新成功时的 state_updates 构建。
|
||||
|
||||
如果当前 key 携带 [REFRESH_FAILED] 标记，保留该标记（配额刷新不等于 token 刷新成功）。
|
||||
"""
|
||||
current_reason = str(getattr(key, "oauth_invalid_reason", None) or "").strip()
|
||||
if current_reason.startswith(OAUTH_REFRESH_FAILED_PREFIX):
|
||||
return {
|
||||
"oauth_invalid_at": getattr(key, "oauth_invalid_at", None),
|
||||
"oauth_invalid_reason": current_reason,
|
||||
}
|
||||
return {
|
||||
"oauth_invalid_at": None,
|
||||
"oauth_invalid_reason": None,
|
||||
}
|
||||
@@ -0,0 +1,136 @@
|
||||
"""
|
||||
Antigravity 配额刷新策略。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.models.database import Provider, ProviderAPIKey, ProviderEndpoint
|
||||
from src.services.provider.auth import get_provider_auth
|
||||
from src.services.provider_keys.quota_refresh._helpers import build_success_state_update
|
||||
|
||||
|
||||
async def refresh_antigravity_key_quota(
|
||||
*,
|
||||
db: Session,
|
||||
provider: Provider,
|
||||
key: ProviderAPIKey,
|
||||
endpoint: ProviderEndpoint | None,
|
||||
codex_wham_usage_url: str,
|
||||
metadata_updates: dict[str, dict],
|
||||
state_updates: dict[str, dict],
|
||||
) -> dict:
|
||||
"""刷新单个 Antigravity Key 的配额信息。"""
|
||||
_ = db
|
||||
_ = codex_wham_usage_url
|
||||
if endpoint is None:
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": "找不到有效的 gemini:chat/gemini:cli 端点",
|
||||
}
|
||||
|
||||
# 直接调用 /v1internal:fetchAvailableModels 获取 quotaInfo，无需发送真实对话请求
|
||||
auth_info = await get_provider_auth(endpoint, key)
|
||||
if not auth_info:
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": "缺少 OAuth 认证信息，请先授权/刷新 Token",
|
||||
}
|
||||
|
||||
access_token = str(auth_info.auth_value).removeprefix("Bearer ").strip()
|
||||
|
||||
from src.services.model.upstream_fetcher import (
|
||||
UpstreamModelsFetchContext,
|
||||
fetch_models_for_key,
|
||||
)
|
||||
from src.services.provider.adapters.antigravity.client import (
|
||||
AntigravityAccountForbiddenException,
|
||||
)
|
||||
from src.services.proxy_node.resolver import resolve_effective_proxy
|
||||
|
||||
effective_proxy = resolve_effective_proxy(
|
||||
getattr(provider, "proxy", None),
|
||||
getattr(key, "proxy", None),
|
||||
)
|
||||
|
||||
fetch_ctx = UpstreamModelsFetchContext(
|
||||
provider_type="antigravity",
|
||||
api_key_value=access_token,
|
||||
# antigravity fetcher 不依赖 endpoint mapping
|
||||
format_to_endpoint={},
|
||||
proxy_config=effective_proxy,
|
||||
auth_config=auth_info.decrypted_auth_config,
|
||||
)
|
||||
|
||||
try:
|
||||
_models, errors, ok, upstream_meta = await fetch_models_for_key(
|
||||
fetch_ctx, timeout_seconds=10.0
|
||||
)
|
||||
except AntigravityAccountForbiddenException as e:
|
||||
# 对齐 AM：所有 403 一律标记 is_forbidden；手动启用状态保持不变。
|
||||
state_updates[key.id] = {
|
||||
"oauth_invalid_at": datetime.now(timezone.utc),
|
||||
"oauth_invalid_reason": f"账户访问被禁止: {e.reason or e.message}",
|
||||
}
|
||||
# 更新 upstream_metadata 标记封禁状态
|
||||
metadata_updates[key.id] = {
|
||||
"antigravity": {
|
||||
"is_forbidden": True,
|
||||
"forbidden_reason": e.reason or e.message,
|
||||
"forbidden_at": int(time.time()),
|
||||
"updated_at": int(time.time()),
|
||||
}
|
||||
}
|
||||
logger.warning(
|
||||
"[ANTIGRAVITY_QUOTA] Key {} 账户访问被禁止，已更新账号状态: {}",
|
||||
key.id,
|
||||
e.reason or e.message,
|
||||
)
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "forbidden",
|
||||
"message": f"账户访问被禁止: {e.reason or e.message}",
|
||||
"is_forbidden": True,
|
||||
"auto_disabled": False,
|
||||
}
|
||||
|
||||
if ok and upstream_meta:
|
||||
# 刷新成功时清除之前的封禁标记（如果账户已恢复）
|
||||
if "antigravity" in upstream_meta:
|
||||
upstream_meta["antigravity"]["is_forbidden"] = False
|
||||
upstream_meta["antigravity"]["forbidden_reason"] = None
|
||||
upstream_meta["antigravity"]["forbidden_at"] = None
|
||||
metadata_updates[key.id] = upstream_meta
|
||||
state_updates[key.id] = build_success_state_update(key)
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "success",
|
||||
"metadata": upstream_meta,
|
||||
}
|
||||
|
||||
if ok and not upstream_meta:
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "no_metadata",
|
||||
"message": "响应中未包含配额信息",
|
||||
}
|
||||
|
||||
error_msg = "; ".join(errors) if errors else "fetchAvailableModels failed"
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": error_msg,
|
||||
}
|
||||
@@ -0,0 +1,507 @@
|
||||
"""
|
||||
Codex 配额刷新策略。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.crypto import crypto_service
|
||||
from src.core.logger import logger
|
||||
from src.models.database import Provider, ProviderAPIKey, ProviderEndpoint
|
||||
from src.services.provider.auth import get_provider_auth
|
||||
from src.services.provider.oauth_token import looks_like_token_invalidated
|
||||
from src.services.provider.pool.account_state import (
|
||||
OAUTH_ACCOUNT_BLOCK_PREFIX,
|
||||
OAUTH_EXPIRED_PREFIX,
|
||||
OAUTH_REQUEST_FAILED_PREFIX,
|
||||
)
|
||||
from src.services.provider_keys.auth_type import normalize_auth_type
|
||||
from src.services.provider_keys.codex_usage_parser import (
|
||||
parse_codex_usage_headers,
|
||||
parse_codex_wham_usage_response,
|
||||
)
|
||||
from src.services.provider_keys.quota_refresh._helpers import build_success_state_update
|
||||
|
||||
|
||||
def _normalize_plan_type(value: Any) -> str | None:
|
||||
if not isinstance(value, str):
|
||||
return None
|
||||
normalized = value.strip().lower()
|
||||
return normalized or None
|
||||
|
||||
|
||||
def _build_quota_exhausted_fallback_metadata(plan_type: str | None) -> dict[str, Any]:
|
||||
"""Build conservative Codex quota metadata when wham/usage returns 402."""
|
||||
normalized_plan = _normalize_plan_type(plan_type)
|
||||
metadata: dict[str, Any] = {"updated_at": int(time.time())}
|
||||
if normalized_plan:
|
||||
metadata["plan_type"] = normalized_plan
|
||||
# primary_* = weekly, secondary_* = 5H (aligned with parser semantics)
|
||||
metadata["primary_used_percent"] = 100.0
|
||||
if normalized_plan != "free":
|
||||
metadata["secondary_used_percent"] = 100.0
|
||||
return metadata
|
||||
|
||||
|
||||
def _extract_error_message_from_response(response: httpx.Response) -> str:
|
||||
"""Best-effort extraction of upstream error message for diagnostics."""
|
||||
try:
|
||||
payload = response.json()
|
||||
if isinstance(payload, dict):
|
||||
err = payload.get("error")
|
||||
if isinstance(err, dict):
|
||||
message = str(err.get("message", "")).strip()
|
||||
if message:
|
||||
return message
|
||||
if isinstance(err, str) and err.strip():
|
||||
return err.strip()
|
||||
message = str(payload.get("message", "")).strip()
|
||||
if message:
|
||||
return message
|
||||
except Exception:
|
||||
pass
|
||||
text = str(getattr(response, "text", "") or "").strip()
|
||||
return text[:300] if text else ""
|
||||
|
||||
|
||||
async def _build_codex_proxy_snapshot(
|
||||
effective_proxy: dict[str, Any] | None,
|
||||
) -> Any:
|
||||
from importlib import import_module
|
||||
|
||||
from src.services.request.execution_runtime_plan import ExecutionProxySnapshot
|
||||
|
||||
resolver_module = import_module("src.services.proxy_node.resolver")
|
||||
build_proxy_url_async = getattr(resolver_module, "build_proxy_url_async", None)
|
||||
get_system_proxy_config_async = getattr(resolver_module, "get_system_proxy_config_async", None)
|
||||
resolve_delegate_config_async = getattr(resolver_module, "resolve_delegate_config_async", None)
|
||||
resolve_proxy_info_async = getattr(resolver_module, "resolve_proxy_info_async", None)
|
||||
|
||||
proxy_config = effective_proxy
|
||||
if (not proxy_config or not proxy_config.get("enabled", True)) and callable(
|
||||
get_system_proxy_config_async
|
||||
):
|
||||
proxy_config = await get_system_proxy_config_async()
|
||||
if not proxy_config:
|
||||
return None
|
||||
|
||||
try:
|
||||
delegate_cfg = (
|
||||
await resolve_delegate_config_async(proxy_config)
|
||||
if callable(resolve_delegate_config_async)
|
||||
else None
|
||||
)
|
||||
proxy_url: str | None = None
|
||||
if (
|
||||
proxy_config
|
||||
and not (delegate_cfg and delegate_cfg.get("tunnel"))
|
||||
and callable(build_proxy_url_async)
|
||||
):
|
||||
proxy_url = await build_proxy_url_async(proxy_config)
|
||||
proxy_info = (
|
||||
await resolve_proxy_info_async(proxy_config)
|
||||
if callable(resolve_proxy_info_async)
|
||||
else {"url": proxy_url}
|
||||
)
|
||||
return ExecutionProxySnapshot.from_proxy_info(
|
||||
proxy_info,
|
||||
proxy_url=proxy_url,
|
||||
mode_override="tunnel" if delegate_cfg and delegate_cfg.get("tunnel") else None,
|
||||
node_id_override=(
|
||||
str(delegate_cfg.get("node_id") or "").strip() or None
|
||||
if delegate_cfg and delegate_cfg.get("tunnel")
|
||||
else None
|
||||
),
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning("Codex quota proxy snapshot build failed: {}", exc)
|
||||
return None
|
||||
|
||||
|
||||
async def _try_rust_codex_quota_response(
|
||||
*,
|
||||
key: ProviderAPIKey,
|
||||
provider: Provider,
|
||||
endpoint: ProviderEndpoint,
|
||||
url: str,
|
||||
headers: dict[str, str],
|
||||
proxy_snapshot: Any,
|
||||
) -> httpx.Response | None:
|
||||
from src.config import config
|
||||
from src.services.request.execution_runtime_plan import (
|
||||
ExecutionPlan,
|
||||
ExecutionPlanBody,
|
||||
ExecutionPlanTimeouts,
|
||||
)
|
||||
from src.services.request.execution_runtime_client import (
|
||||
ExecutionRuntimeClient,
|
||||
ExecutionRuntimeClientError,
|
||||
)
|
||||
|
||||
if config.execution_runtime_backend != "rust":
|
||||
return None
|
||||
|
||||
try:
|
||||
result = await ExecutionRuntimeClient().execute_sync_json(
|
||||
ExecutionPlan(
|
||||
request_id=f"codex-quota:{key.id}",
|
||||
candidate_id=None,
|
||||
provider_name="codex",
|
||||
provider_id=str(getattr(provider, "id", "") or ""),
|
||||
endpoint_id=str(getattr(endpoint, "id", "") or ""),
|
||||
key_id=str(getattr(key, "id", "") or ""),
|
||||
method="GET",
|
||||
url=url,
|
||||
headers=dict(headers),
|
||||
body=ExecutionPlanBody(),
|
||||
stream=False,
|
||||
provider_api_format="openai:cli",
|
||||
client_api_format="openai:cli",
|
||||
model_name="codex-wham-usage",
|
||||
proxy=proxy_snapshot,
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=30_000,
|
||||
read_ms=30_000,
|
||||
write_ms=30_000,
|
||||
pool_ms=30_000,
|
||||
total_ms=30_000,
|
||||
),
|
||||
)
|
||||
)
|
||||
except (
|
||||
ExecutionRuntimeClientError,
|
||||
httpx.HTTPError,
|
||||
json.JSONDecodeError,
|
||||
ValueError,
|
||||
AttributeError,
|
||||
TypeError,
|
||||
) as exc:
|
||||
logger.warning("Codex quota Rust fallback key_id={} url={}: {}", key.id, url, exc)
|
||||
return None
|
||||
|
||||
response_headers = dict(result.headers)
|
||||
if result.response_json is not None:
|
||||
response_headers.setdefault("content-type", "application/json")
|
||||
response_body = json.dumps(result.response_json, ensure_ascii=False).encode("utf-8")
|
||||
elif result.response_body_bytes is not None:
|
||||
response_body = result.response_body_bytes
|
||||
else:
|
||||
response_body = b""
|
||||
|
||||
return httpx.Response(
|
||||
status_code=result.status_code,
|
||||
request=httpx.Request("GET", url, headers=headers),
|
||||
headers=response_headers,
|
||||
content=response_body,
|
||||
)
|
||||
|
||||
|
||||
def _looks_like_account_deactivated(message: str | None) -> bool:
|
||||
lowered = str(message or "").strip().lower()
|
||||
return "account has been deactivated" in lowered or "account deactivated" in lowered
|
||||
|
||||
|
||||
def _looks_like_workspace_deactivated(message: str | None) -> bool:
|
||||
lowered = str(message or "").strip().lower()
|
||||
return "deactivated_workspace" in lowered or (
|
||||
"workspace" in lowered and "deactivated" in lowered
|
||||
)
|
||||
|
||||
|
||||
def _build_structured_invalid_reason(*, status_code: int, upstream_message: str | None) -> str:
|
||||
message = str(upstream_message or "").strip()
|
||||
|
||||
if status_code == 402 and _looks_like_workspace_deactivated(message):
|
||||
return f"{OAUTH_ACCOUNT_BLOCK_PREFIX}工作区已停用 (deactivated_workspace)"
|
||||
|
||||
if _looks_like_account_deactivated(message):
|
||||
detail = message or "OpenAI 账号已停用"
|
||||
return f"{OAUTH_ACCOUNT_BLOCK_PREFIX}{detail}"
|
||||
|
||||
# Codex 某些场景会返回 403，但语义仍是 access token 已失效/被轮换。
|
||||
# 这类异常可通过 refresh_token 恢复，不应落成账号级 block。
|
||||
if looks_like_token_invalidated(message):
|
||||
detail = message or "Codex Token 无效或已过期"
|
||||
return f"{OAUTH_EXPIRED_PREFIX}{detail}"
|
||||
|
||||
if status_code == 401:
|
||||
detail = message or "Codex Token 无效或已过期 (401)"
|
||||
return f"{OAUTH_EXPIRED_PREFIX}{detail}"
|
||||
|
||||
if status_code == 403:
|
||||
detail = message or "Codex 账户访问受限 (403)"
|
||||
return f"{OAUTH_ACCOUNT_BLOCK_PREFIX}{detail}"
|
||||
|
||||
return message
|
||||
|
||||
|
||||
def _build_soft_request_failure_reason(*, status_code: int, upstream_message: str | None) -> str:
|
||||
detail = str(upstream_message or "").strip() or f"Codex 请求失败 ({status_code})"
|
||||
return f"{OAUTH_REQUEST_FAILED_PREFIX}{detail}"
|
||||
|
||||
|
||||
def _get_current_invalid_reason(key: ProviderAPIKey) -> str:
|
||||
return str(getattr(key, "oauth_invalid_reason", None) or "").strip()
|
||||
|
||||
|
||||
def _merge_invalid_reason(current: str, candidate_reason: str) -> str:
|
||||
if not current:
|
||||
return candidate_reason
|
||||
if current.startswith(OAUTH_ACCOUNT_BLOCK_PREFIX):
|
||||
return current
|
||||
if current.startswith(OAUTH_EXPIRED_PREFIX) and candidate_reason.startswith(
|
||||
OAUTH_REQUEST_FAILED_PREFIX
|
||||
):
|
||||
return current
|
||||
return candidate_reason
|
||||
|
||||
|
||||
def _build_invalid_state_update(
|
||||
key: ProviderAPIKey,
|
||||
*,
|
||||
candidate_reason: str,
|
||||
) -> dict[str, Any]:
|
||||
current_reason = _get_current_invalid_reason(key)
|
||||
merged_reason = _merge_invalid_reason(current_reason, candidate_reason)
|
||||
if merged_reason == current_reason:
|
||||
return {
|
||||
"oauth_invalid_at": getattr(key, "oauth_invalid_at", None),
|
||||
"oauth_invalid_reason": merged_reason,
|
||||
}
|
||||
return {
|
||||
"oauth_invalid_at": datetime.now(timezone.utc),
|
||||
"oauth_invalid_reason": merged_reason,
|
||||
}
|
||||
|
||||
|
||||
async def refresh_codex_key_quota(
|
||||
*,
|
||||
db: Session,
|
||||
provider: Provider,
|
||||
key: ProviderAPIKey,
|
||||
endpoint: ProviderEndpoint | None,
|
||||
codex_wham_usage_url: str,
|
||||
metadata_updates: dict[str, dict],
|
||||
state_updates: dict[str, dict],
|
||||
) -> dict:
|
||||
"""刷新单个 Codex Key 的限额信息。"""
|
||||
_ = db
|
||||
if endpoint is None:
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": "找不到有效的 openai:cli 端点",
|
||||
}
|
||||
|
||||
# 获取认证信息（用于刷新 OAuth token）
|
||||
auth_info = await get_provider_auth(endpoint, key)
|
||||
|
||||
# 构建请求头
|
||||
headers: dict[str, Any] = {
|
||||
"Accept": "application/json",
|
||||
}
|
||||
if auth_info:
|
||||
headers[auth_info.auth_header] = auth_info.auth_value
|
||||
else:
|
||||
# 标准 API Key
|
||||
decrypted_key = crypto_service.decrypt(key.api_key)
|
||||
headers["Authorization"] = f"Bearer {decrypted_key}"
|
||||
|
||||
# 从 auth_config 中解密获取 plan_type 和 account_id
|
||||
oauth_plan_type = None
|
||||
oauth_account_id = None
|
||||
auth_type = normalize_auth_type(getattr(key, "auth_type", "api_key"))
|
||||
if auth_type == "oauth" and key.auth_config:
|
||||
try:
|
||||
decrypted_config = crypto_service.decrypt(key.auth_config)
|
||||
auth_config_data = json.loads(decrypted_config)
|
||||
if isinstance(auth_config_data, dict):
|
||||
oauth_plan_type = _normalize_plan_type(auth_config_data.get("plan_type"))
|
||||
raw_account_id = auth_config_data.get("account_id")
|
||||
if isinstance(raw_account_id, str):
|
||||
oauth_account_id = raw_account_id.strip() or None
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 如果有 account_id 且不是 free 账号（plan_type 缺失时默认携带，增强兼容性）
|
||||
if oauth_account_id and oauth_plan_type != "free":
|
||||
headers["chatgpt-account-id"] = oauth_account_id
|
||||
|
||||
# 解析代理配置（key 级别 > provider 级别 > 系统默认）
|
||||
from src.services.proxy_node.resolver import resolve_effective_proxy
|
||||
|
||||
effective_proxy = resolve_effective_proxy(
|
||||
getattr(provider, "proxy", None),
|
||||
getattr(key, "proxy", None),
|
||||
)
|
||||
proxy_snapshot = await _build_codex_proxy_snapshot(effective_proxy)
|
||||
|
||||
# 使用 wham/usage API 获取限额信息
|
||||
response = await _try_rust_codex_quota_response(
|
||||
key=key,
|
||||
provider=provider,
|
||||
endpoint=endpoint,
|
||||
url=codex_wham_usage_url,
|
||||
headers=headers,
|
||||
proxy_snapshot=proxy_snapshot,
|
||||
)
|
||||
if response is None:
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": "Codex 配额刷新仅支持 Rust executor",
|
||||
"status_code": 503,
|
||||
}
|
||||
|
||||
if response.status_code != 200:
|
||||
status_code = int(response.status_code)
|
||||
err_msg = _extract_error_message_from_response(response)
|
||||
|
||||
header_quota = parse_codex_usage_headers(dict(response.headers) if response.headers else {})
|
||||
if isinstance(header_quota, dict) and header_quota:
|
||||
metadata_updates[key.id] = {"codex": header_quota}
|
||||
|
||||
if status_code == 401:
|
||||
state_updates[key.id] = _build_invalid_state_update(
|
||||
key,
|
||||
candidate_reason=_build_structured_invalid_reason(
|
||||
status_code=401,
|
||||
upstream_message=err_msg,
|
||||
),
|
||||
)
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "auth_invalid",
|
||||
"message": f"wham/usage API 返回状态码 401{f': {err_msg}' if err_msg else ''}",
|
||||
"status_code": 401,
|
||||
"auto_disabled": False,
|
||||
}
|
||||
|
||||
if status_code == 402:
|
||||
if _looks_like_workspace_deactivated(err_msg):
|
||||
codex_meta = metadata_updates.get(key.id, {}).get("codex")
|
||||
if not isinstance(codex_meta, dict):
|
||||
codex_meta = {}
|
||||
codex_meta = {
|
||||
**codex_meta,
|
||||
"updated_at": int(time.time()),
|
||||
"account_disabled": True,
|
||||
"reason": "deactivated_workspace",
|
||||
"message": err_msg or "deactivated_workspace",
|
||||
}
|
||||
if oauth_plan_type and not codex_meta.get("plan_type"):
|
||||
codex_meta["plan_type"] = oauth_plan_type
|
||||
metadata_updates[key.id] = {"codex": codex_meta}
|
||||
state_updates[key.id] = _build_invalid_state_update(
|
||||
key,
|
||||
candidate_reason=_build_structured_invalid_reason(
|
||||
status_code=402,
|
||||
upstream_message=err_msg,
|
||||
),
|
||||
)
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "workspace_deactivated",
|
||||
"message": f"wham/usage API 返回状态码 402{f': {err_msg}' if err_msg else ''}",
|
||||
"status_code": 402,
|
||||
}
|
||||
|
||||
if key.id not in metadata_updates:
|
||||
metadata_updates[key.id] = {
|
||||
"codex": _build_quota_exhausted_fallback_metadata(oauth_plan_type)
|
||||
}
|
||||
state_updates[key.id] = build_success_state_update(key)
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "quota_exhausted",
|
||||
"message": f"wham/usage API 返回状态码 402{f': {err_msg}' if err_msg else ''}",
|
||||
"status_code": 402,
|
||||
}
|
||||
|
||||
if status_code == 403:
|
||||
candidate_reason = _build_structured_invalid_reason(
|
||||
status_code=403,
|
||||
upstream_message=err_msg,
|
||||
)
|
||||
if not looks_like_token_invalidated(err_msg):
|
||||
candidate_reason = _build_soft_request_failure_reason(
|
||||
status_code=403,
|
||||
upstream_message=err_msg,
|
||||
)
|
||||
state_updates[key.id] = _build_invalid_state_update(
|
||||
key,
|
||||
candidate_reason=candidate_reason,
|
||||
)
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "forbidden",
|
||||
"message": f"wham/usage API 返回状态码 403{f': {err_msg}' if err_msg else ''}",
|
||||
"status_code": 403,
|
||||
"auto_disabled": False,
|
||||
}
|
||||
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": (
|
||||
f"wham/usage API 返回状态码 {status_code}{f': {err_msg}' if err_msg else ''}"
|
||||
),
|
||||
"status_code": status_code,
|
||||
}
|
||||
|
||||
# 解析 JSON 响应
|
||||
try:
|
||||
data = response.json()
|
||||
except Exception:
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": "无法解析 wham/usage API 响应",
|
||||
}
|
||||
|
||||
# 解析限额信息
|
||||
try:
|
||||
metadata = parse_codex_wham_usage_response(data)
|
||||
except Exception as exc:
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": f"wham/usage 响应结构异常: {exc}",
|
||||
"status_code": response.status_code,
|
||||
}
|
||||
|
||||
if metadata:
|
||||
# 收集元数据，稍后统一更新数据库（存储到 codex 子对象）
|
||||
metadata_updates[key.id] = {"codex": metadata}
|
||||
state_updates[key.id] = build_success_state_update(key)
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "success",
|
||||
"metadata": metadata,
|
||||
}
|
||||
|
||||
# 响应成功但没有限额信息
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "no_metadata",
|
||||
"message": "响应中未包含限额信息",
|
||||
"status_code": response.status_code,
|
||||
}
|
||||
@@ -0,0 +1,159 @@
|
||||
"""
|
||||
Kiro 配额刷新策略。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.crypto import crypto_service
|
||||
from src.core.logger import logger
|
||||
from src.models.database import Provider, ProviderAPIKey, ProviderEndpoint
|
||||
from src.services.provider_keys.quota_refresh._helpers import build_success_state_update
|
||||
|
||||
|
||||
async def refresh_kiro_key_quota(
|
||||
*,
|
||||
db: Session,
|
||||
provider: Provider,
|
||||
key: ProviderAPIKey,
|
||||
endpoint: ProviderEndpoint | None,
|
||||
codex_wham_usage_url: str,
|
||||
metadata_updates: dict[str, dict],
|
||||
state_updates: dict[str, dict],
|
||||
) -> dict:
|
||||
"""刷新单个 Kiro Key 的配额信息。"""
|
||||
_ = db
|
||||
_ = endpoint
|
||||
_ = codex_wham_usage_url
|
||||
|
||||
from src.services.provider.adapters.kiro.usage import (
|
||||
KiroAccountBannedException,
|
||||
)
|
||||
from src.services.provider.adapters.kiro.usage import (
|
||||
fetch_kiro_usage_limits as _fetch_kiro_usage_limits,
|
||||
)
|
||||
from src.services.provider.adapters.kiro.usage import (
|
||||
parse_kiro_usage_response as _parse_kiro_usage_response,
|
||||
)
|
||||
|
||||
# Kiro: 直接使用 auth_config 调用 getUsageLimits API
|
||||
if not key.auth_config:
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": "缺少 Kiro 认证配置 (auth_config)",
|
||||
}
|
||||
|
||||
# 解密 auth_config
|
||||
try:
|
||||
decrypted_config = crypto_service.decrypt(key.auth_config)
|
||||
auth_config_data = json.loads(decrypted_config)
|
||||
except Exception:
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": "无法解密 auth_config，可能是加密密钥已更改",
|
||||
}
|
||||
|
||||
# 获取代理配置（key 级别 > provider 级别）
|
||||
from src.services.proxy_node.resolver import resolve_effective_proxy
|
||||
|
||||
proxy_config = resolve_effective_proxy(
|
||||
getattr(provider, "proxy", None),
|
||||
getattr(key, "proxy", None),
|
||||
)
|
||||
|
||||
# 调用 Kiro getUsageLimits API
|
||||
try:
|
||||
result = await _fetch_kiro_usage_limits(
|
||||
auth_config=auth_config_data,
|
||||
proxy_config=proxy_config,
|
||||
)
|
||||
except KiroAccountBannedException as e:
|
||||
# 账户被封禁，记录账号状态；手动启用状态保持不变。
|
||||
state_updates[key.id] = {
|
||||
"oauth_invalid_at": datetime.now(timezone.utc),
|
||||
"oauth_invalid_reason": f"账户已封禁: {e.reason or e.message}",
|
||||
}
|
||||
# 更新 upstream_metadata 标记封禁状态
|
||||
metadata_updates[key.id] = {
|
||||
"kiro": {
|
||||
"is_banned": True,
|
||||
"ban_reason": e.reason or e.message,
|
||||
"banned_at": int(time.time()),
|
||||
"updated_at": int(time.time()),
|
||||
}
|
||||
}
|
||||
logger.warning(
|
||||
"[KIRO_QUOTA] Key {} 账户已封禁，已更新账号状态: {}",
|
||||
key.id,
|
||||
e.reason or e.message,
|
||||
)
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "banned",
|
||||
"message": f"账户已封禁: {e.reason or e.message}",
|
||||
"is_banned": True,
|
||||
"auto_disabled": False,
|
||||
}
|
||||
except RuntimeError as e:
|
||||
error_msg = str(e)
|
||||
# 检查是否需要标记账号异常
|
||||
if "401" in error_msg or "认证失败" in error_msg:
|
||||
state_updates[key.id] = {
|
||||
"oauth_invalid_at": datetime.now(timezone.utc),
|
||||
"oauth_invalid_reason": "Kiro Token 无效或已过期",
|
||||
}
|
||||
logger.warning("[KIRO_QUOTA] Key {} Token 无效，已标记为异常", key.id)
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": error_msg,
|
||||
}
|
||||
|
||||
usage_data = result.get("usage_data")
|
||||
updated_auth_config = result.get("updated_auth_config")
|
||||
|
||||
# 解析限额信息
|
||||
metadata = _parse_kiro_usage_response(usage_data)
|
||||
|
||||
if metadata:
|
||||
# 刷新成功时清除之前的封禁标记（如果账户已恢复）
|
||||
metadata["is_banned"] = False
|
||||
metadata["ban_reason"] = None
|
||||
metadata["banned_at"] = None
|
||||
# 收集元数据，稍后统一更新数据库（存储到 kiro 子对象）
|
||||
metadata_updates[key.id] = {"kiro": metadata}
|
||||
state_updates[key.id] = build_success_state_update(key)
|
||||
|
||||
# 如果 auth_config 有更新（例如 token 刷新），也需要更新
|
||||
if updated_auth_config:
|
||||
try:
|
||||
new_auth_config_json = json.dumps(updated_auth_config)
|
||||
state_updates[key.id]["auth_config"] = crypto_service.encrypt(new_auth_config_json)
|
||||
except Exception as exc:
|
||||
logger.warning("更新 auth_config 失败 (key={}): {}", key.id, exc)
|
||||
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "success",
|
||||
"metadata": metadata,
|
||||
}
|
||||
|
||||
# 响应成功但没有限额信息
|
||||
return {
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "no_metadata",
|
||||
"message": "响应中未包含限额信息",
|
||||
}
|
||||
167
_deprecated_py_src/services/provider_keys/response_builder.py
Normal file
167
_deprecated_py_src/services/provider_keys/response_builder.py
Normal file
@@ -0,0 +1,167 @@
|
||||
"""
|
||||
Provider Key 响应对象构建器。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import asdict
|
||||
from typing import Any
|
||||
|
||||
from src.core.crypto import crypto_service
|
||||
from src.core.logger import logger
|
||||
from src.core.provider_oauth_utils import normalize_oauth_organizations
|
||||
from src.models.database import ProviderAPIKey
|
||||
from src.models.endpoint_models import EndpointAPIKeyResponse
|
||||
from src.services.provider_keys.auth_type import normalize_auth_type
|
||||
from src.services.provider_keys.status_snapshot_store import (
|
||||
normalize_oauth_expires_at,
|
||||
resolve_provider_key_status_snapshot,
|
||||
)
|
||||
|
||||
|
||||
def build_key_response(
|
||||
key: ProviderAPIKey,
|
||||
api_key_plain: str | None = None,
|
||||
*,
|
||||
provider_type: str | None = None,
|
||||
) -> EndpointAPIKeyResponse:
|
||||
"""构建 Key 响应对象。"""
|
||||
auth_type = normalize_auth_type(getattr(key, "auth_type", "api_key"))
|
||||
encrypted_api_key = str(getattr(key, "api_key", "") or "")
|
||||
request_count = int(getattr(key, "request_count", 0) or 0)
|
||||
success_count = int(getattr(key, "success_count", 0) or 0)
|
||||
total_response_time_ms = float(getattr(key, "total_response_time_ms", 0) or 0.0)
|
||||
rpm_limit = getattr(key, "rpm_limit", None)
|
||||
|
||||
if auth_type in ("service_account", "vertex_ai"):
|
||||
# Service Account 不显示占位符
|
||||
masked_key = "[Service Account]"
|
||||
elif auth_type == "oauth":
|
||||
masked_key = "[OAuth Token]"
|
||||
else:
|
||||
try:
|
||||
decrypted_key = crypto_service.decrypt(encrypted_api_key)
|
||||
masked_key = f"{decrypted_key[:8]}***{decrypted_key[-4:]}"
|
||||
except Exception:
|
||||
masked_key = "***ERROR***"
|
||||
|
||||
success_rate = success_count / request_count if request_count > 0 else 0.0
|
||||
avg_response_time_ms = total_response_time_ms / success_count if success_count > 0 else 0.0
|
||||
|
||||
is_adaptive = rpm_limit is None
|
||||
key_dict: dict[str, Any] = dict(getattr(key, "__dict__", {}))
|
||||
key_dict.pop("_sa_instance_state", None)
|
||||
key_dict.pop("api_key", None) # 移除敏感字段，避免泄露
|
||||
key_dict["auth_type"] = auth_type
|
||||
|
||||
# 提取 OAuth 元数据（如果是 OAuth 类型）
|
||||
oauth_expires_at = None
|
||||
oauth_email = None
|
||||
oauth_plan_type = None
|
||||
oauth_account_id = None
|
||||
oauth_account_name = None
|
||||
oauth_account_user_id = None
|
||||
auth_config: dict[str, Any] | None = None
|
||||
oauth_organizations: list[dict[str, object]] = []
|
||||
encrypted_auth_config = key_dict.pop("auth_config", None) # 移除敏感字段，避免泄露
|
||||
if auth_type == "oauth" and isinstance(encrypted_auth_config, str) and encrypted_auth_config:
|
||||
try:
|
||||
decrypted_config = crypto_service.decrypt(encrypted_auth_config)
|
||||
auth_config = json.loads(decrypted_config)
|
||||
oauth_expires_at = normalize_oauth_expires_at(auth_config.get("expires_at"))
|
||||
oauth_email = auth_config.get("email")
|
||||
oauth_plan_type = auth_config.get("plan_type") # Codex: plus/free/team/enterprise
|
||||
# Antigravity 使用 "tier" 字段（如 "PAID"/"FREE"），做小写化 fallback
|
||||
if not oauth_plan_type:
|
||||
ag_tier = auth_config.get("tier")
|
||||
if ag_tier and isinstance(ag_tier, str):
|
||||
oauth_plan_type = ag_tier.lower()
|
||||
oauth_account_id = auth_config.get("account_id") # Codex: chatgpt_account_id
|
||||
oauth_account_name = auth_config.get("account_name")
|
||||
oauth_account_user_id = auth_config.get("account_user_id")
|
||||
oauth_organizations = normalize_oauth_organizations(auth_config.get("organizations"))
|
||||
except Exception as e:
|
||||
logger.error("Failed to decrypt auth_config for key {}: {}", key.id, e)
|
||||
|
||||
if not provider_type:
|
||||
provider_rel = getattr(key, "provider", None)
|
||||
provider_type = (
|
||||
str(getattr(provider_rel, "provider_type", None) or "").strip()
|
||||
or str(getattr(provider_rel, "type", None) or "").strip()
|
||||
or None
|
||||
)
|
||||
|
||||
status_snapshot = resolve_provider_key_status_snapshot(
|
||||
key,
|
||||
provider_type=provider_type,
|
||||
auth_config=auth_config,
|
||||
oauth_expires_at=oauth_expires_at,
|
||||
)
|
||||
|
||||
# 从 health_by_format 计算汇总字段（便于列表展示）
|
||||
raw_health_by_format = getattr(key, "health_by_format", None)
|
||||
health_by_format = raw_health_by_format if isinstance(raw_health_by_format, dict) else {}
|
||||
raw_circuit_by_format = getattr(key, "circuit_breaker_by_format", None)
|
||||
circuit_by_format = raw_circuit_by_format if isinstance(raw_circuit_by_format, dict) else {}
|
||||
|
||||
# 计算整体健康度（取所有格式中的最低值）
|
||||
if health_by_format:
|
||||
health_scores = [float(h.get("health_score") or 1.0) for h in health_by_format.values()]
|
||||
min_health_score = min(health_scores) if health_scores else 1.0
|
||||
# 取最大的连续失败次数
|
||||
max_consecutive = max(
|
||||
(int(h.get("consecutive_failures") or 0) for h in health_by_format.values()),
|
||||
default=0,
|
||||
)
|
||||
# 取最近的失败时间
|
||||
failure_times = [
|
||||
h.get("last_failure_at") for h in health_by_format.values() if h.get("last_failure_at")
|
||||
]
|
||||
last_failure = max(failure_times) if failure_times else None
|
||||
else:
|
||||
min_health_score = 1.0
|
||||
max_consecutive = 0
|
||||
last_failure = None
|
||||
|
||||
# 检查是否有任何格式的熔断器打开
|
||||
any_circuit_open = any(c.get("open", False) for c in circuit_by_format.values())
|
||||
|
||||
key_dict.update(
|
||||
{
|
||||
"api_key_masked": masked_key,
|
||||
"api_key_plain": api_key_plain,
|
||||
"success_rate": success_rate,
|
||||
"avg_response_time_ms": round(avg_response_time_ms, 2),
|
||||
"is_adaptive": is_adaptive,
|
||||
"effective_limit": (
|
||||
getattr(
|
||||
key, "learned_rpm_limit", None
|
||||
) # 自适应模式：使用学习值，未学习时为 None（不限制）
|
||||
if is_adaptive
|
||||
else rpm_limit
|
||||
),
|
||||
# 汇总字段
|
||||
"health_score": min_health_score,
|
||||
"consecutive_failures": max_consecutive,
|
||||
"last_failure_at": last_failure,
|
||||
"circuit_breaker_open": any_circuit_open,
|
||||
# OAuth 相关
|
||||
"oauth_expires_at": oauth_expires_at,
|
||||
"oauth_email": oauth_email,
|
||||
"oauth_plan_type": oauth_plan_type,
|
||||
"oauth_account_id": oauth_account_id,
|
||||
"oauth_account_name": oauth_account_name,
|
||||
"oauth_account_user_id": oauth_account_user_id,
|
||||
"oauth_organizations": oauth_organizations,
|
||||
"oauth_invalid_at": status_snapshot.oauth.invalid_at,
|
||||
"oauth_invalid_reason": getattr(key, "oauth_invalid_reason", None),
|
||||
"status_snapshot": asdict(status_snapshot),
|
||||
}
|
||||
)
|
||||
|
||||
# 防御性：确保 api_formats 存在（历史数据可能为空/缺失）
|
||||
if "api_formats" not in key_dict or key_dict["api_formats"] is None:
|
||||
key_dict["api_formats"] = []
|
||||
|
||||
return EndpointAPIKeyResponse(**key_dict)
|
||||
@@ -0,0 +1,293 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import asdict
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.engine import Connection
|
||||
|
||||
from src.core.crypto import crypto_service
|
||||
from src.models.database import Provider, ProviderAPIKey
|
||||
from src.services.provider.pool.account_state import (
|
||||
AccountStatusSnapshot,
|
||||
OAuthStatusSnapshot,
|
||||
ProviderKeyStatusSnapshot,
|
||||
QuotaStatusSnapshot,
|
||||
build_provider_key_status_snapshot,
|
||||
resolve_oauth_status_snapshot,
|
||||
)
|
||||
|
||||
|
||||
def _clean_text(value: Any) -> str | None:
|
||||
if not isinstance(value, str):
|
||||
return None
|
||||
text = value.strip()
|
||||
return text or None
|
||||
|
||||
|
||||
def _coerce_bool(value: Any) -> bool:
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
if isinstance(value, (int, float)):
|
||||
return value != 0
|
||||
if isinstance(value, str):
|
||||
return value.strip().lower() in {"1", "true", "yes", "y"}
|
||||
return False
|
||||
|
||||
|
||||
def _coerce_int(value: Any) -> int | None:
|
||||
if isinstance(value, bool):
|
||||
return None
|
||||
if isinstance(value, int):
|
||||
return value
|
||||
if isinstance(value, float):
|
||||
return int(value)
|
||||
if isinstance(value, str):
|
||||
text = value.strip()
|
||||
if not text:
|
||||
return None
|
||||
try:
|
||||
return int(float(text))
|
||||
except ValueError:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _coerce_float(value: Any) -> float | None:
|
||||
if isinstance(value, bool):
|
||||
return None
|
||||
if isinstance(value, (int, float)):
|
||||
return float(value)
|
||||
if isinstance(value, str):
|
||||
text = value.strip()
|
||||
if not text:
|
||||
return None
|
||||
try:
|
||||
return float(text)
|
||||
except ValueError:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def extract_oauth_auth_config(key: ProviderAPIKey) -> dict[str, Any] | None:
|
||||
if str(getattr(key, "auth_type", "") or "").strip().lower() != "oauth":
|
||||
return None
|
||||
|
||||
auth_config_raw = getattr(key, "auth_config", None)
|
||||
if not auth_config_raw:
|
||||
return None
|
||||
|
||||
try:
|
||||
decrypted = crypto_service.decrypt(auth_config_raw)
|
||||
if isinstance(decrypted, str) and decrypted.strip():
|
||||
parsed = json.loads(decrypted)
|
||||
if isinstance(parsed, dict):
|
||||
return parsed
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def normalize_oauth_expires_at(raw: Any) -> int | None:
|
||||
value = _coerce_float(raw)
|
||||
if value is None or value <= 0:
|
||||
return None
|
||||
if value > 1_000_000_000_000:
|
||||
value /= 1000
|
||||
return int(value)
|
||||
|
||||
|
||||
def hydrate_provider_key_status_snapshot(raw: Any) -> ProviderKeyStatusSnapshot | None:
|
||||
if not isinstance(raw, dict):
|
||||
return None
|
||||
|
||||
oauth_raw = raw.get("oauth") if isinstance(raw.get("oauth"), dict) else {}
|
||||
account_raw = raw.get("account") if isinstance(raw.get("account"), dict) else {}
|
||||
quota_raw = raw.get("quota") if isinstance(raw.get("quota"), dict) else {}
|
||||
|
||||
return ProviderKeyStatusSnapshot(
|
||||
oauth=OAuthStatusSnapshot(
|
||||
code=_clean_text(oauth_raw.get("code")) or "none",
|
||||
label=_clean_text(oauth_raw.get("label")),
|
||||
reason=_clean_text(oauth_raw.get("reason")),
|
||||
expires_at=_coerce_int(oauth_raw.get("expires_at")),
|
||||
invalid_at=_coerce_int(oauth_raw.get("invalid_at")),
|
||||
source=_clean_text(oauth_raw.get("source")),
|
||||
requires_reauth=_coerce_bool(oauth_raw.get("requires_reauth")),
|
||||
expiring_soon=_coerce_bool(oauth_raw.get("expiring_soon")),
|
||||
),
|
||||
account=AccountStatusSnapshot(
|
||||
code=_clean_text(account_raw.get("code")) or "ok",
|
||||
label=_clean_text(account_raw.get("label")),
|
||||
reason=_clean_text(account_raw.get("reason")),
|
||||
blocked=_coerce_bool(account_raw.get("blocked")),
|
||||
source=_clean_text(account_raw.get("source")),
|
||||
recoverable=_coerce_bool(account_raw.get("recoverable")),
|
||||
),
|
||||
quota=QuotaStatusSnapshot(
|
||||
code=_clean_text(quota_raw.get("code")) or "unknown",
|
||||
label=_clean_text(quota_raw.get("label")),
|
||||
reason=_clean_text(quota_raw.get("reason")),
|
||||
exhausted=_coerce_bool(quota_raw.get("exhausted")),
|
||||
usage_ratio=_coerce_float(quota_raw.get("usage_ratio")),
|
||||
updated_at=_coerce_int(quota_raw.get("updated_at")),
|
||||
reset_seconds=_coerce_float(quota_raw.get("reset_seconds")),
|
||||
plan_type=_clean_text(quota_raw.get("plan_type")),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def resolve_provider_type_for_key(
|
||||
key: ProviderAPIKey,
|
||||
*,
|
||||
provider_type: str | None = None,
|
||||
connection: Connection | None = None,
|
||||
) -> str | None:
|
||||
normalized = _clean_text(provider_type)
|
||||
if normalized:
|
||||
return normalized
|
||||
|
||||
provider_rel = getattr(key, "__dict__", {}).get("provider")
|
||||
rel_type = _clean_text(getattr(provider_rel, "provider_type", None)) or _clean_text(
|
||||
getattr(provider_rel, "type", None)
|
||||
)
|
||||
if rel_type:
|
||||
return rel_type
|
||||
|
||||
provider_id = _clean_text(getattr(key, "provider_id", None))
|
||||
if provider_id and connection is not None:
|
||||
result = connection.execute(
|
||||
select(Provider.provider_type).where(Provider.id == provider_id)
|
||||
).scalar_one_or_none()
|
||||
return _clean_text(result)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def derive_oauth_expires_at(
|
||||
key: ProviderAPIKey,
|
||||
*,
|
||||
auth_config: dict[str, Any] | None = None,
|
||||
) -> int | None:
|
||||
if str(getattr(key, "auth_type", "") or "").strip().lower() != "oauth":
|
||||
return None
|
||||
|
||||
cfg = auth_config if isinstance(auth_config, dict) else extract_oauth_auth_config(key)
|
||||
if cfg:
|
||||
for field in ("expires_at", "expiresAt", "expiry", "exp"):
|
||||
expires_at = normalize_oauth_expires_at(cfg.get(field))
|
||||
if expires_at is not None:
|
||||
return expires_at
|
||||
|
||||
expires_dt = getattr(key, "expires_at", None)
|
||||
if isinstance(expires_dt, datetime):
|
||||
return int(expires_dt.timestamp())
|
||||
return None
|
||||
|
||||
|
||||
def resolve_provider_key_status_snapshot(
|
||||
key: ProviderAPIKey,
|
||||
*,
|
||||
provider_type: str | None = None,
|
||||
connection: Connection | None = None,
|
||||
auth_config: dict[str, Any] | None = None,
|
||||
oauth_expires_at: int | None = None,
|
||||
now_ts: int | None = None,
|
||||
) -> ProviderKeyStatusSnapshot:
|
||||
persisted_snapshot = hydrate_provider_key_status_snapshot(getattr(key, "status_snapshot", None))
|
||||
current_snapshot = _build_snapshot_from_current_fields(
|
||||
key,
|
||||
provider_type=provider_type,
|
||||
connection=connection,
|
||||
auth_config=auth_config,
|
||||
oauth_expires_at=oauth_expires_at,
|
||||
now_ts=now_ts,
|
||||
)
|
||||
if persisted_snapshot is None:
|
||||
return current_snapshot
|
||||
|
||||
resolved_oauth_expires_at = current_snapshot.oauth.expires_at
|
||||
if resolved_oauth_expires_at is None and persisted_snapshot.oauth.expires_at is not None:
|
||||
resolved_oauth_expires_at = int(persisted_snapshot.oauth.expires_at)
|
||||
resolved_oauth_invalid_at = current_snapshot.oauth.invalid_at
|
||||
if resolved_oauth_invalid_at is None and persisted_snapshot.oauth.invalid_at is not None:
|
||||
resolved_oauth_invalid_at = int(persisted_snapshot.oauth.invalid_at)
|
||||
oauth_invalid_reason = _clean_text(getattr(key, "oauth_invalid_reason", None)) or (
|
||||
persisted_snapshot.oauth.reason if persisted_snapshot is not None else None
|
||||
)
|
||||
|
||||
return ProviderKeyStatusSnapshot(
|
||||
oauth=resolve_oauth_status_snapshot(
|
||||
auth_type=str(getattr(key, "auth_type", "api_key") or "api_key"),
|
||||
oauth_expires_at=resolved_oauth_expires_at,
|
||||
oauth_invalid_at=resolved_oauth_invalid_at,
|
||||
oauth_invalid_reason=oauth_invalid_reason,
|
||||
now_ts=now_ts,
|
||||
),
|
||||
account=persisted_snapshot.account,
|
||||
quota=persisted_snapshot.quota,
|
||||
)
|
||||
|
||||
|
||||
def _build_snapshot_from_current_fields(
|
||||
key: ProviderAPIKey,
|
||||
*,
|
||||
provider_type: str | None = None,
|
||||
connection: Connection | None = None,
|
||||
auth_config: dict[str, Any] | None = None,
|
||||
oauth_expires_at: int | None = None,
|
||||
now_ts: int | None = None,
|
||||
) -> ProviderKeyStatusSnapshot:
|
||||
resolved_provider_type = resolve_provider_type_for_key(
|
||||
key, provider_type=provider_type, connection=connection
|
||||
)
|
||||
oauth_auth_config = (
|
||||
auth_config if isinstance(auth_config, dict) else extract_oauth_auth_config(key)
|
||||
)
|
||||
normalized_oauth_expires_at = normalize_oauth_expires_at(oauth_expires_at)
|
||||
resolved_oauth_expires_at = (
|
||||
normalized_oauth_expires_at
|
||||
if normalized_oauth_expires_at is not None
|
||||
else derive_oauth_expires_at(
|
||||
key,
|
||||
auth_config=oauth_auth_config,
|
||||
)
|
||||
)
|
||||
raw_invalid_at = getattr(key, "oauth_invalid_at", None)
|
||||
oauth_invalid_at = (
|
||||
int(raw_invalid_at.timestamp()) if isinstance(raw_invalid_at, datetime) else None
|
||||
)
|
||||
oauth_invalid_reason = _clean_text(getattr(key, "oauth_invalid_reason", None))
|
||||
|
||||
return build_provider_key_status_snapshot(
|
||||
auth_type=str(getattr(key, "auth_type", "api_key") or "api_key"),
|
||||
oauth_expires_at=resolved_oauth_expires_at,
|
||||
oauth_invalid_at=oauth_invalid_at,
|
||||
oauth_invalid_reason=oauth_invalid_reason,
|
||||
provider_type=resolved_provider_type,
|
||||
upstream_metadata=getattr(key, "upstream_metadata", None),
|
||||
now_ts=now_ts,
|
||||
)
|
||||
|
||||
|
||||
def sync_provider_key_status_snapshot(
|
||||
key: ProviderAPIKey,
|
||||
*,
|
||||
provider_type: str | None = None,
|
||||
connection: Connection | None = None,
|
||||
auth_config: dict[str, Any] | None = None,
|
||||
oauth_expires_at: int | None = None,
|
||||
) -> dict[str, Any]:
|
||||
snapshot = _build_snapshot_from_current_fields(
|
||||
key,
|
||||
provider_type=provider_type,
|
||||
connection=connection,
|
||||
auth_config=auth_config,
|
||||
oauth_expires_at=oauth_expires_at,
|
||||
)
|
||||
snapshot_dict = asdict(snapshot)
|
||||
key.status_snapshot = snapshot_dict
|
||||
return snapshot_dict
|
||||
Reference in New Issue
Block a user