mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +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
@@ -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": "响应中未包含限额信息",
|
||||
}
|
||||
Reference in New Issue
Block a user