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:
45
_deprecated_py_src/services/provider/__init__.py
Normal file
45
_deprecated_py_src/services/provider/__init__.py
Normal file
@@ -0,0 +1,45 @@
|
||||
"""
|
||||
Provider 服务模块
|
||||
|
||||
包含 Provider 管理、格式处理、传输层等功能。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from src.services.provider.format import normalize_endpoint_signature
|
||||
from src.services.provider.service import ProviderService
|
||||
from src.services.provider.transport import build_provider_url
|
||||
|
||||
__all__ = [
|
||||
"ProviderService",
|
||||
"normalize_endpoint_signature",
|
||||
"build_provider_url",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
"""Lazy attribute access to avoid import-time side effects.
|
||||
|
||||
Importing `src.services.provider` should not eagerly import the whole provider
|
||||
service stack (which can create circular imports during test collection).
|
||||
"""
|
||||
|
||||
if name == "ProviderService":
|
||||
from src.services.provider.service import ProviderService as _ProviderService
|
||||
|
||||
return _ProviderService
|
||||
if name == "normalize_endpoint_signature":
|
||||
from src.services.provider.format import (
|
||||
normalize_endpoint_signature as _normalize_endpoint_signature,
|
||||
)
|
||||
|
||||
return _normalize_endpoint_signature
|
||||
if name == "build_provider_url":
|
||||
from src.services.provider.transport import build_provider_url as _build_provider_url
|
||||
|
||||
return _build_provider_url
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -0,0 +1 @@
|
||||
"""Provider-specific adapters (antigravity, codex, etc.)."""
|
||||
@@ -0,0 +1 @@
|
||||
"""Antigravity integration package."""
|
||||
@@ -0,0 +1,627 @@
|
||||
"""Antigravity API 客户端(与 Antigravity-Manager 对齐)。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import random
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
from src.clients.http_client import HTTPClientPool
|
||||
from src.core.logger import logger
|
||||
from src.services.provider.adapters.antigravity.constants import (
|
||||
DAILY_BASE_URL,
|
||||
PROD_BASE_URL,
|
||||
SANDBOX_BASE_URL,
|
||||
VERSION_FETCH_URL,
|
||||
get_v1internal_extra_headers,
|
||||
parse_version_string,
|
||||
update_user_agent_version,
|
||||
)
|
||||
from src.services.provider.adapters.antigravity.rust_http import (
|
||||
execute_antigravity_rust_http_request,
|
||||
)
|
||||
from src.services.provider.adapters.antigravity.url_availability import url_availability
|
||||
|
||||
# loadCodeAssist 请求体 metadata
|
||||
_CODE_ASSIST_METADATA = {
|
||||
"ideType": "ANTIGRAVITY",
|
||||
}
|
||||
|
||||
# Duration 解析正则(与 AM 的 retry.rs 对齐)
|
||||
_DURATION_RE = re.compile(r"([\d.]+)\s*(ms|s|m|h)")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Retry-After / Duration 解析工具(对齐 AM upstream/retry.rs)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _parse_duration_ms(duration_str: str) -> int | None:
|
||||
"""解析 Duration 字符串 (e.g. '1.5s', '200ms', '1h16m0.667s'),返回毫秒。"""
|
||||
total_ms = 0.0
|
||||
matched = False
|
||||
for m in _DURATION_RE.finditer(duration_str):
|
||||
matched = True
|
||||
value = float(m.group(1))
|
||||
unit = m.group(2)
|
||||
if unit == "ms":
|
||||
total_ms += value
|
||||
elif unit == "s":
|
||||
total_ms += value * 1000
|
||||
elif unit == "m":
|
||||
total_ms += value * 60_000
|
||||
elif unit == "h":
|
||||
total_ms += value * 3_600_000
|
||||
return round(total_ms) if matched else None
|
||||
|
||||
|
||||
def parse_retry_delay(error_body: str | bytes | dict[str, Any]) -> float | None:
|
||||
"""从 429 错误响应体中提取 retry delay(秒)。
|
||||
|
||||
支持两种格式(与 AM 对齐):
|
||||
1. error.details[].@type="...RetryInfo" → retryDelay
|
||||
2. error.details[].metadata.quotaResetDelay
|
||||
"""
|
||||
try:
|
||||
data: dict[str, Any]
|
||||
if isinstance(error_body, (str, bytes)):
|
||||
data = json.loads(error_body)
|
||||
elif isinstance(error_body, dict):
|
||||
data = error_body
|
||||
else:
|
||||
return None
|
||||
|
||||
details = data.get("error", {}).get("details", [])
|
||||
if not isinstance(details, list):
|
||||
return None
|
||||
|
||||
# 方式 1: RetryInfo.retryDelay
|
||||
for detail in details:
|
||||
if not isinstance(detail, dict):
|
||||
continue
|
||||
type_str = detail.get("@type", "")
|
||||
if isinstance(type_str, str) and "RetryInfo" in type_str:
|
||||
retry_delay = detail.get("retryDelay")
|
||||
if isinstance(retry_delay, str):
|
||||
ms = _parse_duration_ms(retry_delay)
|
||||
if ms is not None:
|
||||
return min((ms + 200) / 1000.0, 30.0)
|
||||
|
||||
# 方式 2: metadata.quotaResetDelay
|
||||
for detail in details:
|
||||
if not isinstance(detail, dict):
|
||||
continue
|
||||
metadata = detail.get("metadata")
|
||||
if not isinstance(metadata, dict):
|
||||
continue
|
||||
quota_delay = metadata.get("quotaResetDelay")
|
||||
if isinstance(quota_delay, str):
|
||||
ms = _parse_duration_ms(quota_delay)
|
||||
if ms is not None:
|
||||
return min((ms + 200) / 1000.0, 30.0)
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fallback 判断(对齐 AM upstream/client.rs: should_try_next_endpoint)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _should_fallback_status(status_code: int) -> bool:
|
||||
"""判断是否应该 fallback 到下一个端点。
|
||||
|
||||
与 AM 对齐:429 + 408(超时) + 404 + 所有 5xx。
|
||||
4xx 客户端错误(400/401/403 等)不 fallback——换 URL 也不会成功。
|
||||
"""
|
||||
if status_code == 429:
|
||||
return True
|
||||
if status_code in (404, 408):
|
||||
return True
|
||||
if 500 <= status_code < 600: # noqa: PLR2004
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 账户封禁异常(对齐 AM: is_forbidden 标志)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class AntigravityAccountForbiddenException(Exception):
|
||||
"""Antigravity 账户被封禁/禁止访问异常。
|
||||
|
||||
当 API 返回 403 Forbidden 时抛出,表示账户权限被撤销。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
message: str = "账户访问被禁止",
|
||||
status_code: int = 403,
|
||||
reason: str | None = None,
|
||||
):
|
||||
super().__init__(message)
|
||||
self.message = message
|
||||
self.status_code = status_code
|
||||
self.reason = reason
|
||||
|
||||
|
||||
def _extract_forbidden_reason(response_text: str) -> str | None:
|
||||
"""从 403 响应体中提取封禁原因。
|
||||
|
||||
尝试解析 JSON 响应中的 error.message 字段。
|
||||
"""
|
||||
if not response_text:
|
||||
return None
|
||||
try:
|
||||
data = json.loads(response_text)
|
||||
if isinstance(data, dict):
|
||||
error = data.get("error")
|
||||
if isinstance(error, dict):
|
||||
message = error.get("message")
|
||||
if isinstance(message, str) and message.strip():
|
||||
return message.strip()
|
||||
# 直接在顶层查找 message
|
||||
message = data.get("message")
|
||||
if isinstance(message, str) and message.strip():
|
||||
return message.strip()
|
||||
except Exception:
|
||||
pass
|
||||
# 如果无法解析,返回原始文本的前 100 个字符
|
||||
if len(response_text) > 100:
|
||||
return response_text[:100] + "..."
|
||||
return response_text if response_text.strip() else None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 从 loadCodeAssist 响应中提取信息
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def extract_tier_id(data: dict[str, Any]) -> str:
|
||||
"""从 loadCodeAssist 响应中提取 tier ID。
|
||||
|
||||
优先选 allowedTiers 中 isDefault=true 的,fallback 到第一个,最终 fallback 到 "LEGACY"。
|
||||
"""
|
||||
allowed_tiers = data.get("allowedTiers")
|
||||
if not isinstance(allowed_tiers, list):
|
||||
return "LEGACY"
|
||||
|
||||
# 第一轮:找 isDefault
|
||||
for tier in allowed_tiers:
|
||||
if isinstance(tier, dict) and tier.get("isDefault") is True:
|
||||
tier_id = tier.get("id", "")
|
||||
if isinstance(tier_id, str) and tier_id.strip():
|
||||
return tier_id.strip()
|
||||
|
||||
# 第二轮:取第一个有 id 的
|
||||
for tier in allowed_tiers:
|
||||
if isinstance(tier, dict):
|
||||
tier_id = tier.get("id", "")
|
||||
if isinstance(tier_id, str) and tier_id.strip():
|
||||
return tier_id.strip()
|
||||
|
||||
return "LEGACY"
|
||||
|
||||
|
||||
def extract_project_id(data: dict[str, Any]) -> str:
|
||||
"""从响应中提取 project_id,兼容 string 和 {"id": "..."} 两种格式。"""
|
||||
raw = data.get("cloudaicompanionProject")
|
||||
if isinstance(raw, str) and raw.strip():
|
||||
return raw.strip()
|
||||
if isinstance(raw, dict):
|
||||
pid = raw.get("id", "")
|
||||
if isinstance(pid, str) and pid.strip():
|
||||
return pid.strip()
|
||||
return ""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 核心 API 客户端函数
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def load_code_assist(
|
||||
access_token: str,
|
||||
proxy_config: dict[str, Any] | None = None,
|
||||
*,
|
||||
timeout_seconds: float = 10.0,
|
||||
) -> dict[str, Any]:
|
||||
"""调用 /v1internal:loadCodeAssist 获取账户信息。
|
||||
|
||||
对齐 AM project_resolver.rs:Sandbox 优先,避免 Prod 429。
|
||||
|
||||
注意:
|
||||
- email 需通过 Google userinfo API 获取(由 enrich_auth_config 复用已有逻辑)
|
||||
- 这里仅负责 project_id / tier 等信息
|
||||
"""
|
||||
if not access_token:
|
||||
raise ValueError("missing access_token")
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {access_token}",
|
||||
"Content-Type": "application/json",
|
||||
**get_v1internal_extra_headers(),
|
||||
}
|
||||
body = {"metadata": _CODE_ASSIST_METADATA}
|
||||
|
||||
# Sandbox 优先(与 AM 对齐:避免 Prod 429)
|
||||
urls = url_availability.get_ordered_urls(prefer_daily=True)
|
||||
if not urls:
|
||||
urls = [SANDBOX_BASE_URL, DAILY_BASE_URL, PROD_BASE_URL]
|
||||
last_exc: Exception | None = None
|
||||
|
||||
for base_url in urls:
|
||||
try:
|
||||
url = f"{base_url}/v1internal:loadCodeAssist"
|
||||
resp = await execute_antigravity_rust_http_request(
|
||||
method="POST",
|
||||
url=url,
|
||||
headers=headers,
|
||||
body=body,
|
||||
proxy_config=proxy_config,
|
||||
request_id=f"antigravity:load-code-assist:{base_url}",
|
||||
provider_api_format="antigravity:load_code_assist",
|
||||
timeout_seconds=timeout_seconds,
|
||||
content_type="application/json",
|
||||
)
|
||||
if resp is None:
|
||||
client = await HTTPClientPool.get_proxy_client(proxy_config)
|
||||
resp = await client.post(
|
||||
url,
|
||||
json=body,
|
||||
headers=headers,
|
||||
timeout=timeout_seconds,
|
||||
)
|
||||
if 200 <= resp.status_code < 300:
|
||||
url_availability.mark_success(base_url)
|
||||
data = resp.json()
|
||||
return data if isinstance(data, dict) else {}
|
||||
|
||||
# 可 fallback 的错误:标记不可用,尝试下一个 URL
|
||||
if _should_fallback_status(resp.status_code):
|
||||
url_availability.mark_unavailable(base_url)
|
||||
|
||||
# 429:尝试解析 Retry-After 并等待
|
||||
if resp.status_code == 429:
|
||||
delay = parse_retry_delay(resp.text)
|
||||
if delay and delay > 0:
|
||||
logger.debug(
|
||||
"[antigravity] loadCodeAssist 429, waiting {:.1f}s before fallback",
|
||||
delay,
|
||||
)
|
||||
await asyncio.sleep(min(delay, 5.0))
|
||||
|
||||
last_exc = RuntimeError(
|
||||
f"loadCodeAssist failed: status={resp.status_code} base_url={base_url}"
|
||||
)
|
||||
continue
|
||||
|
||||
# 不可 fallback 的 4xx(400/401/403 等):直接抛出,换 URL 也不会成功
|
||||
raise RuntimeError(
|
||||
f"loadCodeAssist failed: status={resp.status_code} base_url={base_url} "
|
||||
f"body={resp.text[:200] if resp.text else ''}"
|
||||
)
|
||||
except RuntimeError:
|
||||
raise
|
||||
except Exception as e:
|
||||
url_availability.mark_unavailable(base_url)
|
||||
last_exc = e
|
||||
continue
|
||||
|
||||
raise last_exc or RuntimeError("loadCodeAssist failed: all endpoints exhausted")
|
||||
|
||||
|
||||
async def onboard_user(
|
||||
access_token: str,
|
||||
tier_id: str = "LEGACY",
|
||||
proxy_config: dict[str, Any] | None = None,
|
||||
*,
|
||||
max_attempts: int = 5,
|
||||
poll_interval: float = 2.0,
|
||||
timeout_seconds: float = 30.0,
|
||||
) -> str:
|
||||
"""调用 /v1internal:onboardUser 激活账号并获取 project_id。
|
||||
|
||||
当 loadCodeAssist 返回 allowedTiers 但没有 cloudaicompanionProject 时,
|
||||
需要先通过 onboardUser 选择 tier 来分配 project。
|
||||
|
||||
改进(对齐 AM):
|
||||
- 使用 url_availability 做多 URL fallback
|
||||
- 轮询期间的临时网络错误会 continue 而非直接终止
|
||||
|
||||
Args:
|
||||
access_token: OAuth access token
|
||||
tier_id: 要选择的 tier ID(从 allowedTiers 中提取)
|
||||
proxy_config: 代理配置
|
||||
max_attempts: 最大轮询次数(onboardUser 是异步操作)
|
||||
poll_interval: 轮询间隔(秒)
|
||||
timeout_seconds: 单次请求超时
|
||||
|
||||
Returns:
|
||||
project_id
|
||||
|
||||
Raises:
|
||||
RuntimeError: 激活失败
|
||||
"""
|
||||
if not access_token:
|
||||
raise ValueError("missing access_token")
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {access_token}",
|
||||
"Content-Type": "application/json",
|
||||
**get_v1internal_extra_headers(),
|
||||
}
|
||||
body = {
|
||||
"tierId": tier_id,
|
||||
"metadata": _CODE_ASSIST_METADATA,
|
||||
}
|
||||
|
||||
# 使用 url_availability 选择端点(与其他函数一致)
|
||||
urls = url_availability.get_ordered_urls(prefer_daily=True)
|
||||
if not urls:
|
||||
urls = [SANDBOX_BASE_URL, DAILY_BASE_URL, PROD_BASE_URL]
|
||||
base_url = urls[0]
|
||||
|
||||
logger.info("[antigravity] onboardUser: 开始激活账号, tier={}, endpoint={}", tier_id, base_url)
|
||||
|
||||
for attempt in range(1, max_attempts + 1):
|
||||
try:
|
||||
url = f"{base_url}/v1internal:onboardUser"
|
||||
resp = await execute_antigravity_rust_http_request(
|
||||
method="POST",
|
||||
url=url,
|
||||
headers=headers,
|
||||
body=body,
|
||||
proxy_config=proxy_config,
|
||||
request_id=f"antigravity:onboard-user:{tier_id}:{attempt}",
|
||||
provider_api_format="antigravity:onboard_user",
|
||||
timeout_seconds=timeout_seconds,
|
||||
content_type="application/json",
|
||||
)
|
||||
if resp is None:
|
||||
client = await HTTPClientPool.get_proxy_client(proxy_config)
|
||||
resp = await client.post(
|
||||
url,
|
||||
json=body,
|
||||
headers=headers,
|
||||
timeout=timeout_seconds,
|
||||
)
|
||||
if resp.status_code < 200 or resp.status_code >= 300:
|
||||
# 标记可用性
|
||||
if _should_fallback_status(resp.status_code):
|
||||
url_availability.mark_unavailable(base_url)
|
||||
raise RuntimeError(
|
||||
f"onboardUser failed: status={resp.status_code}, "
|
||||
f"body={resp.text[:200] if resp.text else ''}"
|
||||
)
|
||||
|
||||
url_availability.mark_success(base_url)
|
||||
data = resp.json()
|
||||
if not isinstance(data, dict):
|
||||
raise RuntimeError(f"onboardUser: unexpected response type: {type(data)}")
|
||||
|
||||
done = data.get("done")
|
||||
if done is True:
|
||||
# 从 response.cloudaicompanionProject 提取 project_id
|
||||
response_data = data.get("response")
|
||||
if isinstance(response_data, dict):
|
||||
project_id = extract_project_id(response_data)
|
||||
if project_id:
|
||||
logger.info(
|
||||
"[antigravity] onboardUser: 激活成功, project_id={}",
|
||||
project_id[:8] + "...",
|
||||
)
|
||||
return project_id
|
||||
|
||||
# done=true 但 project_id 为空(常见:cloudaicompanionProject: {})
|
||||
# 返回空串,由调用方 fallback 到随机 project_id
|
||||
logger.debug("[antigravity] onboardUser: done=true 但 project_id 为空")
|
||||
return ""
|
||||
|
||||
# done != true,继续轮询
|
||||
logger.debug(
|
||||
"[antigravity] onboardUser: 轮询 {}/{}, 等待完成...",
|
||||
attempt,
|
||||
max_attempts,
|
||||
)
|
||||
if attempt < max_attempts:
|
||||
await asyncio.sleep(poll_interval)
|
||||
|
||||
except RuntimeError:
|
||||
raise
|
||||
except Exception as e:
|
||||
# 临时网络错误:记录并继续轮询(而非直接终止)
|
||||
logger.warning(
|
||||
"[antigravity] onboardUser: 轮询 {}/{} 网络错误: {}, 继续重试...",
|
||||
attempt,
|
||||
max_attempts,
|
||||
e,
|
||||
)
|
||||
if attempt < max_attempts:
|
||||
await asyncio.sleep(poll_interval)
|
||||
continue
|
||||
raise RuntimeError(
|
||||
f"onboardUser request failed after {max_attempts} attempts: {e}"
|
||||
) from e
|
||||
|
||||
raise RuntimeError(f"onboardUser: 超时,已轮询 {max_attempts} 次仍未完成")
|
||||
|
||||
|
||||
async def fetch_available_models(
|
||||
access_token: str,
|
||||
*,
|
||||
project_id: str,
|
||||
proxy_config: dict[str, Any] | None = None,
|
||||
timeout_seconds: float = 10.0,
|
||||
) -> dict[str, Any]:
|
||||
"""调用 /v1internal:fetchAvailableModels 获取可用模型(包含配额信息)。
|
||||
|
||||
Antigravity 的此接口会返回类似:
|
||||
{"models": {"claude-sonnet-4": {"displayName": "...", "quotaInfo": {...}}, ...}}
|
||||
|
||||
对齐 AM:Sandbox 优先 + 正确的 fallback 逻辑。
|
||||
"""
|
||||
if not access_token:
|
||||
raise ValueError("missing access_token")
|
||||
if not project_id:
|
||||
raise ValueError("missing project_id")
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {access_token}",
|
||||
"Content-Type": "application/json",
|
||||
"Accept": "application/json",
|
||||
**get_v1internal_extra_headers(),
|
||||
}
|
||||
body = {"project": project_id}
|
||||
|
||||
# Sandbox 优先
|
||||
urls = url_availability.get_ordered_urls(prefer_daily=True)
|
||||
if not urls:
|
||||
urls = [SANDBOX_BASE_URL, DAILY_BASE_URL, PROD_BASE_URL]
|
||||
last_exc: Exception | None = None
|
||||
|
||||
for base_url in urls:
|
||||
try:
|
||||
url = f"{base_url}/v1internal:fetchAvailableModels"
|
||||
resp = await execute_antigravity_rust_http_request(
|
||||
method="POST",
|
||||
url=url,
|
||||
headers=headers,
|
||||
body=body,
|
||||
proxy_config=proxy_config,
|
||||
request_id=f"antigravity:fetch-available-models:{base_url}",
|
||||
provider_api_format="antigravity:fetch_available_models",
|
||||
timeout_seconds=timeout_seconds,
|
||||
content_type="application/json",
|
||||
)
|
||||
if resp is None:
|
||||
client = await HTTPClientPool.get_proxy_client(proxy_config)
|
||||
resp = await client.post(
|
||||
url,
|
||||
json=body,
|
||||
headers=headers,
|
||||
timeout=timeout_seconds,
|
||||
)
|
||||
if 200 <= resp.status_code < 300:
|
||||
url_availability.mark_success(base_url)
|
||||
data = resp.json()
|
||||
return data if isinstance(data, dict) else {}
|
||||
|
||||
# 可 fallback 的错误
|
||||
if _should_fallback_status(resp.status_code):
|
||||
url_availability.mark_unavailable(base_url)
|
||||
|
||||
# 429:尝试等待
|
||||
if resp.status_code == 429:
|
||||
delay = parse_retry_delay(resp.text)
|
||||
if delay and delay > 0:
|
||||
logger.debug(
|
||||
"[antigravity] fetchAvailableModels 429, waiting {:.1f}s",
|
||||
delay,
|
||||
)
|
||||
await asyncio.sleep(min(delay, 5.0))
|
||||
|
||||
last_exc = RuntimeError(
|
||||
f"fetchAvailableModels failed: status={resp.status_code} base_url={base_url}"
|
||||
)
|
||||
continue
|
||||
|
||||
# 403 Forbidden:账户权限被禁止(对齐 AM is_forbidden 标志)
|
||||
if resp.status_code == 403:
|
||||
reason = _extract_forbidden_reason(resp.text)
|
||||
logger.warning(
|
||||
"[antigravity] fetchAvailableModels 403 Forbidden: {}",
|
||||
reason or "unknown",
|
||||
)
|
||||
raise AntigravityAccountForbiddenException(
|
||||
message="账户访问被禁止",
|
||||
status_code=403,
|
||||
reason=reason,
|
||||
)
|
||||
|
||||
# 不可 fallback 的 4xx:直接抛出
|
||||
raise RuntimeError(
|
||||
f"fetchAvailableModels failed: status={resp.status_code} base_url={base_url} "
|
||||
f"body={resp.text[:200] if resp.text else ''}"
|
||||
)
|
||||
except (RuntimeError, AntigravityAccountForbiddenException):
|
||||
raise
|
||||
except Exception as e:
|
||||
url_availability.mark_unavailable(base_url)
|
||||
last_exc = e
|
||||
continue
|
||||
|
||||
raise last_exc or RuntimeError("fetchAvailableModels failed: all endpoints exhausted")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# User-Agent 版本号动态更新
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def refresh_user_agent(
|
||||
proxy_config: dict[str, Any] | None = None,
|
||||
*,
|
||||
timeout_seconds: float = 5.0,
|
||||
) -> str | None:
|
||||
"""从远程 API 获取最新 Antigravity 版本号并更新 User-Agent。
|
||||
|
||||
对齐 AM constants.rs:
|
||||
1. 尝试 VERSION_FETCH_URL
|
||||
2. Fallback 到 changelog 页面
|
||||
3. 最终 fallback 到 _FALLBACK_VERSION
|
||||
|
||||
Returns:
|
||||
更新后的版本号,或 None(如果获取失败但 fallback 到默认版本)。
|
||||
"""
|
||||
try:
|
||||
client = await HTTPClientPool.get_proxy_client(proxy_config)
|
||||
resp = await client.get(VERSION_FETCH_URL, timeout=timeout_seconds)
|
||||
if 200 <= resp.status_code < 300:
|
||||
version = parse_version_string(resp.text)
|
||||
if version:
|
||||
update_user_agent_version(version)
|
||||
logger.info("[antigravity] User-Agent 版本已更新: {}", version)
|
||||
return version
|
||||
except Exception as e:
|
||||
logger.debug("[antigravity] 获取远程版本失败: {}", e)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fallback Project ID 生成
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_ADJECTIVES = ("useful", "bright", "swift", "calm", "bold")
|
||||
_NOUNS = ("fuze", "wave", "spark", "flow", "core")
|
||||
|
||||
|
||||
def generate_fallback_project_id() -> str:
|
||||
"""生成随机 project_id(与 AM project_resolver.rs 的 generateProjectID 一致)。
|
||||
|
||||
格式: {adjective}-{noun}-{5位随机字符(base36)}
|
||||
Antigravity API 不严格校验 project 字段,当所有正常获取途径都失败时用作 fallback。
|
||||
"""
|
||||
adj = random.choice(_ADJECTIVES) # noqa: S311
|
||||
noun = random.choice(_NOUNS) # noqa: S311
|
||||
chars = "abcdefghijklmnopqrstuvwxyz0123456789"
|
||||
rand_part = "".join(random.choice(chars) for _ in range(5)) # noqa: S311
|
||||
return f"{adj}-{noun}-{rand_part}"
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AntigravityAccountForbiddenException",
|
||||
"extract_project_id",
|
||||
"extract_tier_id",
|
||||
"fetch_available_models",
|
||||
"generate_fallback_project_id",
|
||||
"load_code_assist",
|
||||
"onboard_user",
|
||||
"parse_retry_delay",
|
||||
"refresh_user_agent",
|
||||
]
|
||||
@@ -0,0 +1,279 @@
|
||||
"""Antigravity 全局常量定义。
|
||||
|
||||
注意:这里的 PROVIDER_TYPE 指的是 Provider.provider_type(用于路由与特判),
|
||||
不是 endpoint signature(family:kind)。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import platform
|
||||
import re
|
||||
import threading
|
||||
import uuid
|
||||
|
||||
# ============== API 端点 ==============
|
||||
# 唯一定义在 core 层,此处 re-export 保持向后兼容
|
||||
from src.core.provider_templates.fixed_providers import ANTIGRAVITY_PROD_URL as PROD_BASE_URL
|
||||
|
||||
DAILY_BASE_URL = "https://daily-cloudcode-pa.googleapis.com"
|
||||
SANDBOX_BASE_URL = "https://daily-cloudcode-pa.sandbox.googleapis.com"
|
||||
|
||||
# ============== User-Agent ==============
|
||||
VERSION_FETCH_URL = "https://antigravity-auto-updater-974169037036.us-central1.run.app"
|
||||
_FALLBACK_VERSION = "1.18.4"
|
||||
_FALLBACK_CHROME = "132.0.6834.160"
|
||||
_FALLBACK_ELECTRON = "39.2.3"
|
||||
_VERSION_RE = re.compile(r"\d+\.\d+\.\d+")
|
||||
|
||||
|
||||
def _detect_platform_info() -> str:
|
||||
"""检测当前运行平台,格式对齐 AM constants.rs 的 Electron UA。"""
|
||||
os_name = platform.system().lower()
|
||||
if os_name == "darwin":
|
||||
return "Macintosh; Intel Mac OS X 10_15_7"
|
||||
elif os_name == "windows":
|
||||
return "Windows NT 10.0; Win64; x64"
|
||||
else:
|
||||
return "X11; Linux x86_64"
|
||||
|
||||
|
||||
_PLATFORM_INFO = _detect_platform_info()
|
||||
|
||||
|
||||
def _build_antigravity_http_user_agent(
|
||||
*,
|
||||
platform_token: str,
|
||||
version: str,
|
||||
chrome_version: str,
|
||||
electron_version: str,
|
||||
) -> str:
|
||||
return (
|
||||
f"Mozilla/5.0 ({platform_token}) AppleWebKit/537.36 (KHTML, like Gecko) "
|
||||
f"Antigravity/{version} Chrome/{chrome_version} "
|
||||
f"Electron/{electron_version} Safari/537.36"
|
||||
)
|
||||
|
||||
|
||||
# HTTP Header User-Agent(对齐 AM constants.rs: 完整 Electron 浏览器格式)
|
||||
HTTP_USER_AGENT = _build_antigravity_http_user_agent(
|
||||
platform_token=_PLATFORM_INFO,
|
||||
version=_FALLBACK_VERSION,
|
||||
chrome_version=_FALLBACK_CHROME,
|
||||
electron_version=_FALLBACK_ELECTRON,
|
||||
)
|
||||
|
||||
# V1InternalRequest.userAgent 字段(固定值)
|
||||
REQUEST_USER_AGENT = "antigravity"
|
||||
|
||||
# --- 动态 User-Agent 支持 ---
|
||||
_ua_lock = threading.Lock()
|
||||
_ua_version: str = _FALLBACK_VERSION
|
||||
|
||||
|
||||
def get_http_user_agent() -> str:
|
||||
"""返回当前 HTTP User-Agent 字符串(对齐 AM Electron UA 格式)。"""
|
||||
with _ua_lock:
|
||||
return _build_antigravity_http_user_agent(
|
||||
platform_token=_PLATFORM_INFO,
|
||||
version=_ua_version,
|
||||
chrome_version=_FALLBACK_CHROME,
|
||||
electron_version=_FALLBACK_ELECTRON,
|
||||
)
|
||||
|
||||
|
||||
def update_user_agent_version(version: str) -> None:
|
||||
"""更新 User-Agent 中的版本号(由 refresh_user_agent 调用)。"""
|
||||
global HTTP_USER_AGENT, _ua_version # noqa: PLW0603
|
||||
version = str(version or "").strip()
|
||||
if not version:
|
||||
return
|
||||
with _ua_lock:
|
||||
_ua_version = version
|
||||
HTTP_USER_AGENT = _build_antigravity_http_user_agent(
|
||||
platform_token=_PLATFORM_INFO,
|
||||
version=_ua_version,
|
||||
chrome_version=_FALLBACK_CHROME,
|
||||
electron_version=_FALLBACK_ELECTRON,
|
||||
)
|
||||
|
||||
|
||||
def parse_version_string(text: str) -> str | None:
|
||||
"""从任意文本中提取 X.Y.Z 格式的版本号。"""
|
||||
m = _VERSION_RE.search(text)
|
||||
return m.group(0) if m else None
|
||||
|
||||
|
||||
# ============== URL 可用性 ==============
|
||||
URL_UNAVAILABLE_TTL_SECONDS = 300 # 5 分钟
|
||||
|
||||
# ============== AM Client Identity Headers ==============
|
||||
# 对齐 AM upstream/client.rs: 伪装为官方 Antigravity 客户端
|
||||
# 缺少这些 header 会导致新模型(如 gemini-3.1-pro-preview)返回 404
|
||||
_SESSION_ID = uuid.uuid4().hex # 每次进程启动生成一个固定 session ID
|
||||
|
||||
|
||||
def get_v1internal_extra_headers() -> dict[str, str]:
|
||||
"""构建 v1internal 请求需要的额外 header(对齐 AM upstream/client.rs)。"""
|
||||
with _ua_lock:
|
||||
version = _ua_version
|
||||
|
||||
return {
|
||||
"User-Agent": get_http_user_agent(),
|
||||
"x-client-name": "antigravity",
|
||||
"x-client-version": version,
|
||||
"x-vscode-sessionid": _SESSION_ID,
|
||||
"x-goog-api-client": "gl-node/18.18.2 fire/0.8.6 grpc/1.10.x",
|
||||
}
|
||||
|
||||
|
||||
# ============== Thinking Signature ==============
|
||||
# 统一从 core 层导入,避免多处定义
|
||||
from src.core.api_format.conversion.constants import DUMMY_THOUGHT_SIGNATURE # noqa: E402
|
||||
from src.core.api_format.conversion.thinking_cache import MIN_SIGNATURE_LENGTH # noqa: E402, F401
|
||||
|
||||
# ============== Thinking Budget ==============
|
||||
THINKING_BUDGET_AUTO_CAP = 24576
|
||||
THINKING_BUDGET_DEFAULT_INJECT = 24576 # 对齐 AM wrapper.rs (was 16000)
|
||||
# 给输出留的空间(对齐 Antigravity-Manager:普通模型 8192,图像模型 2048)
|
||||
OUTPUT_OVERHEAD = 8192
|
||||
OUTPUT_OVERHEAD_IMAGE = 2048
|
||||
# 模型最大输出限制(防止超限)
|
||||
MODEL_MAX_OUTPUT_LIMIT = 65536
|
||||
# 包含这些关键字的模型会自动注入 thinkingConfig(如果缺失)
|
||||
THINKING_MODELS_AUTO_INJECT_KEYWORDS = (
|
||||
"thinking",
|
||||
"gemini-2.0-pro",
|
||||
"gemini-3-pro",
|
||||
"gemini-3.1-pro",
|
||||
)
|
||||
|
||||
# ============== Google Search (Grounding) ==============
|
||||
# 对齐 AM common_utils.rs: 仅 gemini-2.5-flash 支持 googleSearch tool
|
||||
WEB_SEARCH_MODEL = "gemini-2.5-flash"
|
||||
# 联网工具检测关键字(对齐 AM detects_networking_tool)
|
||||
NETWORKING_TOOL_KEYWORDS = frozenset(
|
||||
{
|
||||
"web_search",
|
||||
"google_search",
|
||||
"web_search_20250305",
|
||||
"google_search_retrieval",
|
||||
}
|
||||
)
|
||||
|
||||
# ============== Image Generation ==============
|
||||
# 上游图像生成模型的固定名称
|
||||
IMAGE_GEN_UPSTREAM_MODEL = "gemini-3-pro-image"
|
||||
# 模型后缀 → 宽高比映射
|
||||
IMAGE_ASPECT_RATIO_SUFFIXES: dict[str, str] = {
|
||||
"-21x9": "21:9",
|
||||
"-21-9": "21:9",
|
||||
"-16x9": "16:9",
|
||||
"-16-9": "16:9",
|
||||
"-9x16": "9:16",
|
||||
"-9-16": "9:16",
|
||||
"-4x3": "4:3",
|
||||
"-4-3": "4:3",
|
||||
"-3x4": "3:4",
|
||||
"-3-4": "3:4",
|
||||
"-3x2": "3:2",
|
||||
"-3-2": "3:2",
|
||||
"-2x3": "2:3",
|
||||
"-2-3": "2:3",
|
||||
"-5x4": "5:4",
|
||||
"-5-4": "5:4",
|
||||
"-4x5": "4:5",
|
||||
"-4-5": "4:5",
|
||||
"-1x1": "1:1",
|
||||
"-1-1": "1:1",
|
||||
}
|
||||
# 标准宽高比字符串(用于直接匹配 size 参数)
|
||||
STANDARD_ASPECT_RATIOS = frozenset(
|
||||
{
|
||||
"21:9",
|
||||
"16:9",
|
||||
"9:16",
|
||||
"4:3",
|
||||
"3:4",
|
||||
"3:2",
|
||||
"2:3",
|
||||
"5:4",
|
||||
"4:5",
|
||||
"1:1",
|
||||
}
|
||||
)
|
||||
# 宽高比容差匹配表:(ratio, label)
|
||||
ASPECT_RATIO_TABLE: tuple[tuple[float, str], ...] = (
|
||||
(21.0 / 9.0, "21:9"),
|
||||
(16.0 / 9.0, "16:9"),
|
||||
(4.0 / 3.0, "4:3"),
|
||||
(3.0 / 4.0, "3:4"),
|
||||
(9.0 / 16.0, "9:16"),
|
||||
(3.0 / 2.0, "3:2"),
|
||||
(2.0 / 3.0, "2:3"),
|
||||
(5.0 / 4.0, "5:4"),
|
||||
(4.0 / 5.0, "4:5"),
|
||||
(1.0, "1:1"),
|
||||
)
|
||||
|
||||
# ============== Retry ==============
|
||||
RETRY_429_BASE_SECONDS = 5.0
|
||||
RETRY_503_BASE_SECONDS = 10.0
|
||||
RETRY_503_MAX_SECONDS = 60.0
|
||||
RETRY_500_BASE_SECONDS = 3.0
|
||||
|
||||
# ============== v1internal 路径 ==============
|
||||
V1INTERNAL_PATH_TEMPLATE = "/v1internal:{action}"
|
||||
|
||||
# ============== Signature 错误关键字(用于 400 错误检测) ==============
|
||||
SIGNATURE_ERROR_KEYWORDS = (
|
||||
"Invalid `signature`",
|
||||
"thinking.signature",
|
||||
"thinking.thinking",
|
||||
"Corrupted thought signature",
|
||||
)
|
||||
|
||||
# ============== Antigravity System Instruction ==============
|
||||
ANTIGRAVITY_SYSTEM_INSTRUCTION = (
|
||||
"You are Antigravity, a powerful agentic AI coding assistant designed by the "
|
||||
"Google Deepmind team working on Advanced Agentic Coding.\n"
|
||||
"You are pair programming with a USER to solve their coding task. The task may "
|
||||
"require creating a new codebase, modifying or debugging an existing codebase, "
|
||||
"or simply answering a question.\n"
|
||||
"**Absolute paths only**\n"
|
||||
"**Proactiveness**"
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"ANTIGRAVITY_SYSTEM_INSTRUCTION",
|
||||
"ASPECT_RATIO_TABLE",
|
||||
"DAILY_BASE_URL",
|
||||
"DUMMY_THOUGHT_SIGNATURE",
|
||||
"HTTP_USER_AGENT",
|
||||
"IMAGE_ASPECT_RATIO_SUFFIXES",
|
||||
"IMAGE_GEN_UPSTREAM_MODEL",
|
||||
"MIN_SIGNATURE_LENGTH",
|
||||
"MODEL_MAX_OUTPUT_LIMIT",
|
||||
"NETWORKING_TOOL_KEYWORDS",
|
||||
"OUTPUT_OVERHEAD",
|
||||
"OUTPUT_OVERHEAD_IMAGE",
|
||||
"PROD_BASE_URL",
|
||||
"REQUEST_USER_AGENT",
|
||||
"RETRY_429_BASE_SECONDS",
|
||||
"RETRY_500_BASE_SECONDS",
|
||||
"RETRY_503_BASE_SECONDS",
|
||||
"RETRY_503_MAX_SECONDS",
|
||||
"SANDBOX_BASE_URL",
|
||||
"SIGNATURE_ERROR_KEYWORDS",
|
||||
"STANDARD_ASPECT_RATIOS",
|
||||
"THINKING_BUDGET_AUTO_CAP",
|
||||
"THINKING_BUDGET_DEFAULT_INJECT",
|
||||
"THINKING_MODELS_AUTO_INJECT_KEYWORDS",
|
||||
"URL_UNAVAILABLE_TTL_SECONDS",
|
||||
"V1INTERNAL_PATH_TEMPLATE",
|
||||
"VERSION_FETCH_URL",
|
||||
"WEB_SEARCH_MODEL",
|
||||
"get_http_user_agent",
|
||||
"get_v1internal_extra_headers",
|
||||
"parse_version_string",
|
||||
"update_user_agent_version",
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,388 @@
|
||||
"""Antigravity provider plugin — 统一注册入口。
|
||||
|
||||
将 Antigravity 对各通用 registry / capability registry 的注册集中在一个文件中:
|
||||
- Envelope (v1internal 信封)
|
||||
- Transport Hook (URL 构建)
|
||||
- Auth Enricher (OAuth enrichment)
|
||||
- Model Fetcher (模型获取)
|
||||
- Provider Format Capability(跨格式变体)
|
||||
|
||||
新增 provider 时参照此文件创建对应的 plugin.py 即可。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from typing import Any
|
||||
from urllib.parse import urlencode
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.services.provider.adapters.antigravity.constants import V1INTERNAL_PATH_TEMPLATE
|
||||
from src.services.provider.adapters.antigravity.url_availability import url_availability
|
||||
from src.services.provider.request_context import set_selected_base_url
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Transport Hook
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def build_antigravity_url(
|
||||
endpoint: Any,
|
||||
*,
|
||||
is_stream: bool,
|
||||
effective_query_params: dict[str, Any],
|
||||
**_kwargs: Any,
|
||||
) -> str:
|
||||
"""构建 Antigravity v1internal URL。
|
||||
|
||||
使用 url_availability 选择最优端点,构建 v1internal:generateContent 或
|
||||
v1internal:streamGenerateContent URL。
|
||||
"""
|
||||
ordered_urls = url_availability.get_ordered_urls(prefer_daily=True)
|
||||
base_url = ordered_urls[0] if ordered_urls else endpoint.base_url
|
||||
|
||||
# 存入 contextvars(供后续 Handler 层 envelope 获取)
|
||||
set_selected_base_url(str(base_url) if base_url is not None else None)
|
||||
|
||||
action = "streamGenerateContent" if is_stream else "generateContent"
|
||||
path = V1INTERNAL_PATH_TEMPLATE.format(action=action)
|
||||
|
||||
query_params = dict(effective_query_params)
|
||||
|
||||
# v1internal 流式请求同样支持 ?alt=sse
|
||||
if is_stream:
|
||||
query_params.setdefault("alt", "sse")
|
||||
|
||||
# 移除 v1internal 不支持的查询参数
|
||||
query_params.pop("beta", None)
|
||||
|
||||
url = f"{str(base_url).rstrip('/')}{path}"
|
||||
if query_params:
|
||||
query_string = urlencode(query_params, doseq=True)
|
||||
if query_string:
|
||||
url = f"{url}?{query_string}"
|
||||
|
||||
return url
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Auth Enricher
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def enrich_antigravity(
|
||||
auth_config: dict[str, Any],
|
||||
token_response: dict[str, Any],
|
||||
access_token: str,
|
||||
proxy_config: dict[str, Any] | None,
|
||||
) -> dict[str, Any]:
|
||||
"""Antigravity auth_config enrichment.
|
||||
|
||||
1. 通过 Google userinfo API 获取 email
|
||||
2. 通过 loadCodeAssist 获取 project_id / tier
|
||||
3. 未激活账号尝试 onboardUser
|
||||
4. fallback 到随机 project_id
|
||||
"""
|
||||
from src.core.provider_oauth_utils import fetch_google_email
|
||||
from src.services.provider.adapters.antigravity.client import (
|
||||
extract_project_id,
|
||||
extract_tier_id,
|
||||
generate_fallback_project_id,
|
||||
load_code_assist,
|
||||
onboard_user,
|
||||
)
|
||||
|
||||
# Email(仅在缺失时获取)
|
||||
if not auth_config.get("email"):
|
||||
email = await fetch_google_email(
|
||||
access_token,
|
||||
proxy_config=proxy_config,
|
||||
timeout_seconds=10.0,
|
||||
)
|
||||
if email:
|
||||
auth_config["email"] = email
|
||||
|
||||
# Project ID + Tier(需要 loadCodeAssist 时一起获取)
|
||||
need_project = not auth_config.get("project_id")
|
||||
# Tier 每次 enrich 都重新获取(确保归一化为 Free/Pro/Ultra)
|
||||
need_tier = True
|
||||
|
||||
if need_project or need_tier:
|
||||
project_id = auth_config.get("project_id", "")
|
||||
try:
|
||||
code_assist = await load_code_assist(access_token, proxy_config=proxy_config)
|
||||
|
||||
# 提取 tier 信息(对齐 sub2api:优先 paidTier,fallback currentTier)
|
||||
tier_str = _extract_tier_from_code_assist(code_assist)
|
||||
if tier_str:
|
||||
auth_config["tier"] = tier_str
|
||||
logger.info("[enrich] Antigravity tier: {}", tier_str)
|
||||
|
||||
if need_project:
|
||||
project_id = extract_project_id(code_assist)
|
||||
|
||||
# 未激活:尝试 onboardUser
|
||||
if not project_id and code_assist.get("allowedTiers"):
|
||||
tier_id = extract_tier_id(code_assist)
|
||||
logger.info("[enrich] Antigravity onboardUser tier={}", tier_id)
|
||||
project_id = await onboard_user(
|
||||
access_token,
|
||||
tier_id=tier_id,
|
||||
proxy_config=proxy_config,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("[enrich] Antigravity loadCodeAssist/onboardUser 失败: {}", e)
|
||||
|
||||
if need_project:
|
||||
if project_id:
|
||||
auth_config["project_id"] = project_id
|
||||
logger.info("[enrich] Antigravity project_id: {}", project_id[:8] + "...")
|
||||
else:
|
||||
fallback = generate_fallback_project_id()
|
||||
auth_config["project_id"] = fallback
|
||||
logger.info("[enrich] Antigravity 随机 project_id fallback: {}", fallback)
|
||||
|
||||
return auth_config
|
||||
|
||||
|
||||
def _extract_tier_from_code_assist(code_assist: dict[str, Any]) -> str:
|
||||
"""从 loadCodeAssist 响应中提取用户层级,返回 Free/Pro/Ultra。
|
||||
|
||||
对齐 sub2api/CLIProxyAPI:优先 paidTier,fallback currentTier。
|
||||
兼容 tier 为字符串或 {"id": "...", "tierType": "..."} 两种格式。
|
||||
"""
|
||||
# 优先 paidTier(付费订阅级别)
|
||||
paid_tier = code_assist.get("paidTier")
|
||||
tier = _normalize_tier(_extract_tier_raw(paid_tier))
|
||||
if tier:
|
||||
return tier
|
||||
|
||||
# fallback currentTier
|
||||
current_tier = code_assist.get("currentTier")
|
||||
tier = _normalize_tier(_extract_tier_raw(current_tier))
|
||||
if tier:
|
||||
return tier
|
||||
|
||||
return "Free"
|
||||
|
||||
|
||||
def _extract_tier_raw(tier_obj: Any) -> str:
|
||||
"""从 tier 对象中提取原始标识,兼容字符串和字典两种格式。"""
|
||||
if isinstance(tier_obj, str) and tier_obj.strip():
|
||||
return tier_obj.strip()
|
||||
if isinstance(tier_obj, dict):
|
||||
for key in ("id", "tierType"):
|
||||
val = tier_obj.get(key)
|
||||
if isinstance(val, str) and val.strip():
|
||||
return val.strip()
|
||||
return ""
|
||||
|
||||
|
||||
def _normalize_tier(raw: str) -> str:
|
||||
"""将上游 tier 标识归一化为 Free/Pro/Ultra。
|
||||
|
||||
上游格式示例:
|
||||
- id: "free-tier", "g1-pro-tier", "g1-ultra-tier"
|
||||
- tierType: "FREE", "PAID"
|
||||
"""
|
||||
if not raw:
|
||||
return ""
|
||||
lower = raw.lower()
|
||||
if "ultra" in lower:
|
||||
return "Ultra"
|
||||
if "pro" in lower or "paid" in lower:
|
||||
return "Pro"
|
||||
if "free" in lower or "legacy" in lower:
|
||||
return "Free"
|
||||
return raw
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Model Fetcher
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Antigravity 内部测试/调试模型,不应暴露给用户
|
||||
BLOCKED_MODELS: frozenset[str] = frozenset({"chat_23310", "chat_20706"})
|
||||
|
||||
|
||||
async def fetch_models_antigravity(
|
||||
ctx: Any,
|
||||
timeout_seconds: float,
|
||||
) -> tuple[list[dict], list[str], bool, dict[str, Any] | None]:
|
||||
"""Antigravity 模型获取策略。
|
||||
|
||||
调用 v1internal:fetchAvailableModels 获取可用模型列表,
|
||||
解析配额信息,过滤黑名单模型。
|
||||
"""
|
||||
from src.services.provider.adapters.antigravity.client import (
|
||||
AntigravityAccountForbiddenException,
|
||||
fetch_available_models,
|
||||
)
|
||||
|
||||
auth_config = ctx.auth_config or {}
|
||||
project_id = auth_config.get("project_id")
|
||||
if not isinstance(project_id, str) or not project_id.strip():
|
||||
return [], ["antigravity: missing auth_config.project_id (please re-auth)"], False, None
|
||||
|
||||
try:
|
||||
data = await fetch_available_models(
|
||||
ctx.api_key_value,
|
||||
project_id=project_id.strip(),
|
||||
proxy_config=ctx.proxy_config,
|
||||
timeout_seconds=timeout_seconds,
|
||||
)
|
||||
except AntigravityAccountForbiddenException:
|
||||
# 重新抛出,由上层(keys.py)处理自动停用逻辑
|
||||
raise
|
||||
except Exception as e:
|
||||
return [], [f"antigravity: fetchAvailableModels error: {e}"], False, None
|
||||
|
||||
raw_models = data.get("models")
|
||||
if not isinstance(raw_models, dict):
|
||||
return [], ["antigravity: invalid response (missing models)"], False, None
|
||||
|
||||
models: list[dict] = []
|
||||
quota_by_model: dict[str, dict[str, Any]] = {}
|
||||
|
||||
for model_id, model_data in raw_models.items():
|
||||
if not isinstance(model_id, str) or not model_id.strip():
|
||||
continue
|
||||
if model_id.strip() in BLOCKED_MODELS:
|
||||
continue
|
||||
if not isinstance(model_data, dict):
|
||||
model_data = {}
|
||||
|
||||
display_name = model_data.get("displayName")
|
||||
if not isinstance(display_name, str) or not display_name:
|
||||
display_name = model_id
|
||||
|
||||
models.append(
|
||||
{
|
||||
"id": model_id,
|
||||
"owned_by": "antigravity",
|
||||
"display_name": display_name,
|
||||
"api_format": "gemini:chat",
|
||||
}
|
||||
)
|
||||
|
||||
quota_info = model_data.get("quotaInfo")
|
||||
if not isinstance(quota_info, dict):
|
||||
# 没有 quotaInfo 视为配额耗尽
|
||||
quota_by_model[model_id] = {
|
||||
"remaining_fraction": 0.0,
|
||||
"used_percent": 100.0,
|
||||
}
|
||||
continue
|
||||
|
||||
remaining = quota_info.get("remainingFraction")
|
||||
reset_time = quota_info.get("resetTime")
|
||||
|
||||
remaining_fraction: float | None = None
|
||||
try:
|
||||
if remaining is not None:
|
||||
remaining_fraction = float(remaining)
|
||||
except Exception:
|
||||
remaining_fraction = None
|
||||
|
||||
if remaining_fraction is None:
|
||||
# remainingFraction 缺失视为配额耗尽
|
||||
payload: dict[str, Any] = {
|
||||
"remaining_fraction": 0.0,
|
||||
"used_percent": 100.0,
|
||||
}
|
||||
if isinstance(reset_time, str) and reset_time.strip():
|
||||
payload["reset_time"] = reset_time.strip()
|
||||
quota_by_model[model_id] = payload
|
||||
continue
|
||||
|
||||
used_percent = (1.0 - remaining_fraction) * 100.0
|
||||
if used_percent < 0:
|
||||
used_percent = 0.0
|
||||
if used_percent > 100: # noqa: PLR2004
|
||||
used_percent = 100.0
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"remaining_fraction": remaining_fraction,
|
||||
"used_percent": used_percent,
|
||||
}
|
||||
if isinstance(reset_time, str) and reset_time.strip():
|
||||
payload["reset_time"] = reset_time.strip()
|
||||
quota_by_model[model_id] = payload
|
||||
|
||||
upstream_metadata: dict[str, Any] | None = None
|
||||
if quota_by_model:
|
||||
upstream_metadata = {
|
||||
"antigravity": {
|
||||
"updated_at": int(time.time()),
|
||||
"quota_by_model": quota_by_model,
|
||||
}
|
||||
}
|
||||
|
||||
return models, [], True, upstream_metadata
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Export builder
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_AG_SKIP_KEYS = frozenset(
|
||||
{
|
||||
"access_token",
|
||||
"expires_at",
|
||||
"updated_at",
|
||||
"token_type",
|
||||
"scope",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def antigravity_export_builder(
|
||||
auth_config: dict[str, Any],
|
||||
upstream_metadata: dict[str, Any] | None,
|
||||
) -> dict[str, Any]:
|
||||
"""Antigravity 导出:保留 refresh_token / email / project_id / tier。"""
|
||||
return {
|
||||
k: v for k, v in auth_config.items() if k not in _AG_SKIP_KEYS and v is not None and v != ""
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unified Registration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def register_all() -> None:
|
||||
"""一次性注册 Antigravity 的所有 hooks 到各通用 registry。"""
|
||||
from src.core.api_format.capabilities import register_provider_behavior_variant
|
||||
from src.core.provider_oauth_utils import register_auth_enricher
|
||||
from src.services.model.upstream_fetcher import UpstreamModelsFetcherRegistry
|
||||
from src.services.provider.adapters.antigravity.envelope import antigravity_v1internal_envelope
|
||||
from src.services.provider.envelope import register_envelope
|
||||
from src.services.provider.export import register_export_builder
|
||||
from src.services.provider.transport import register_transport_hook
|
||||
|
||||
# Envelope
|
||||
register_envelope("antigravity", "gemini:chat", antigravity_v1internal_envelope)
|
||||
# Backward compat: allow existing endpoints that still use the old signature.
|
||||
register_envelope("antigravity", "gemini:cli", antigravity_v1internal_envelope)
|
||||
register_envelope("antigravity", "", antigravity_v1internal_envelope)
|
||||
|
||||
# Transport
|
||||
register_transport_hook("antigravity", "gemini:chat", build_antigravity_url)
|
||||
# Backward compat: allow existing endpoints that still use the old signature.
|
||||
register_transport_hook("antigravity", "gemini:cli", build_antigravity_url)
|
||||
|
||||
# Auth
|
||||
register_auth_enricher("antigravity", enrich_antigravity)
|
||||
|
||||
# Export
|
||||
register_export_builder("antigravity", antigravity_export_builder)
|
||||
|
||||
# Model Fetcher
|
||||
UpstreamModelsFetcherRegistry.register(
|
||||
provider_types=["antigravity"],
|
||||
fetcher=fetch_models_antigravity,
|
||||
)
|
||||
|
||||
# Provider Format Capability
|
||||
register_provider_behavior_variant("antigravity", cross_format=True)
|
||||
@@ -0,0 +1,108 @@
|
||||
"""Shared Rust executor HTTP helper for Antigravity side calls."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from src.config.settings import config
|
||||
from src.core.logger import logger
|
||||
|
||||
|
||||
async def execute_antigravity_rust_http_request(
|
||||
*,
|
||||
method: str,
|
||||
url: str,
|
||||
headers: dict[str, str],
|
||||
body: Any,
|
||||
proxy_config: dict[str, Any] | None,
|
||||
request_id: str,
|
||||
provider_api_format: str,
|
||||
timeout_seconds: float,
|
||||
content_type: str | None = None,
|
||||
) -> httpx.Response | None:
|
||||
from src.services.request.execution_runtime_plan import (
|
||||
ExecutionPlan,
|
||||
ExecutionPlanBody,
|
||||
ExecutionPlanTimeouts,
|
||||
build_execution_plan_body,
|
||||
build_proxy_snapshot,
|
||||
)
|
||||
from src.services.request.execution_runtime_client import (
|
||||
ExecutionRuntimeClient,
|
||||
ExecutionRuntimeClientError,
|
||||
)
|
||||
|
||||
if config.execution_runtime_backend != "rust":
|
||||
return None
|
||||
|
||||
final_headers = dict(headers)
|
||||
if (
|
||||
body is not None
|
||||
and content_type
|
||||
and not any(str(key).lower() == "content-type" for key in final_headers)
|
||||
):
|
||||
final_headers["content-type"] = content_type
|
||||
|
||||
timeout_ms = max(int(timeout_seconds * 1000), 1_000)
|
||||
|
||||
try:
|
||||
proxy_snapshot = await build_proxy_snapshot(proxy_config, label="Antigravity")
|
||||
result = await ExecutionRuntimeClient().execute_sync_json(
|
||||
ExecutionPlan(
|
||||
request_id=request_id,
|
||||
candidate_id=None,
|
||||
provider_name="antigravity",
|
||||
provider_id="",
|
||||
endpoint_id="",
|
||||
key_id="",
|
||||
method=method,
|
||||
url=url,
|
||||
headers=final_headers,
|
||||
body=(
|
||||
build_execution_plan_body(body, content_type=content_type)
|
||||
if body is not None
|
||||
else ExecutionPlanBody()
|
||||
),
|
||||
stream=False,
|
||||
provider_api_format=provider_api_format,
|
||||
client_api_format=provider_api_format,
|
||||
model_name="antigravity",
|
||||
content_type=content_type,
|
||||
proxy=proxy_snapshot,
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=timeout_ms,
|
||||
read_ms=timeout_ms,
|
||||
write_ms=timeout_ms,
|
||||
pool_ms=timeout_ms,
|
||||
total_ms=timeout_ms,
|
||||
),
|
||||
)
|
||||
)
|
||||
except (ExecutionRuntimeClientError, httpx.HTTPError, json.JSONDecodeError) as exc:
|
||||
logger.warning("Antigravity Rust HTTP fallback {} {}: {}", method, url, exc)
|
||||
return None
|
||||
except Exception as exc:
|
||||
logger.warning("Antigravity Rust HTTP unexpected fallback {} {}: {}", method, 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(method, url, headers=final_headers),
|
||||
headers=response_headers,
|
||||
content=response_body,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["execute_antigravity_rust_http_request"]
|
||||
@@ -0,0 +1,13 @@
|
||||
"""Backward-compatible re-export for Antigravity thinking signature cache.
|
||||
|
||||
The implementation moved to `src.core.api_format.conversion.thinking_cache` to eliminate
|
||||
core → services reverse dependencies.
|
||||
"""
|
||||
|
||||
from src.core.api_format.conversion.thinking_cache import (
|
||||
MIN_SIGNATURE_LENGTH,
|
||||
ThinkingSignatureCache,
|
||||
signature_cache,
|
||||
)
|
||||
|
||||
__all__ = ["MIN_SIGNATURE_LENGTH", "ThinkingSignatureCache", "signature_cache"]
|
||||
@@ -0,0 +1,82 @@
|
||||
"""Antigravity URL 可用性管理(带 TTL 自动恢复)。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
|
||||
from src.services.provider.adapters.antigravity.constants import (
|
||||
DAILY_BASE_URL,
|
||||
PROD_BASE_URL,
|
||||
SANDBOX_BASE_URL,
|
||||
URL_UNAVAILABLE_TTL_SECONDS,
|
||||
)
|
||||
|
||||
|
||||
class URLAvailability:
|
||||
"""管理 Antigravity API 端点可用性(进程内)。"""
|
||||
|
||||
_instance: "URLAvailability | None" = None
|
||||
_lock = threading.Lock()
|
||||
|
||||
def __new__(cls) -> "URLAvailability":
|
||||
if cls._instance is None:
|
||||
with cls._lock:
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance._init()
|
||||
return cls._instance
|
||||
|
||||
def _init(self) -> None:
|
||||
self._unavailable: dict[str, float] = {} # url -> recover_at(ts)
|
||||
self._last_success: str | None = None
|
||||
self._mu = threading.RLock()
|
||||
|
||||
def _prune(self, now: float | None = None) -> None:
|
||||
now_ts = time.time() if now is None else now
|
||||
self._unavailable = {u: t for u, t in self._unavailable.items() if t > now_ts}
|
||||
|
||||
def is_available(self, url: str) -> bool:
|
||||
with self._mu:
|
||||
self._prune()
|
||||
return url not in self._unavailable
|
||||
|
||||
def get_ordered_urls(self, *, prefer_daily: bool = True) -> list[str]:
|
||||
"""返回优先级排序的可用 URL 列表。
|
||||
|
||||
- 默认 daily 优先(通常限流更宽松)
|
||||
- 最近成功的 URL 会被提升到最前
|
||||
- 若全部被标记不可用,则返回 base_order(允许继续尝试,等待 TTL 自动恢复)
|
||||
"""
|
||||
with self._mu:
|
||||
self._prune()
|
||||
|
||||
# Antigravity-Manager 顺序:sandbox → daily → prod
|
||||
base_order = (
|
||||
[SANDBOX_BASE_URL, DAILY_BASE_URL, PROD_BASE_URL]
|
||||
if prefer_daily
|
||||
else [PROD_BASE_URL, SANDBOX_BASE_URL, DAILY_BASE_URL]
|
||||
)
|
||||
|
||||
if self._last_success and self._last_success in base_order:
|
||||
base_order.remove(self._last_success)
|
||||
base_order.insert(0, self._last_success)
|
||||
|
||||
available = [u for u in base_order if u not in self._unavailable]
|
||||
return available if available else base_order
|
||||
|
||||
def mark_success(self, url: str) -> None:
|
||||
with self._mu:
|
||||
self._last_success = url
|
||||
self._unavailable.pop(url, None)
|
||||
|
||||
def mark_unavailable(self, url: str) -> None:
|
||||
with self._mu:
|
||||
self._unavailable[url] = time.time() + URL_UNAVAILABLE_TTL_SECONDS
|
||||
if self._last_success == url:
|
||||
self._last_success = None
|
||||
|
||||
|
||||
url_availability = URLAvailability()
|
||||
|
||||
__all__ = ["URLAvailability", "url_availability"]
|
||||
@@ -0,0 +1,3 @@
|
||||
"""Claude Code provider adapter."""
|
||||
|
||||
__all__ = []
|
||||
@@ -0,0 +1,96 @@
|
||||
"""Claude Code CLI client restriction.
|
||||
|
||||
When cli_only_enabled is True, only requests from genuine Claude Code CLI
|
||||
clients are allowed. Non-CLI traffic receives a 403 response.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextvars
|
||||
from typing import Any
|
||||
|
||||
from src.core.logger import logger
|
||||
|
||||
# Contextvar to carry the original request headers into the envelope layer.
|
||||
_original_request_headers: contextvars.ContextVar[dict[str, str] | None] = contextvars.ContextVar(
|
||||
"claude_code_original_request_headers",
|
||||
default=None,
|
||||
)
|
||||
|
||||
|
||||
def set_original_request_headers(headers: dict[str, str] | None) -> None:
|
||||
_original_request_headers.set(headers)
|
||||
|
||||
|
||||
def get_original_request_headers() -> dict[str, str] | None:
|
||||
return _original_request_headers.get()
|
||||
|
||||
|
||||
# Known Claude Code CLI User-Agent patterns.
|
||||
_CLI_USER_AGENT_PATTERNS = (
|
||||
"claude-code",
|
||||
"claudecode",
|
||||
"claude_code",
|
||||
)
|
||||
|
||||
# Known originator / x-app values indicating CLI usage.
|
||||
_CLI_APP_VALUES = {"cli"}
|
||||
|
||||
|
||||
def is_claude_code_client(headers: dict[str, Any]) -> bool:
|
||||
"""Detect whether the request originates from a Claude Code CLI client.
|
||||
|
||||
Detection signals (any match is sufficient):
|
||||
1. User-Agent contains a known Claude Code CLI pattern
|
||||
2. x-app header equals "cli"
|
||||
"""
|
||||
# Normalize header keys to lowercase for case-insensitive matching.
|
||||
lower_headers = {k.lower(): v for k, v in headers.items()}
|
||||
|
||||
# Check User-Agent
|
||||
ua = str(lower_headers.get("user-agent", "")).lower()
|
||||
for pattern in _CLI_USER_AGENT_PATTERNS:
|
||||
if pattern in ua:
|
||||
return True
|
||||
|
||||
# Check x-app header
|
||||
x_app = str(lower_headers.get("x-app", "")).strip().lower()
|
||||
if x_app in _CLI_APP_VALUES:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def enforce_cli_only(cli_only_enabled: bool) -> None:
|
||||
"""Enforce CLI-only restriction if enabled.
|
||||
|
||||
Reads original request headers from contextvar, checks whether the
|
||||
client is a Claude Code CLI, and raises HTTPException(403) if not.
|
||||
"""
|
||||
if not cli_only_enabled:
|
||||
return
|
||||
|
||||
headers = get_original_request_headers()
|
||||
if headers is None:
|
||||
# No headers available; skip enforcement (should not happen in normal flow).
|
||||
logger.debug("CLI-only check skipped: no request headers in context")
|
||||
return
|
||||
|
||||
if is_claude_code_client(headers):
|
||||
return
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
logger.info("CLI-only restriction: rejected non-CLI client")
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail="This endpoint only accepts requests from Claude Code CLI clients.",
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"enforce_cli_only",
|
||||
"get_original_request_headers",
|
||||
"is_claude_code_client",
|
||||
"set_original_request_headers",
|
||||
]
|
||||
@@ -0,0 +1,49 @@
|
||||
"""Claude Code adapter constants."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
CLAUDE_MESSAGES_PATH = "/v1/messages"
|
||||
DEFAULT_ANTHROPIC_VERSION = "2023-06-01"
|
||||
DEFAULT_ACCEPT = "application/json"
|
||||
STREAM_HELPER_METHOD = "stream"
|
||||
SESSION_ID_MASKING_TTL_SECONDS = 15 * 60
|
||||
|
||||
# Claude Code OAuth required betas.
|
||||
BETA_CLAUDE_CODE = "claude-code-20250219"
|
||||
BETA_OAUTH = "oauth-2025-04-20"
|
||||
BETA_INTERLEAVED_THINKING = "interleaved-thinking-2025-05-14"
|
||||
BETA_CONTEXT_1M = "context-1m-2025-08-07"
|
||||
|
||||
CLAUDE_CODE_REQUIRED_BETA_TOKENS: tuple[str, ...] = (
|
||||
BETA_CLAUDE_CODE,
|
||||
BETA_OAUTH,
|
||||
BETA_INTERLEAVED_THINKING,
|
||||
)
|
||||
|
||||
# Mimic headers observed from Claude Code traffic.
|
||||
CLAUDE_CODE_DEFAULT_HEADERS: dict[str, str] = {
|
||||
"X-Stainless-Lang": "js",
|
||||
"X-Stainless-Package-Version": "0.70.0",
|
||||
"X-Stainless-OS": "Linux",
|
||||
"X-Stainless-Arch": "arm64",
|
||||
"X-Stainless-Runtime": "node",
|
||||
"X-Stainless-Runtime-Version": "v24.13.0",
|
||||
"X-Stainless-Retry-Count": "0",
|
||||
"X-Stainless-Timeout": "600",
|
||||
"X-App": "cli",
|
||||
"Anthropic-Dangerous-Direct-Browser-Access": "true",
|
||||
}
|
||||
|
||||
__all__ = [
|
||||
"BETA_CLAUDE_CODE",
|
||||
"BETA_CONTEXT_1M",
|
||||
"BETA_INTERLEAVED_THINKING",
|
||||
"BETA_OAUTH",
|
||||
"CLAUDE_CODE_DEFAULT_HEADERS",
|
||||
"CLAUDE_CODE_REQUIRED_BETA_TOKENS",
|
||||
"CLAUDE_MESSAGES_PATH",
|
||||
"DEFAULT_ACCEPT",
|
||||
"DEFAULT_ANTHROPIC_VERSION",
|
||||
"SESSION_ID_MASKING_TTL_SECONDS",
|
||||
"STREAM_HELPER_METHOD",
|
||||
]
|
||||
@@ -0,0 +1,141 @@
|
||||
"""Claude Code request context using contextvars."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextvars
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.models.admin_requests import ClaudeCodeAdvancedConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from src.services.provider.pool.config import PoolConfig
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ClaudeCodeRequestContext:
|
||||
is_stream: bool = False
|
||||
# 使用 Key 级别作用域,确保会话限制按 OAuth 账号隔离。
|
||||
scope_key: str | None = None
|
||||
key_id: str | None = None
|
||||
max_sessions: int | None = None
|
||||
session_idle_timeout_minutes: int = 5
|
||||
session_id_masking_enabled: bool = False
|
||||
cache_ttl_override_enabled: bool = False
|
||||
cache_ttl_override_target: str = "ephemeral"
|
||||
cli_only_enabled: bool = False
|
||||
# Account Pool fields
|
||||
provider_id: str | None = None
|
||||
pool_config: PoolConfig | None = None
|
||||
session_uuid: str | None = None
|
||||
|
||||
|
||||
_claude_code_request_context: contextvars.ContextVar[ClaudeCodeRequestContext | None] = (
|
||||
contextvars.ContextVar(
|
||||
"claude_code_request_context",
|
||||
default=None,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def set_claude_code_request_context(ctx: ClaudeCodeRequestContext | None) -> None:
|
||||
_claude_code_request_context.set(ctx)
|
||||
|
||||
|
||||
def get_claude_code_request_context() -> ClaudeCodeRequestContext | None:
|
||||
return _claude_code_request_context.get()
|
||||
|
||||
|
||||
def build_claude_code_request_context(
|
||||
*,
|
||||
provider_config: Any,
|
||||
key_id: str | None,
|
||||
is_stream: bool,
|
||||
provider_id: str | None = None,
|
||||
) -> ClaudeCodeRequestContext:
|
||||
"""根据 Provider.config 构建 Claude Code 请求上下文。"""
|
||||
from src.services.provider.pool.config import parse_pool_config
|
||||
|
||||
normalized_key_id = str(key_id or "").strip() or None
|
||||
|
||||
advanced_config: ClaudeCodeAdvancedConfig | None = None
|
||||
provider_config_dict = provider_config if isinstance(provider_config, dict) else {}
|
||||
raw_advanced = provider_config_dict.get("claude_code_advanced")
|
||||
|
||||
if raw_advanced is not None:
|
||||
try:
|
||||
if isinstance(raw_advanced, ClaudeCodeAdvancedConfig):
|
||||
advanced_config = raw_advanced
|
||||
elif isinstance(raw_advanced, dict):
|
||||
advanced_config = ClaudeCodeAdvancedConfig.model_validate(raw_advanced)
|
||||
else:
|
||||
logger.warning(
|
||||
"Claude Code advanced config 类型无效: {},已忽略",
|
||||
type(raw_advanced).__name__,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning("Claude Code advanced config 解析失败,已忽略: {}", str(exc))
|
||||
|
||||
max_sessions = advanced_config.max_sessions if advanced_config else None
|
||||
idle_timeout_minutes = (
|
||||
advanced_config.session_idle_timeout_minutes
|
||||
if advanced_config and advanced_config.session_idle_timeout_minutes is not None
|
||||
else 5
|
||||
)
|
||||
session_id_masking_enabled = (
|
||||
bool(advanced_config.session_id_masking_enabled) if advanced_config else False
|
||||
)
|
||||
cache_ttl_override_enabled = (
|
||||
bool(advanced_config.cache_ttl_override_enabled) if advanced_config else False
|
||||
)
|
||||
cache_ttl_override_target = (
|
||||
str(advanced_config.cache_ttl_override_target or "ephemeral")
|
||||
if advanced_config
|
||||
else "ephemeral"
|
||||
)
|
||||
cli_only_enabled = bool(advanced_config.cli_only_enabled) if advanced_config else False
|
||||
|
||||
# Parse pool config (None = non-pool provider, keep as None for semantic consistency)
|
||||
pool_cfg = parse_pool_config(provider_config_dict)
|
||||
|
||||
return ClaudeCodeRequestContext(
|
||||
is_stream=bool(is_stream),
|
||||
scope_key=f"key:{normalized_key_id}" if normalized_key_id else None,
|
||||
key_id=normalized_key_id,
|
||||
max_sessions=max_sessions,
|
||||
session_idle_timeout_minutes=idle_timeout_minutes,
|
||||
session_id_masking_enabled=session_id_masking_enabled,
|
||||
cache_ttl_override_enabled=cache_ttl_override_enabled,
|
||||
cache_ttl_override_target=cache_ttl_override_target,
|
||||
cli_only_enabled=cli_only_enabled,
|
||||
provider_id=str(provider_id or "").strip() or None,
|
||||
pool_config=pool_cfg,
|
||||
)
|
||||
|
||||
|
||||
def build_and_set_claude_code_request_context(
|
||||
*,
|
||||
provider_config: Any,
|
||||
key_id: str | None,
|
||||
is_stream: bool,
|
||||
provider_id: str | None = None,
|
||||
) -> ClaudeCodeRequestContext:
|
||||
"""构建并写入 Claude Code 上下文。"""
|
||||
ctx = build_claude_code_request_context(
|
||||
provider_config=provider_config,
|
||||
key_id=key_id,
|
||||
is_stream=is_stream,
|
||||
provider_id=provider_id,
|
||||
)
|
||||
set_claude_code_request_context(ctx)
|
||||
return ctx
|
||||
|
||||
|
||||
__all__ = [
|
||||
"build_and_set_claude_code_request_context",
|
||||
"build_claude_code_request_context",
|
||||
"ClaudeCodeRequestContext",
|
||||
"get_claude_code_request_context",
|
||||
"set_claude_code_request_context",
|
||||
]
|
||||
@@ -0,0 +1,677 @@
|
||||
"""Claude Code upstream envelope hooks."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import replace
|
||||
from typing import Any
|
||||
|
||||
from src.clients.redis_client import get_redis_client, get_redis_client_sync
|
||||
from src.config.settings import config
|
||||
from src.core.exceptions import ConcurrencyLimitError
|
||||
from src.core.logger import logger
|
||||
from src.services.provider.adapters.claude_code.constants import (
|
||||
BETA_CONTEXT_1M,
|
||||
CLAUDE_CODE_DEFAULT_HEADERS,
|
||||
CLAUDE_CODE_REQUIRED_BETA_TOKENS,
|
||||
DEFAULT_ACCEPT,
|
||||
DEFAULT_ANTHROPIC_VERSION,
|
||||
SESSION_ID_MASKING_TTL_SECONDS,
|
||||
STREAM_HELPER_METHOD,
|
||||
)
|
||||
from src.services.provider.adapters.claude_code.context import (
|
||||
ClaudeCodeRequestContext,
|
||||
get_claude_code_request_context,
|
||||
set_claude_code_request_context,
|
||||
)
|
||||
from src.services.provider.request_context import get_current_fingerprint
|
||||
|
||||
_SESSION_MARKER = "_session_"
|
||||
_DUMMY_THINKING_SIGNATURE = "skip_thought_signature_validator"
|
||||
_session_runtime_lock = threading.Lock()
|
||||
# key: scope_key -> {session_id -> last_seen_monotonic}
|
||||
_active_sessions: dict[str, dict[str, float]] = {}
|
||||
# key: scope_key -> (masked_session_uuid, expire_at_monotonic)
|
||||
_masked_sessions: dict[str, tuple[str, float]] = {}
|
||||
# 上限保护:scope_key 总数超过此值时触发全局清理
|
||||
_MAX_SCOPE_KEYS = 5000
|
||||
# 全局清理间隔(monotonic 秒),避免高频请求时每次都做全局扫描
|
||||
_LAST_GLOBAL_CLEANUP: float = 0.0
|
||||
_GLOBAL_CLEANUP_INTERVAL = 300.0 # 5 分钟
|
||||
_REDIS_SESSION_KEY_PREFIX = "claude_code:sessions"
|
||||
_REDIS_SESSION_RESERVE_LUA = """
|
||||
local key = KEYS[1]
|
||||
local sid = ARGV[1]
|
||||
local now = tonumber(ARGV[2])
|
||||
local expire_before = tonumber(ARGV[3])
|
||||
local max_sessions = tonumber(ARGV[4])
|
||||
local ttl_seconds = tonumber(ARGV[5])
|
||||
|
||||
redis.call("ZREMRANGEBYSCORE", key, "-inf", expire_before)
|
||||
|
||||
local exists = redis.call("ZSCORE", key, sid)
|
||||
if exists then
|
||||
redis.call("ZADD", key, now, sid)
|
||||
redis.call("EXPIRE", key, ttl_seconds)
|
||||
return {1, redis.call("ZCARD", key)}
|
||||
end
|
||||
|
||||
local active = redis.call("ZCARD", key)
|
||||
if active >= max_sessions then
|
||||
return {0, active}
|
||||
end
|
||||
|
||||
redis.call("ZADD", key, now, sid)
|
||||
redis.call("EXPIRE", key, ttl_seconds)
|
||||
return {1, active + 1}
|
||||
"""
|
||||
|
||||
|
||||
def merge_anthropic_beta_tokens(
|
||||
incoming: str | None,
|
||||
*,
|
||||
required: tuple[str, ...] = CLAUDE_CODE_REQUIRED_BETA_TOKENS,
|
||||
) -> str:
|
||||
"""Merge required beta tokens and incoming anthropic-beta with deduplication."""
|
||||
seen: set[str] = set()
|
||||
merged: list[str] = []
|
||||
|
||||
def _append(token: str) -> None:
|
||||
token = token.strip()
|
||||
if not token or token in seen:
|
||||
return
|
||||
seen.add(token)
|
||||
merged.append(token)
|
||||
|
||||
for token in required:
|
||||
_append(token)
|
||||
for token in str(incoming or "").split(","):
|
||||
_append(token)
|
||||
|
||||
return ",".join(merged)
|
||||
|
||||
|
||||
def _parse_stream_flag(raw_stream: Any) -> bool:
|
||||
if isinstance(raw_stream, bool):
|
||||
return raw_stream
|
||||
return str(raw_stream).strip().lower() in {"1", "true", "yes", "on"}
|
||||
|
||||
|
||||
def _get_metadata_user_id(request_body: dict[str, Any]) -> str | None:
|
||||
metadata = request_body.get("metadata")
|
||||
if not isinstance(metadata, dict):
|
||||
return None
|
||||
user_id = metadata.get("user_id")
|
||||
if not isinstance(user_id, str):
|
||||
return None
|
||||
text = user_id.strip()
|
||||
return text or None
|
||||
|
||||
|
||||
def _set_metadata_user_id(request_body: dict[str, Any], user_id: str) -> None:
|
||||
metadata = request_body.get("metadata")
|
||||
if not isinstance(metadata, dict):
|
||||
metadata = {}
|
||||
request_body["metadata"] = metadata
|
||||
metadata["user_id"] = user_id
|
||||
|
||||
|
||||
def _extract_session_id(user_id: str) -> str | None:
|
||||
idx = user_id.rfind(_SESSION_MARKER)
|
||||
if idx == -1:
|
||||
return None
|
||||
session_id = user_id[idx + len(_SESSION_MARKER) :].strip()
|
||||
return session_id or None
|
||||
|
||||
|
||||
def _is_thinking_enabled(request_body: dict[str, Any]) -> bool:
|
||||
thinking = request_body.get("thinking")
|
||||
if not isinstance(thinking, dict):
|
||||
return False
|
||||
thinking_type = str(thinking.get("type") or "").strip().lower()
|
||||
return thinking_type in {"enabled", "adaptive"}
|
||||
|
||||
|
||||
def _sanitize_thinking_blocks(request_body: dict[str, Any]) -> None:
|
||||
"""过滤可能导致 Claude Code 400 的无效 thinking 块。"""
|
||||
messages = request_body.get("messages")
|
||||
if not isinstance(messages, list) or not messages:
|
||||
return
|
||||
|
||||
thinking_enabled = _is_thinking_enabled(request_body)
|
||||
filtered_messages = 0
|
||||
filtered_blocks = 0
|
||||
|
||||
for message in messages:
|
||||
if not isinstance(message, dict):
|
||||
continue
|
||||
|
||||
role = str(message.get("role") or "")
|
||||
content = message.get("content")
|
||||
if not isinstance(content, list):
|
||||
continue
|
||||
|
||||
new_content: list[Any] = []
|
||||
modified = False
|
||||
for block in content:
|
||||
if not isinstance(block, dict):
|
||||
new_content.append(block)
|
||||
continue
|
||||
|
||||
block_type = str(block.get("type") or "")
|
||||
if block_type in {"thinking", "redacted_thinking"}:
|
||||
keep = False
|
||||
# 仅保留 assistant 且带真实 signature 的 thinking 块。
|
||||
if thinking_enabled and role == "assistant":
|
||||
signature = str(block.get("signature") or "").strip()
|
||||
keep = bool(signature and signature != _DUMMY_THINKING_SIGNATURE)
|
||||
if keep:
|
||||
new_content.append(block)
|
||||
else:
|
||||
modified = True
|
||||
filtered_blocks += 1
|
||||
continue
|
||||
|
||||
# 兼容无 type 但带 thinking 字段的历史块,直接移除。
|
||||
if not block_type and "thinking" in block:
|
||||
modified = True
|
||||
filtered_blocks += 1
|
||||
continue
|
||||
|
||||
new_content.append(block)
|
||||
|
||||
if modified:
|
||||
message["content"] = new_content
|
||||
filtered_messages += 1
|
||||
|
||||
if filtered_blocks:
|
||||
logger.info(
|
||||
"Claude Code thinking 预过滤: messages={}, blocks={}, thinking_enabled={}",
|
||||
filtered_messages,
|
||||
filtered_blocks,
|
||||
thinking_enabled,
|
||||
)
|
||||
|
||||
|
||||
def _get_or_create_masked_session(scope_key: str) -> str:
|
||||
now = time.monotonic()
|
||||
with _session_runtime_lock:
|
||||
existing = _masked_sessions.get(scope_key)
|
||||
if existing and existing[1] > now:
|
||||
masked_session_id = existing[0]
|
||||
else:
|
||||
masked_session_id = str(uuid.uuid4())
|
||||
_masked_sessions[scope_key] = (
|
||||
masked_session_id,
|
||||
now + SESSION_ID_MASKING_TTL_SECONDS,
|
||||
)
|
||||
return masked_session_id
|
||||
|
||||
|
||||
def _apply_session_id_masking(request_body: dict[str, Any], *, scope_key: str) -> None:
|
||||
user_id = _get_metadata_user_id(request_body)
|
||||
if not user_id:
|
||||
return
|
||||
idx = user_id.rfind(_SESSION_MARKER)
|
||||
if idx == -1:
|
||||
return
|
||||
masked_session_id = _get_or_create_masked_session(scope_key)
|
||||
_set_metadata_user_id(
|
||||
request_body,
|
||||
user_id[: idx + len(_SESSION_MARKER)] + masked_session_id,
|
||||
)
|
||||
|
||||
|
||||
# -- Cache TTL Override -------------------------------------------------------
|
||||
|
||||
_VALID_CACHE_TTL_TARGETS = {"ephemeral", "1h"}
|
||||
|
||||
|
||||
def _override_cache_control_in_blocks(blocks: list[Any], target: str) -> int:
|
||||
"""Override cache_control TTL in a list of content blocks. Returns count of overrides.
|
||||
|
||||
According to Anthropic API docs, cache_control format is:
|
||||
{"type": "ephemeral", "ttl": "5m" | "1h"}
|
||||
``type`` is always "ephemeral"; the ``ttl`` field controls the actual duration.
|
||||
When ttl is absent, the default is 5m (ephemeral).
|
||||
"""
|
||||
count = 0
|
||||
for block in blocks:
|
||||
if not isinstance(block, dict):
|
||||
continue
|
||||
cc = block.get("cache_control")
|
||||
if not isinstance(cc, dict):
|
||||
continue
|
||||
# Ensure type is always "ephemeral"
|
||||
if cc.get("type") != "ephemeral":
|
||||
cc["type"] = "ephemeral"
|
||||
if target == "ephemeral":
|
||||
# Target is 5m (default) -- remove explicit ttl so it falls back to default
|
||||
if "ttl" in cc:
|
||||
del cc["ttl"]
|
||||
count += 1
|
||||
else:
|
||||
# Target is "1h" -- set ttl explicitly
|
||||
if cc.get("ttl") != target:
|
||||
cc["ttl"] = target
|
||||
count += 1
|
||||
return count
|
||||
|
||||
|
||||
def _apply_cache_ttl_override(request_body: dict[str, Any], target: str) -> None:
|
||||
"""Force all cache_control entries to use a unified TTL type.
|
||||
|
||||
Prevents multi-user behavioral fingerprinting when sharing an OAuth account.
|
||||
"""
|
||||
if target not in _VALID_CACHE_TTL_TARGETS:
|
||||
return
|
||||
|
||||
overridden = 0
|
||||
|
||||
# system prompt (can be string or list of blocks)
|
||||
system = request_body.get("system")
|
||||
if isinstance(system, list):
|
||||
overridden += _override_cache_control_in_blocks(system, target)
|
||||
|
||||
# messages
|
||||
messages = request_body.get("messages")
|
||||
if isinstance(messages, list):
|
||||
for msg in messages:
|
||||
if not isinstance(msg, dict):
|
||||
continue
|
||||
content = msg.get("content")
|
||||
if isinstance(content, list):
|
||||
overridden += _override_cache_control_in_blocks(content, target)
|
||||
|
||||
# tools
|
||||
tools = request_body.get("tools")
|
||||
if isinstance(tools, list):
|
||||
overridden += _override_cache_control_in_blocks(tools, target)
|
||||
|
||||
if overridden:
|
||||
logger.debug("Cache TTL override: {} block(s) -> {}", overridden, target)
|
||||
|
||||
|
||||
def _cleanup_stale_scope_keys(now: float, idle_seconds: int) -> None:
|
||||
"""清理空 bucket 和过期的 _masked_sessions 条目(调用方须持有锁)。"""
|
||||
global _LAST_GLOBAL_CLEANUP
|
||||
|
||||
if now - _LAST_GLOBAL_CLEANUP < _GLOBAL_CLEANUP_INTERVAL:
|
||||
return
|
||||
_LAST_GLOBAL_CLEANUP = now
|
||||
|
||||
# 清理所有 scope_key 下的过期 session,并删除空 bucket
|
||||
stale_keys = []
|
||||
for sk, bucket in _active_sessions.items():
|
||||
expired = [sid for sid, ts in bucket.items() if now - ts > idle_seconds]
|
||||
for sid in expired:
|
||||
bucket.pop(sid, None)
|
||||
if not bucket:
|
||||
stale_keys.append(sk)
|
||||
for sk in stale_keys:
|
||||
_active_sessions.pop(sk, None)
|
||||
|
||||
# 清理过期的 masked session 条目
|
||||
expired_masked = [sk for sk, (_, exp) in _masked_sessions.items() if exp <= now]
|
||||
for sk in expired_masked:
|
||||
_masked_sessions.pop(sk, None)
|
||||
|
||||
total = len(_active_sessions) + len(_masked_sessions)
|
||||
if total > 0 and (stale_keys or expired_masked):
|
||||
logger.debug(
|
||||
"Session 全局清理: 移除 {} 个 active scope + {} 个 masked scope, 剩余 {}",
|
||||
len(stale_keys),
|
||||
len(expired_masked),
|
||||
total,
|
||||
)
|
||||
|
||||
|
||||
def _register_or_reject_session(
|
||||
*,
|
||||
scope_key: str,
|
||||
session_id: str,
|
||||
max_sessions: int,
|
||||
idle_timeout_minutes: int,
|
||||
) -> tuple[bool, int]:
|
||||
now = time.monotonic()
|
||||
idle_seconds = max(60, int(idle_timeout_minutes * 60))
|
||||
|
||||
with _session_runtime_lock:
|
||||
# scope_key 总数超限或达到清理间隔时,执行全局清理
|
||||
if (
|
||||
len(_active_sessions) + len(_masked_sessions) > _MAX_SCOPE_KEYS
|
||||
or now - _LAST_GLOBAL_CLEANUP >= _GLOBAL_CLEANUP_INTERVAL
|
||||
):
|
||||
_cleanup_stale_scope_keys(now, idle_seconds)
|
||||
|
||||
bucket = _active_sessions.setdefault(scope_key, {})
|
||||
|
||||
# 先清理当前 bucket 的过期会话,避免误判占用。
|
||||
expired = [sid for sid, last_seen in bucket.items() if now - last_seen > idle_seconds]
|
||||
for sid in expired:
|
||||
bucket.pop(sid, None)
|
||||
|
||||
if session_id in bucket:
|
||||
bucket[session_id] = now
|
||||
return True, len(bucket)
|
||||
|
||||
if len(bucket) >= max_sessions:
|
||||
return False, len(bucket)
|
||||
|
||||
bucket[session_id] = now
|
||||
return True, len(bucket)
|
||||
|
||||
|
||||
def _build_session_limit_error(
|
||||
*,
|
||||
max_sessions: int,
|
||||
active_count: int,
|
||||
key_id: str | None,
|
||||
) -> ConcurrencyLimitError:
|
||||
return ConcurrencyLimitError(
|
||||
message=(f"Claude Code 活跃会话数已达上限({max_sessions})。当前活跃会话: {active_count}"),
|
||||
key_id=key_id,
|
||||
)
|
||||
|
||||
|
||||
def _redis_session_key(scope_key: str) -> str:
|
||||
return f"{_REDIS_SESSION_KEY_PREFIX}:{scope_key}"
|
||||
|
||||
|
||||
def _parse_redis_session_result(raw: Any) -> tuple[bool, int] | None:
|
||||
if not isinstance(raw, (list, tuple)) or len(raw) < 2:
|
||||
return None
|
||||
try:
|
||||
allowed = int(raw[0]) == 1
|
||||
active_count = int(raw[1])
|
||||
except Exception:
|
||||
return None
|
||||
return allowed, active_count
|
||||
|
||||
|
||||
def _enforce_session_controls(
|
||||
request_body: dict[str, Any],
|
||||
ctx: ClaudeCodeRequestContext,
|
||||
*,
|
||||
enforce_max_sessions: bool = True,
|
||||
) -> None:
|
||||
"""同步执行会话限制 + masking。
|
||||
|
||||
当 ``enforce_max_sessions=False``(将由 ``enforce_distributed_session_controls``
|
||||
异步接管)时仅做 masking;masking 始终在会话限制检查之后执行,避免
|
||||
用被伪装后的 session_id 做计数。
|
||||
"""
|
||||
scope_key = str(ctx.scope_key or "").strip()
|
||||
if not scope_key:
|
||||
return
|
||||
|
||||
# 先基于真实 session_id 做会话限制检查。
|
||||
if enforce_max_sessions and ctx.max_sessions and ctx.max_sessions > 0:
|
||||
user_id = _get_metadata_user_id(request_body)
|
||||
if user_id:
|
||||
session_id = _extract_session_id(user_id) or user_id
|
||||
allowed, active_count = _register_or_reject_session(
|
||||
scope_key=scope_key,
|
||||
session_id=session_id,
|
||||
max_sessions=ctx.max_sessions,
|
||||
idle_timeout_minutes=ctx.session_idle_timeout_minutes,
|
||||
)
|
||||
if not allowed:
|
||||
raise _build_session_limit_error(
|
||||
max_sessions=ctx.max_sessions,
|
||||
active_count=active_count,
|
||||
key_id=ctx.key_id,
|
||||
)
|
||||
|
||||
if not enforce_max_sessions:
|
||||
# 分布式模式下 masking 延迟到 enforce_distributed_session_controls 中执行,
|
||||
# 避免 wrap_request 提前改写 user_id 导致分布式检查拿到伪装后的 session_id。
|
||||
return
|
||||
|
||||
# 仅在本地模式下立即 masking。
|
||||
if ctx.session_id_masking_enabled:
|
||||
_apply_session_id_masking(request_body, scope_key=scope_key)
|
||||
|
||||
|
||||
def _is_distributed_session_control_available() -> bool:
|
||||
try:
|
||||
return get_redis_client_sync() is not None
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
async def enforce_distributed_session_controls(
|
||||
request_body: dict[str, Any],
|
||||
ctx: ClaudeCodeRequestContext | None,
|
||||
) -> None:
|
||||
"""异步执行会话限制 + masking。
|
||||
|
||||
优先使用 Redis(多实例共享);Redis 不可用时回退到进程内计数。
|
||||
masking 在会话限制检查通过后执行,确保计数使用真实 session_id。
|
||||
"""
|
||||
if ctx is None:
|
||||
return
|
||||
|
||||
scope_key = str(ctx.scope_key or "").strip()
|
||||
if not scope_key:
|
||||
# 即使无 scope_key 也无法做 masking(需要 scope_key 作为 key),直接返回。
|
||||
return
|
||||
|
||||
if not ctx.max_sessions or ctx.max_sessions <= 0:
|
||||
# 无会话限制,仅做 masking。
|
||||
if ctx.session_id_masking_enabled:
|
||||
_apply_session_id_masking(request_body, scope_key=scope_key)
|
||||
return
|
||||
|
||||
# 基于真实 user_id 提取 session_id 做限制检查。
|
||||
user_id = _get_metadata_user_id(request_body)
|
||||
if not user_id:
|
||||
return
|
||||
session_id = _extract_session_id(user_id) or user_id
|
||||
|
||||
idle_seconds = max(60, int(ctx.session_idle_timeout_minutes * 60))
|
||||
redis_ttl = idle_seconds + 300
|
||||
now = int(time.time())
|
||||
expire_before = now - idle_seconds
|
||||
|
||||
redis_client = await get_redis_client(require_redis=False)
|
||||
if redis_client is not None:
|
||||
try:
|
||||
raw_result = await redis_client.eval(
|
||||
_REDIS_SESSION_RESERVE_LUA,
|
||||
1,
|
||||
_redis_session_key(scope_key),
|
||||
session_id,
|
||||
str(now),
|
||||
str(expire_before),
|
||||
str(ctx.max_sessions),
|
||||
str(redis_ttl),
|
||||
)
|
||||
parsed = _parse_redis_session_result(raw_result)
|
||||
if parsed is None:
|
||||
raise ValueError(f"invalid redis eval result: {raw_result!r}")
|
||||
|
||||
allowed, active_count = parsed
|
||||
if not allowed:
|
||||
raise _build_session_limit_error(
|
||||
max_sessions=ctx.max_sessions,
|
||||
active_count=active_count,
|
||||
key_id=ctx.key_id,
|
||||
)
|
||||
# 会话限制通过后再做 masking。
|
||||
if ctx.session_id_masking_enabled:
|
||||
_apply_session_id_masking(request_body, scope_key=scope_key)
|
||||
return
|
||||
except ConcurrencyLimitError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.warning("Claude Code 分布式会话控制失败,回退本地计数: {}", str(exc))
|
||||
|
||||
allowed, active_count = _register_or_reject_session(
|
||||
scope_key=scope_key,
|
||||
session_id=session_id,
|
||||
max_sessions=ctx.max_sessions,
|
||||
idle_timeout_minutes=ctx.session_idle_timeout_minutes,
|
||||
)
|
||||
if not allowed:
|
||||
raise _build_session_limit_error(
|
||||
max_sessions=ctx.max_sessions,
|
||||
active_count=active_count,
|
||||
key_id=ctx.key_id,
|
||||
)
|
||||
|
||||
# 会话限制通过后再做 masking。
|
||||
if ctx.session_id_masking_enabled:
|
||||
_apply_session_id_masking(request_body, scope_key=scope_key)
|
||||
|
||||
|
||||
class ClaudeCodeEnvelope:
|
||||
"""Provider envelope hooks for Claude Code OAuth upstream."""
|
||||
|
||||
name = "claude:cli"
|
||||
|
||||
def extra_headers(self) -> dict[str, str] | None:
|
||||
ctx = get_claude_code_request_context()
|
||||
is_stream = bool(ctx.is_stream) if ctx else False
|
||||
fp = get_current_fingerprint()
|
||||
|
||||
headers = dict(CLAUDE_CODE_DEFAULT_HEADERS)
|
||||
headers["Accept"] = DEFAULT_ACCEPT
|
||||
headers["anthropic-version"] = DEFAULT_ANTHROPIC_VERSION
|
||||
headers["anthropic-beta"] = merge_anthropic_beta_tokens(None)
|
||||
if is_stream:
|
||||
headers["x-stainless-helper-method"] = STREAM_HELPER_METHOD
|
||||
|
||||
if fp:
|
||||
headers["X-Stainless-Package-Version"] = fp.stainless_package_version
|
||||
headers["X-Stainless-OS"] = fp.stainless_os
|
||||
headers["X-Stainless-Arch"] = fp.stainless_arch
|
||||
headers["X-Stainless-Runtime-Version"] = fp.stainless_runtime_version
|
||||
headers["X-Stainless-Timeout"] = fp.stainless_timeout
|
||||
headers["User-Agent"] = fp.user_agent
|
||||
else:
|
||||
ua = str(getattr(config, "internal_user_agent_claude_cli", "") or "").strip()
|
||||
if ua:
|
||||
headers["User-Agent"] = ua
|
||||
|
||||
return headers
|
||||
|
||||
def wrap_request(
|
||||
self,
|
||||
request_body: dict[str, Any],
|
||||
*,
|
||||
model: str, # noqa: ARG002
|
||||
url_model: str | None,
|
||||
decrypted_auth_config: dict[str, Any] | None, # noqa: ARG002
|
||||
) -> tuple[dict[str, Any], str | None]:
|
||||
raw_stream = request_body.get("stream", False)
|
||||
is_stream = _parse_stream_flag(raw_stream)
|
||||
|
||||
ctx = get_claude_code_request_context()
|
||||
if ctx is None:
|
||||
ctx = ClaudeCodeRequestContext()
|
||||
|
||||
# CLI-only restriction: reject non-CLI clients early.
|
||||
if ctx.cli_only_enabled:
|
||||
from src.services.provider.adapters.claude_code.client_restriction import (
|
||||
enforce_cli_only,
|
||||
)
|
||||
|
||||
enforce_cli_only(ctx.cli_only_enabled)
|
||||
|
||||
# Extract session_uuid from metadata.user_id for pool sticky session.
|
||||
session_uuid: str | None = None
|
||||
user_id = _get_metadata_user_id(request_body)
|
||||
if user_id:
|
||||
session_uuid = _extract_session_id(user_id)
|
||||
ctx = replace(ctx, is_stream=is_stream, session_uuid=session_uuid)
|
||||
set_claude_code_request_context(ctx)
|
||||
|
||||
_sanitize_thinking_blocks(request_body)
|
||||
|
||||
# Cache TTL override: unify cache_control types to prevent behavioral fingerprinting.
|
||||
if ctx.cache_ttl_override_enabled:
|
||||
_apply_cache_ttl_override(request_body, ctx.cache_ttl_override_target)
|
||||
|
||||
_enforce_session_controls(
|
||||
request_body,
|
||||
ctx,
|
||||
enforce_max_sessions=not _is_distributed_session_control_available(),
|
||||
)
|
||||
return request_body, url_model
|
||||
|
||||
def unwrap_response(self, data: Any) -> Any:
|
||||
return data
|
||||
|
||||
def postprocess_unwrapped_response(self, *, model: str, data: Any) -> None: # noqa: ARG002
|
||||
return
|
||||
|
||||
def capture_selected_base_url(self) -> str | None:
|
||||
return None
|
||||
|
||||
def on_http_status(self, *, base_url: str | None, status_code: int) -> None: # noqa: ARG002
|
||||
return
|
||||
|
||||
def on_connection_error(self, *, base_url: str | None, exc: Exception) -> None: # noqa: ARG002
|
||||
return
|
||||
|
||||
def force_stream_rewrite(self) -> bool:
|
||||
return False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Optional lifecycle hooks
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def prepare_context(
|
||||
self,
|
||||
*,
|
||||
provider_config: Any,
|
||||
key_id: str,
|
||||
user_api_key_id: str | None = None, # noqa: ARG002
|
||||
is_stream: bool,
|
||||
provider_id: str | None = None,
|
||||
key: Any = None,
|
||||
) -> str | None:
|
||||
from src.services.provider.adapters.claude_code.context import (
|
||||
build_and_set_claude_code_request_context,
|
||||
)
|
||||
|
||||
# 在 envelope 层设置指纹 context var(仅 Claude Code 需要指纹注入)
|
||||
if key is not None:
|
||||
from src.services.provider.fingerprint import ensure_key_fingerprint
|
||||
from src.services.provider.request_context import set_current_fingerprint
|
||||
|
||||
set_current_fingerprint(ensure_key_fingerprint(key, persist_if_missing=True))
|
||||
|
||||
build_and_set_claude_code_request_context(
|
||||
provider_config=provider_config,
|
||||
key_id=key_id,
|
||||
is_stream=is_stream,
|
||||
provider_id=provider_id,
|
||||
)
|
||||
fp = get_current_fingerprint()
|
||||
if fp:
|
||||
return fp.impersonate
|
||||
return None
|
||||
|
||||
async def post_wrap_request(self, request_body: dict[str, Any]) -> None:
|
||||
await enforce_distributed_session_controls(
|
||||
request_body,
|
||||
get_claude_code_request_context(),
|
||||
)
|
||||
|
||||
def excluded_beta_tokens(self) -> frozenset[str]:
|
||||
return frozenset({BETA_CONTEXT_1M})
|
||||
|
||||
|
||||
claude_code_envelope = ClaudeCodeEnvelope()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ClaudeCodeEnvelope",
|
||||
"claude_code_envelope",
|
||||
"enforce_distributed_session_controls",
|
||||
"merge_anthropic_beta_tokens",
|
||||
]
|
||||
@@ -0,0 +1,63 @@
|
||||
"""Claude Code provider plugin."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
from urllib.parse import urlencode
|
||||
|
||||
from src.services.provider.adapters.claude_code.constants import CLAUDE_MESSAGES_PATH
|
||||
from src.services.provider.preset_models import create_preset_models_fetcher
|
||||
|
||||
fetch_models_claude_code = create_preset_models_fetcher("claude_code")
|
||||
|
||||
|
||||
def build_claude_code_url(
|
||||
endpoint: Any,
|
||||
*,
|
||||
is_stream: bool,
|
||||
effective_query_params: dict[str, Any],
|
||||
**_kwargs: Any,
|
||||
) -> str:
|
||||
"""Build Claude Code upstream URL and avoid duplicate /v1/messages suffix."""
|
||||
_ = is_stream
|
||||
|
||||
base = str(getattr(endpoint, "base_url", "") or "").rstrip("/")
|
||||
if base.endswith(CLAUDE_MESSAGES_PATH) or base.endswith("/messages"):
|
||||
url = base
|
||||
elif base.endswith("/v1"):
|
||||
url = f"{base}/messages"
|
||||
else:
|
||||
url = f"{base}{CLAUDE_MESSAGES_PATH}"
|
||||
|
||||
if effective_query_params:
|
||||
query_string = urlencode(effective_query_params, doseq=True)
|
||||
if query_string:
|
||||
url = f"{url}?{query_string}"
|
||||
|
||||
return url
|
||||
|
||||
|
||||
def register_all() -> None:
|
||||
"""Register Claude Code hooks into shared registries."""
|
||||
from src.services.model.upstream_fetcher import UpstreamModelsFetcherRegistry
|
||||
from src.services.provider.adapters.claude_code.envelope import claude_code_envelope
|
||||
from src.services.provider.envelope import register_envelope
|
||||
from src.services.provider.transport import register_transport_hook
|
||||
|
||||
register_envelope("claude_code", "claude:cli", claude_code_envelope)
|
||||
register_envelope("claude_code", "", claude_code_envelope)
|
||||
|
||||
register_transport_hook("claude_code", "claude:cli", build_claude_code_url)
|
||||
|
||||
UpstreamModelsFetcherRegistry.register(
|
||||
provider_types=["claude_code"],
|
||||
fetcher=fetch_models_claude_code,
|
||||
)
|
||||
|
||||
from src.services.provider.adapters.claude_code.pool_hook import claude_code_pool_hook
|
||||
from src.services.provider.pool.hooks import register_pool_hook
|
||||
|
||||
register_pool_hook("claude_code", claude_code_pool_hook)
|
||||
|
||||
|
||||
__all__ = ["build_claude_code_url", "fetch_models_claude_code", "register_all"]
|
||||
@@ -0,0 +1,6 @@
|
||||
"""Backward compat shim -- canonical definitions moved to src.services.provider.pool.config."""
|
||||
|
||||
from src.services.provider.pool.config import * # noqa: F401,F403
|
||||
from src.services.provider.pool.config import PoolConfig, UnschedulableRule, parse_pool_config
|
||||
|
||||
__all__ = ["PoolConfig", "UnschedulableRule", "parse_pool_config"]
|
||||
@@ -0,0 +1,10 @@
|
||||
"""Backward compat shim -- canonical definitions moved to src.services.provider.pool.cost_tracker."""
|
||||
|
||||
from src.services.provider.pool.cost_tracker import ( # noqa: F401
|
||||
get_window_usage,
|
||||
is_approaching_limit,
|
||||
is_at_limit,
|
||||
record_usage,
|
||||
)
|
||||
|
||||
__all__ = ["record_usage", "get_window_usage", "is_at_limit", "is_approaching_limit"]
|
||||
@@ -0,0 +1,6 @@
|
||||
"""Backward compat shim -- canonical definitions moved to src.services.provider.pool.health_policy."""
|
||||
|
||||
from src.services.provider.pool.health_policy import * # noqa: F401,F403
|
||||
from src.services.provider.pool.health_policy import apply_health_policy # noqa: F811
|
||||
|
||||
__all__ = ["apply_health_policy"]
|
||||
@@ -0,0 +1,27 @@
|
||||
"""Claude Code pool scheduling hook."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
class ClaudeCodePoolHook:
|
||||
"""Pool scheduling hook for Claude Code providers.
|
||||
|
||||
Extracts the session UUID from ``metadata.user_id`` which follows the
|
||||
pattern ``<user>_session_<uuid>``.
|
||||
"""
|
||||
|
||||
name = "claude_code"
|
||||
|
||||
def extract_session_uuid(self, request_body: dict[str, Any]) -> str | None:
|
||||
metadata = request_body.get("metadata")
|
||||
if isinstance(metadata, dict):
|
||||
user_id = metadata.get("user_id")
|
||||
if isinstance(user_id, str) and "_session_" in user_id:
|
||||
idx = user_id.rfind("_session_")
|
||||
return user_id[idx + len("_session_") :].strip() or None
|
||||
return None
|
||||
|
||||
|
||||
claude_code_pool_hook = ClaudeCodePoolHook()
|
||||
@@ -0,0 +1,8 @@
|
||||
"""Backward compat shim -- canonical definitions moved to src.services.provider.pool.manager."""
|
||||
|
||||
from src.services.provider.pool.manager import * # noqa: F401,F403
|
||||
from src.services.provider.pool.manager import PoolManager
|
||||
|
||||
ClaudeCodePoolManager = PoolManager # noqa: F811
|
||||
|
||||
__all__ = ["PoolManager", "ClaudeCodePoolManager"]
|
||||
@@ -0,0 +1,9 @@
|
||||
"""Backward compat shim -- canonical definitions moved to src.services.provider.pool.oauth_cache."""
|
||||
|
||||
from src.services.provider.pool.oauth_cache import ( # noqa: F401
|
||||
cache_token,
|
||||
get_cached_token,
|
||||
invalidate_token,
|
||||
)
|
||||
|
||||
__all__ = ["get_cached_token", "cache_token", "invalidate_token"]
|
||||
@@ -0,0 +1,23 @@
|
||||
"""Backward compat shim -- canonical definitions moved to src.services.provider.pool.redis_ops."""
|
||||
|
||||
from src.services.provider.pool.redis_ops import * # noqa: F401,F403
|
||||
from src.services.provider.pool.redis_ops import (
|
||||
add_cost_entry,
|
||||
batch_get_cooldowns,
|
||||
cache_oauth_token,
|
||||
clear_cooldown,
|
||||
clear_cost,
|
||||
delete_sticky_binding,
|
||||
get_cached_oauth_token,
|
||||
get_cooldown,
|
||||
get_cooldown_ttl,
|
||||
get_cost_window_total,
|
||||
get_key_sticky_count,
|
||||
get_lru_scores,
|
||||
get_sticky_binding,
|
||||
get_sticky_session_count,
|
||||
invalidate_oauth_token_cache,
|
||||
set_cooldown,
|
||||
set_sticky_binding,
|
||||
touch_lru,
|
||||
)
|
||||
@@ -0,0 +1,3 @@
|
||||
"""Codex provider integration package."""
|
||||
|
||||
__all__ = []
|
||||
@@ -0,0 +1,55 @@
|
||||
"""Codex request-scoped context.
|
||||
|
||||
Codex only needs a small amount of per-request runtime state that does not belong in
|
||||
the outbound payload itself. Today that state is the compact-mode flag used by the
|
||||
transport and upstream stream-policy layers.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextvars
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CodexRequestContext:
|
||||
"""Per-request context for the Codex adapter."""
|
||||
|
||||
is_compact: bool = False
|
||||
|
||||
|
||||
_codex_request_context: contextvars.ContextVar[CodexRequestContext | None] = contextvars.ContextVar(
|
||||
"codex_request_context",
|
||||
default=None,
|
||||
)
|
||||
|
||||
|
||||
def set_codex_request_context(ctx: CodexRequestContext | None) -> None:
|
||||
_codex_request_context.set(ctx)
|
||||
|
||||
|
||||
def get_codex_request_context() -> CodexRequestContext | None:
|
||||
return _codex_request_context.get()
|
||||
|
||||
|
||||
def is_codex_compact_request(*, endpoint_sig: str | None = None) -> bool:
|
||||
"""Return whether the current Codex request should use compact semantics.
|
||||
|
||||
Modern configurations use a dedicated ``openai:compact`` endpoint. Older ones may
|
||||
still route compact traffic through ``openai:cli`` and rely on request-scoped
|
||||
context instead.
|
||||
"""
|
||||
normalized_sig = str(endpoint_sig or "").strip().lower()
|
||||
if normalized_sig == "openai:compact":
|
||||
return True
|
||||
|
||||
ctx = get_codex_request_context()
|
||||
return bool(ctx and ctx.is_compact)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"CodexRequestContext",
|
||||
"get_codex_request_context",
|
||||
"is_codex_compact_request",
|
||||
"set_codex_request_context",
|
||||
]
|
||||
371
_deprecated_py_src/services/provider/adapters/codex/plugin.py
Normal file
371
_deprecated_py_src/services/provider/adapters/codex/plugin.py
Normal file
@@ -0,0 +1,371 @@
|
||||
"""Codex provider plugin — 统一注册入口。
|
||||
|
||||
将 Codex 对各通用 registry / capability registry 的注册集中在一个文件中:
|
||||
- Transport Hook (URL 构建)
|
||||
- Auth Enricher (OAuth enrichment)
|
||||
- Provider Format Capability(默认 body_rules)
|
||||
- Model Fetcher (fixed catalog — Codex has no /v1/models endpoint)
|
||||
|
||||
新增 provider 时参照此文件创建对应的 plugin.py 即可。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
from collections.abc import Mapping
|
||||
from typing import Any
|
||||
from urllib.parse import urlencode
|
||||
|
||||
from src.core.logger import logger
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Preset model catalog
|
||||
# ---------------------------------------------------------------------------
|
||||
# Codex upstream (chatgpt.com/backend-api/codex) has no /v1/models endpoint.
|
||||
# We use the unified preset models registry from preset_models.py.
|
||||
from src.services.provider.preset_models import create_preset_models_fetcher
|
||||
|
||||
fetch_models_codex = create_preset_models_fetcher("codex")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Transport Hook
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _get_header_value(headers: Mapping[str, Any] | None, header_name: str) -> str | None:
|
||||
if not isinstance(headers, Mapping):
|
||||
return None
|
||||
|
||||
target = str(header_name or "").strip().lower()
|
||||
if not target:
|
||||
return None
|
||||
|
||||
for name, value in headers.items():
|
||||
if str(name or "").strip().lower() != target:
|
||||
continue
|
||||
normalized = str(value or "").strip()
|
||||
return normalized or None
|
||||
return None
|
||||
|
||||
|
||||
def _build_short_header_id(seed: str) -> str:
|
||||
return hashlib.sha256(seed.encode()).hexdigest()[:16]
|
||||
|
||||
|
||||
def _normalize_codex_debug_organizations(value: Any) -> list[dict[str, Any]]:
|
||||
if not isinstance(value, list):
|
||||
return []
|
||||
|
||||
items: list[dict[str, Any]] = []
|
||||
for raw in value:
|
||||
if not isinstance(raw, Mapping):
|
||||
continue
|
||||
|
||||
normalized: dict[str, Any] = {}
|
||||
org_id = str(raw.get("id") or "").strip()
|
||||
if org_id:
|
||||
normalized["id"] = org_id
|
||||
|
||||
title = str(raw.get("title") or "").strip()
|
||||
if title:
|
||||
normalized["title"] = title
|
||||
|
||||
role = str(raw.get("role") or "").strip()
|
||||
if role:
|
||||
normalized["role"] = role
|
||||
|
||||
if "is_default" in raw:
|
||||
normalized["is_default"] = bool(raw.get("is_default"))
|
||||
|
||||
if normalized:
|
||||
items.append(normalized)
|
||||
|
||||
return items
|
||||
|
||||
|
||||
def _build_safe_codex_debug_snapshot(values: Mapping[str, Any] | None) -> dict[str, Any]:
|
||||
if not isinstance(values, Mapping):
|
||||
return {}
|
||||
|
||||
snapshot: dict[str, Any] = {}
|
||||
for field in (
|
||||
"email",
|
||||
"account_id",
|
||||
"account_user_id",
|
||||
"plan_type",
|
||||
"user_id",
|
||||
"account_name",
|
||||
):
|
||||
raw = values.get(field)
|
||||
if isinstance(raw, str):
|
||||
normalized = raw.strip()
|
||||
if normalized:
|
||||
snapshot[field] = normalized
|
||||
elif raw is not None:
|
||||
snapshot[field] = raw
|
||||
|
||||
organizations = _normalize_codex_debug_organizations(values.get("organizations"))
|
||||
if organizations:
|
||||
snapshot["organizations"] = organizations
|
||||
|
||||
return snapshot
|
||||
|
||||
|
||||
def _build_codex_headers(
|
||||
request_body: Any,
|
||||
original_headers: Mapping[str, Any] | None,
|
||||
*,
|
||||
include_conversation_id: bool,
|
||||
decrypted_auth_config: dict[str, Any] | None = None,
|
||||
) -> dict[str, str]:
|
||||
"""Codex upstream: chatgpt-account-id + session_id + conversation_id."""
|
||||
headers: dict[str, str] = {}
|
||||
|
||||
if decrypted_auth_config:
|
||||
account_id = str(decrypted_auth_config.get("account_id") or "").strip()
|
||||
if account_id:
|
||||
headers["chatgpt-account-id"] = account_id
|
||||
|
||||
if isinstance(request_body, dict):
|
||||
cache_key = str(request_body.get("prompt_cache_key") or "").strip()
|
||||
if cache_key:
|
||||
short_id = _build_short_header_id(cache_key)
|
||||
if not _get_header_value(original_headers, "session_id"):
|
||||
headers["session_id"] = short_id
|
||||
if include_conversation_id and not _get_header_value(
|
||||
original_headers, "conversation_id"
|
||||
):
|
||||
headers["conversation_id"] = short_id
|
||||
|
||||
return headers
|
||||
|
||||
|
||||
def build_codex_cli_headers(
|
||||
request_body: Any,
|
||||
original_headers: Mapping[str, Any] | None,
|
||||
*,
|
||||
decrypted_auth_config: dict[str, Any] | None = None,
|
||||
) -> dict[str, str]:
|
||||
from src.services.provider.adapters.codex.context import is_codex_compact_request
|
||||
|
||||
return _build_codex_headers(
|
||||
request_body,
|
||||
original_headers,
|
||||
include_conversation_id=not is_codex_compact_request(endpoint_sig="openai:cli"),
|
||||
decrypted_auth_config=decrypted_auth_config,
|
||||
)
|
||||
|
||||
|
||||
def build_codex_compact_headers(
|
||||
request_body: Any,
|
||||
original_headers: Mapping[str, Any] | None,
|
||||
*,
|
||||
decrypted_auth_config: dict[str, Any] | None = None,
|
||||
) -> dict[str, str]:
|
||||
return _build_codex_headers(
|
||||
request_body,
|
||||
original_headers,
|
||||
include_conversation_id=False,
|
||||
decrypted_auth_config=decrypted_auth_config,
|
||||
)
|
||||
|
||||
|
||||
def build_codex_url(
|
||||
endpoint: Any,
|
||||
*,
|
||||
is_stream: bool,
|
||||
effective_query_params: dict[str, Any],
|
||||
**_kwargs: Any,
|
||||
) -> str:
|
||||
"""构建 Codex OAuth URL。
|
||||
|
||||
Codex upstream (chatgpt.com/backend-api/codex) 使用 /responses
|
||||
而非标准 OpenAI 的 /v1/responses。compact 模式使用 /responses/compact。
|
||||
"""
|
||||
_ = is_stream # Codex 不需要根据 stream 切换路径
|
||||
|
||||
endpoint_sig = str(getattr(endpoint, "api_format", "") or "").strip().lower()
|
||||
from src.services.provider.adapters.codex.context import is_codex_compact_request
|
||||
|
||||
is_compact = is_codex_compact_request(endpoint_sig=endpoint_sig)
|
||||
|
||||
base = str(endpoint.base_url).rstrip("/")
|
||||
# 如果用户已在 base_url 中包含了 /responses,不要重复追加
|
||||
if base.endswith("/responses"):
|
||||
url = f"{base}/compact" if is_compact else base
|
||||
elif base.endswith("/responses/compact"):
|
||||
url = base if is_compact else base.removesuffix("/compact")
|
||||
else:
|
||||
suffix = "/responses/compact" if is_compact else "/responses"
|
||||
url = f"{base}{suffix}"
|
||||
if effective_query_params:
|
||||
query_string = urlencode(effective_query_params, doseq=True)
|
||||
if query_string:
|
||||
url = f"{url}?{query_string}"
|
||||
return url
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Auth Enricher
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def enrich_codex(
|
||||
auth_config: dict[str, Any],
|
||||
token_response: dict[str, Any],
|
||||
access_token: str,
|
||||
proxy_config: dict[str, Any] | None,
|
||||
) -> dict[str, Any]:
|
||||
"""Codex auth_config enrichment: parse token claims -> account/team identity metadata."""
|
||||
from src.core.provider_oauth_utils import (
|
||||
fetch_openai_account_name,
|
||||
parse_codex_id_token,
|
||||
)
|
||||
|
||||
def _read_non_empty_str(*values: Any) -> str | None:
|
||||
for value in values:
|
||||
if isinstance(value, str):
|
||||
normalized = value.strip()
|
||||
if normalized:
|
||||
return normalized
|
||||
return None
|
||||
|
||||
# Prefer explicit fields if token endpoint returns them directly.
|
||||
direct_account_id = _read_non_empty_str(
|
||||
token_response.get("account_id"),
|
||||
token_response.get("accountId"),
|
||||
token_response.get("chatgpt_account_id"),
|
||||
token_response.get("chatgptAccountId"),
|
||||
)
|
||||
if direct_account_id and not auth_config.get("account_id"):
|
||||
auth_config["account_id"] = direct_account_id
|
||||
|
||||
direct_account_user_id = _read_non_empty_str(
|
||||
token_response.get("account_user_id"),
|
||||
token_response.get("accountUserId"),
|
||||
token_response.get("chatgpt_account_user_id"),
|
||||
token_response.get("chatgptAccountUserId"),
|
||||
)
|
||||
if direct_account_user_id and not auth_config.get("account_user_id"):
|
||||
auth_config["account_user_id"] = direct_account_user_id
|
||||
|
||||
direct_plan_type = _read_non_empty_str(
|
||||
token_response.get("plan_type"),
|
||||
token_response.get("planType"),
|
||||
token_response.get("chatgpt_plan_type"),
|
||||
token_response.get("chatgptPlanType"),
|
||||
)
|
||||
if direct_plan_type and not auth_config.get("plan_type"):
|
||||
auth_config["plan_type"] = direct_plan_type
|
||||
|
||||
direct_user_id = _read_non_empty_str(
|
||||
token_response.get("user_id"),
|
||||
token_response.get("userId"),
|
||||
token_response.get("chatgpt_user_id"),
|
||||
token_response.get("chatgptUserId"),
|
||||
)
|
||||
if direct_user_id and not auth_config.get("user_id"):
|
||||
auth_config["user_id"] = direct_user_id
|
||||
|
||||
direct_email = _read_non_empty_str(token_response.get("email"))
|
||||
if direct_email and not auth_config.get("email"):
|
||||
auth_config["email"] = direct_email
|
||||
|
||||
direct_snapshot = _build_safe_codex_debug_snapshot(
|
||||
{
|
||||
"email": direct_email,
|
||||
"account_id": direct_account_id,
|
||||
"account_user_id": direct_account_user_id,
|
||||
"plan_type": direct_plan_type,
|
||||
"user_id": direct_user_id,
|
||||
}
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
"Codex enrich_auth_config: id_token_present={} access_token_present={} token_keys={}",
|
||||
bool(token_response.get("id_token") or token_response.get("idToken")),
|
||||
bool(token_response.get("access_token") or token_response.get("accessToken")),
|
||||
list(token_response.keys()),
|
||||
)
|
||||
logger.debug("Codex direct token response values: {}", direct_snapshot)
|
||||
|
||||
token_candidates = [
|
||||
("id_token", token_response.get("id_token")),
|
||||
("idToken", token_response.get("idToken")),
|
||||
("access_token", token_response.get("access_token")),
|
||||
("accessToken", token_response.get("accessToken")),
|
||||
]
|
||||
for source_name, token_payload in token_candidates:
|
||||
codex_info = parse_codex_id_token(token_payload)
|
||||
if not codex_info:
|
||||
continue
|
||||
logger.debug(
|
||||
"Codex parsed token values: source={} fields={} values={}",
|
||||
source_name,
|
||||
list(codex_info.keys()),
|
||||
_build_safe_codex_debug_snapshot(codex_info),
|
||||
)
|
||||
for key, value in codex_info.items():
|
||||
if not auth_config.get(key):
|
||||
auth_config[key] = value
|
||||
|
||||
account_id = _read_non_empty_str(auth_config.get("account_id"))
|
||||
if account_id:
|
||||
account_name = await fetch_openai_account_name(
|
||||
access_token,
|
||||
account_id,
|
||||
proxy_config=proxy_config,
|
||||
timeout_seconds=10.0,
|
||||
)
|
||||
logger.debug(
|
||||
"Codex account_name lookup: account_id={} resolved_account_name={}",
|
||||
account_id,
|
||||
account_name,
|
||||
)
|
||||
if account_name:
|
||||
auth_config["account_name"] = account_name
|
||||
|
||||
logger.debug(
|
||||
"Codex enrich_auth_config final metadata: {}",
|
||||
_build_safe_codex_debug_snapshot(auth_config),
|
||||
)
|
||||
|
||||
return auth_config
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unified Registration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def register_all() -> None:
|
||||
"""一次性注册 Codex 的所有 hooks 到各通用 registry。"""
|
||||
from src.core.api_format.capabilities import register_provider_default_body_rules
|
||||
from src.core.provider_oauth_utils import register_auth_enricher
|
||||
from src.services.model.upstream_fetcher import UpstreamModelsFetcherRegistry
|
||||
from src.services.provider.transport import register_transport_hook
|
||||
from src.services.provider.upstream_headers import register_upstream_headers_hook
|
||||
|
||||
# Transport
|
||||
register_transport_hook("codex", "openai:cli", build_codex_url)
|
||||
register_transport_hook("codex", "openai:compact", build_codex_url)
|
||||
register_upstream_headers_hook("codex", "openai:cli", build_codex_cli_headers)
|
||||
register_upstream_headers_hook("codex", "openai:compact", build_codex_compact_headers)
|
||||
|
||||
# Auth
|
||||
register_auth_enricher("codex", enrich_codex)
|
||||
|
||||
# Provider Format Capability:默认 body_rules
|
||||
from src.core.api_format.metadata import CODEX_DEFAULT_BODY_RULES
|
||||
|
||||
register_provider_default_body_rules("codex", "openai:cli", CODEX_DEFAULT_BODY_RULES)
|
||||
|
||||
# Export: Codex uses the default export builder (strip null + temp fields)
|
||||
# No need to register a custom one — the default in export.py suffices.
|
||||
|
||||
# Model Fetcher
|
||||
UpstreamModelsFetcherRegistry.register(
|
||||
provider_types=["codex"],
|
||||
fetcher=fetch_models_codex,
|
||||
)
|
||||
@@ -0,0 +1,5 @@
|
||||
"""GeminiCLI provider adapter package."""
|
||||
|
||||
from .plugin import register_all
|
||||
|
||||
__all__ = ["register_all"]
|
||||
@@ -0,0 +1,221 @@
|
||||
"""GeminiCLI upstream client helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import Any
|
||||
|
||||
from src.clients.http_client import HTTPClientPool
|
||||
from src.core.logger import logger
|
||||
from src.services.provider.adapters.gemini_cli.constants import (
|
||||
PROD_BASE_URL,
|
||||
get_v1internal_extra_headers,
|
||||
)
|
||||
from src.services.provider.adapters.gemini_cli.rust_http import (
|
||||
execute_gemini_cli_rust_http_request,
|
||||
)
|
||||
|
||||
_CODE_ASSIST_METADATA = {
|
||||
"ideType": "ANTIGRAVITY",
|
||||
"platform": "PLATFORM_UNSPECIFIED",
|
||||
"pluginType": "GEMINI",
|
||||
}
|
||||
|
||||
|
||||
def _extract_tier_raw(tier_obj: Any) -> str:
|
||||
"""Extract raw tier string from loadCodeAssist response objects."""
|
||||
if isinstance(tier_obj, str) and tier_obj.strip():
|
||||
return tier_obj.strip()
|
||||
if isinstance(tier_obj, dict):
|
||||
for key in ("id", "tierType"):
|
||||
value = tier_obj.get(key)
|
||||
if isinstance(value, str) and value.strip():
|
||||
return value.strip()
|
||||
return ""
|
||||
|
||||
|
||||
def extract_plan_type(data: dict[str, Any]) -> str | None:
|
||||
"""Best-effort normalized plan type for GeminiCLI OAuth accounts."""
|
||||
from src.core.oauth_plan import normalize_oauth_plan_type
|
||||
|
||||
for key in ("paidTier", "currentTier"):
|
||||
raw = _extract_tier_raw(data.get(key))
|
||||
normalized = normalize_oauth_plan_type(raw)
|
||||
if normalized:
|
||||
return normalized
|
||||
return None
|
||||
|
||||
|
||||
def extract_project_id(data: dict[str, Any]) -> str:
|
||||
"""Extract project_id from loadCodeAssist/onboardUser responses."""
|
||||
raw = data.get("cloudaicompanionProject")
|
||||
if isinstance(raw, str) and raw.strip():
|
||||
return raw.strip()
|
||||
if isinstance(raw, dict):
|
||||
project_id = raw.get("id", "")
|
||||
if isinstance(project_id, str) and project_id.strip():
|
||||
return project_id.strip()
|
||||
return ""
|
||||
|
||||
|
||||
def extract_tier_id(data: dict[str, Any]) -> str:
|
||||
"""Choose a tier ID for onboarding when the account is not activated."""
|
||||
allowed_tiers = data.get("allowedTiers")
|
||||
if not isinstance(allowed_tiers, list):
|
||||
return ""
|
||||
|
||||
for tier in allowed_tiers:
|
||||
if isinstance(tier, dict) and tier.get("isDefault") is True:
|
||||
tier_id = tier.get("id", "")
|
||||
if isinstance(tier_id, str) and tier_id.strip():
|
||||
return tier_id.strip()
|
||||
|
||||
for tier in allowed_tiers:
|
||||
if isinstance(tier, dict):
|
||||
tier_id = tier.get("id", "")
|
||||
if isinstance(tier_id, str) and tier_id.strip():
|
||||
return tier_id.strip()
|
||||
|
||||
return ""
|
||||
|
||||
|
||||
async def load_code_assist(
|
||||
access_token: str,
|
||||
proxy_config: dict[str, Any] | None = None,
|
||||
*,
|
||||
timeout_seconds: float = 10.0,
|
||||
) -> dict[str, Any]:
|
||||
"""Load GeminiCLI account metadata from v1internal:loadCodeAssist."""
|
||||
if not access_token:
|
||||
raise ValueError("missing access_token")
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {access_token}",
|
||||
"Content-Type": "application/json",
|
||||
**get_v1internal_extra_headers(),
|
||||
}
|
||||
url = f"{PROD_BASE_URL.rstrip('/')}/v1internal:loadCodeAssist"
|
||||
body = {"metadata": _CODE_ASSIST_METADATA}
|
||||
resp = await execute_gemini_cli_rust_http_request(
|
||||
method="POST",
|
||||
url=url,
|
||||
headers=headers,
|
||||
body=body,
|
||||
proxy_config=proxy_config,
|
||||
request_id="gemini-cli:load-code-assist",
|
||||
provider_api_format="gemini_cli:load_code_assist",
|
||||
timeout_seconds=timeout_seconds,
|
||||
content_type="application/json",
|
||||
)
|
||||
if resp is None:
|
||||
client = await HTTPClientPool.get_proxy_client(proxy_config)
|
||||
resp = await client.post(
|
||||
url,
|
||||
json=body,
|
||||
headers=headers,
|
||||
timeout=timeout_seconds,
|
||||
)
|
||||
if 200 <= resp.status_code < 300:
|
||||
data = resp.json()
|
||||
return data if isinstance(data, dict) else {}
|
||||
raise RuntimeError(
|
||||
f"loadCodeAssist failed: status={resp.status_code} body={resp.text[:200] if resp.text else ''}"
|
||||
)
|
||||
|
||||
|
||||
async def onboard_user(
|
||||
access_token: str,
|
||||
*,
|
||||
tier_id: str,
|
||||
proxy_config: dict[str, Any] | None = None,
|
||||
timeout_seconds: float = 30.0,
|
||||
max_attempts: int = 5,
|
||||
poll_interval: float = 2.0,
|
||||
) -> str:
|
||||
"""Activate GeminiCLI user and fetch project_id via v1internal:onboardUser."""
|
||||
if not access_token:
|
||||
raise ValueError("missing access_token")
|
||||
if not tier_id:
|
||||
raise ValueError("missing tier_id")
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {access_token}",
|
||||
"Content-Type": "application/json",
|
||||
**get_v1internal_extra_headers(),
|
||||
}
|
||||
body = {
|
||||
"tierId": tier_id,
|
||||
"metadata": _CODE_ASSIST_METADATA,
|
||||
}
|
||||
url = f"{PROD_BASE_URL.rstrip('/')}/v1internal:onboardUser"
|
||||
|
||||
for attempt in range(1, max_attempts + 1):
|
||||
resp = await execute_gemini_cli_rust_http_request(
|
||||
method="POST",
|
||||
url=url,
|
||||
headers=headers,
|
||||
body=body,
|
||||
proxy_config=proxy_config,
|
||||
request_id=f"gemini-cli:onboard-user:{tier_id}:{attempt}",
|
||||
provider_api_format="gemini_cli:onboard_user",
|
||||
timeout_seconds=timeout_seconds,
|
||||
content_type="application/json",
|
||||
)
|
||||
if resp is None:
|
||||
client = await HTTPClientPool.get_proxy_client(proxy_config)
|
||||
resp = await client.post(url, json=body, headers=headers, timeout=timeout_seconds)
|
||||
if not (200 <= resp.status_code < 300):
|
||||
raise RuntimeError(
|
||||
f"onboardUser failed: status={resp.status_code} body={resp.text[:200] if resp.text else ''}"
|
||||
)
|
||||
|
||||
data = resp.json()
|
||||
if not isinstance(data, dict):
|
||||
raise RuntimeError(f"onboardUser: unexpected response type: {type(data)}")
|
||||
|
||||
if data.get("done") is True:
|
||||
response_data = data.get("response")
|
||||
if isinstance(response_data, dict):
|
||||
return extract_project_id(response_data)
|
||||
return ""
|
||||
|
||||
if attempt < max_attempts:
|
||||
await asyncio.sleep(poll_interval)
|
||||
|
||||
raise RuntimeError(f"onboardUser: timeout after {max_attempts} attempts")
|
||||
|
||||
|
||||
async def enrich_project_id(
|
||||
access_token: str,
|
||||
proxy_config: dict[str, Any] | None = None,
|
||||
) -> str | None:
|
||||
"""Best-effort project_id resolution for GeminiCLI OAuth keys."""
|
||||
code_assist = await load_code_assist(access_token, proxy_config=proxy_config)
|
||||
project_id = extract_project_id(code_assist)
|
||||
if project_id:
|
||||
return project_id
|
||||
|
||||
tier_id = extract_tier_id(code_assist)
|
||||
if tier_id:
|
||||
try:
|
||||
project_id = await onboard_user(
|
||||
access_token,
|
||||
tier_id=tier_id,
|
||||
proxy_config=proxy_config,
|
||||
)
|
||||
if project_id:
|
||||
return project_id
|
||||
except Exception as exc:
|
||||
logger.warning("GeminiCLI onboardUser failed: {}", exc)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
__all__ = [
|
||||
"enrich_project_id",
|
||||
"extract_plan_type",
|
||||
"extract_project_id",
|
||||
"extract_tier_id",
|
||||
"load_code_assist",
|
||||
"onboard_user",
|
||||
]
|
||||
@@ -0,0 +1,25 @@
|
||||
"""GeminiCLI provider constants."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from src.config.settings import config
|
||||
from src.core.provider_templates.fixed_providers import FIXED_PROVIDERS
|
||||
from src.core.provider_templates.types import ProviderType
|
||||
|
||||
PROD_BASE_URL = FIXED_PROVIDERS[ProviderType.GEMINI_CLI].api_base_url
|
||||
V1INTERNAL_PATH_TEMPLATE = "/v1internal:{action}"
|
||||
|
||||
|
||||
def get_v1internal_extra_headers() -> dict[str, str]:
|
||||
"""Headers required by GeminiCLI upstream requests."""
|
||||
return {
|
||||
"Accept-Encoding": "identity",
|
||||
"User-Agent": config.internal_user_agent_gemini_cli,
|
||||
}
|
||||
|
||||
|
||||
__all__ = [
|
||||
"PROD_BASE_URL",
|
||||
"V1INTERNAL_PATH_TEMPLATE",
|
||||
"get_v1internal_extra_headers",
|
||||
]
|
||||
@@ -0,0 +1,96 @@
|
||||
"""GeminiCLI v1internal request envelope."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from src.services.provider.adapters.gemini_cli.constants import get_v1internal_extra_headers
|
||||
from src.services.provider.request_context import get_selected_base_url
|
||||
|
||||
|
||||
def wrap_v1internal_request(
|
||||
gemini_request: dict[str, Any],
|
||||
*,
|
||||
project_id: str,
|
||||
model: str,
|
||||
) -> dict[str, Any]:
|
||||
"""Wrap a Gemini request into GeminiCLI v1internal format."""
|
||||
inner_request = dict(gemini_request)
|
||||
inner_request.pop("model", None)
|
||||
inner_request.pop("stream", None)
|
||||
return {
|
||||
"model": model,
|
||||
"project": project_id,
|
||||
"request": inner_request,
|
||||
}
|
||||
|
||||
|
||||
class GeminiCliV1InternalEnvelope:
|
||||
name = "gemini_cli:v1internal"
|
||||
|
||||
def extra_headers(self) -> dict[str, str] | None:
|
||||
return get_v1internal_extra_headers()
|
||||
|
||||
def wrap_request(
|
||||
self,
|
||||
request_body: dict[str, Any],
|
||||
*,
|
||||
model: str,
|
||||
url_model: str | None,
|
||||
decrypted_auth_config: dict[str, Any] | None,
|
||||
) -> tuple[dict[str, Any], str | None]:
|
||||
_ = url_model
|
||||
project_id = (decrypted_auth_config or {}).get("project_id")
|
||||
if not isinstance(project_id, str) or not project_id:
|
||||
from src.core.exceptions import ProviderNotAvailableException
|
||||
|
||||
raise ProviderNotAvailableException(
|
||||
"GeminiCLI OAuth 配置缺少 project_id,请重新授权",
|
||||
provider_name="gemini_cli",
|
||||
upstream_response="missing auth_config.project_id",
|
||||
)
|
||||
|
||||
wrapped = wrap_v1internal_request(
|
||||
request_body,
|
||||
project_id=project_id,
|
||||
model=model,
|
||||
)
|
||||
return wrapped, None
|
||||
|
||||
def unwrap_response(self, data: Any) -> Any:
|
||||
if not isinstance(data, dict):
|
||||
return data
|
||||
|
||||
response_obj = data.get("response")
|
||||
if "candidates" not in data and isinstance(response_obj, dict):
|
||||
return response_obj
|
||||
|
||||
return data
|
||||
|
||||
def postprocess_unwrapped_response(self, *, model: str, data: Any) -> None:
|
||||
_ = model, data
|
||||
return
|
||||
|
||||
def capture_selected_base_url(self) -> str | None:
|
||||
return get_selected_base_url()
|
||||
|
||||
def on_http_status(self, *, base_url: str | None, status_code: int) -> None:
|
||||
_ = base_url, status_code
|
||||
return
|
||||
|
||||
def on_connection_error(self, *, base_url: str | None, exc: Exception) -> None:
|
||||
_ = base_url, exc
|
||||
return
|
||||
|
||||
def force_stream_rewrite(self) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
gemini_cli_v1internal_envelope = GeminiCliV1InternalEnvelope()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"GeminiCliV1InternalEnvelope",
|
||||
"gemini_cli_v1internal_envelope",
|
||||
"wrap_v1internal_request",
|
||||
]
|
||||
@@ -0,0 +1,158 @@
|
||||
"""GeminiCLI provider plugin — unified registration entry."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from typing import Any
|
||||
from urllib.parse import urlencode
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.services.provider.adapters.gemini_cli.constants import V1INTERNAL_PATH_TEMPLATE
|
||||
from src.services.provider.preset_models import get_preset_models
|
||||
from src.services.provider.request_context import set_selected_base_url
|
||||
|
||||
|
||||
async def fetch_models_gemini_cli(
|
||||
ctx: Any,
|
||||
timeout_seconds: float,
|
||||
) -> tuple[list[dict], list[str], bool, dict[str, Any] | None]:
|
||||
"""GeminiCLI model fetcher.
|
||||
|
||||
Upstream does not expose a stable public models endpoint for OAuth CLI access,
|
||||
so we currently return a curated preset model catalog and enrich account metadata
|
||||
from loadCodeAssist when possible.
|
||||
"""
|
||||
from src.services.provider.adapters.gemini_cli.client import (
|
||||
extract_plan_type,
|
||||
load_code_assist,
|
||||
)
|
||||
|
||||
models = get_preset_models("gemini_cli")
|
||||
upstream_metadata: dict[str, Any] | None = None
|
||||
|
||||
access_token = str(getattr(ctx, "api_key_value", "") or "").strip()
|
||||
if access_token:
|
||||
try:
|
||||
code_assist = await load_code_assist(
|
||||
access_token,
|
||||
proxy_config=getattr(ctx, "proxy_config", None),
|
||||
timeout_seconds=timeout_seconds,
|
||||
)
|
||||
provider_meta: dict[str, Any] = {"updated_at": int(time.time())}
|
||||
plan_type = extract_plan_type(code_assist)
|
||||
if plan_type:
|
||||
provider_meta["plan_type"] = plan_type
|
||||
project_id = (getattr(ctx, "auth_config", None) or {}).get("project_id")
|
||||
if isinstance(project_id, str) and project_id:
|
||||
provider_meta["project_id"] = project_id
|
||||
upstream_metadata = {"gemini_cli": provider_meta}
|
||||
except Exception as exc:
|
||||
logger.debug("GeminiCLI model metadata fetch failed: {}", exc)
|
||||
|
||||
return models, [], True, upstream_metadata
|
||||
|
||||
|
||||
def build_gemini_cli_url(
|
||||
endpoint: Any,
|
||||
*,
|
||||
is_stream: bool,
|
||||
effective_query_params: dict[str, Any],
|
||||
**_kwargs: Any,
|
||||
) -> str:
|
||||
"""Build GeminiCLI v1internal URL."""
|
||||
base_url = str(getattr(endpoint, "base_url", "") or "").rstrip("/")
|
||||
set_selected_base_url(base_url)
|
||||
|
||||
action = "streamGenerateContent" if is_stream else "generateContent"
|
||||
path = V1INTERNAL_PATH_TEMPLATE.format(action=action)
|
||||
url = f"{base_url}{path}"
|
||||
if is_stream:
|
||||
effective_query_params.setdefault("alt", "sse")
|
||||
if effective_query_params:
|
||||
query_string = urlencode(effective_query_params, doseq=True)
|
||||
if query_string:
|
||||
url = f"{url}?{query_string}"
|
||||
return url
|
||||
|
||||
|
||||
async def enrich_gemini_cli(
|
||||
auth_config: dict[str, Any],
|
||||
token_response: dict[str, Any],
|
||||
access_token: str,
|
||||
proxy_config: dict[str, Any] | None,
|
||||
) -> dict[str, Any]:
|
||||
"""GeminiCLI auth_config enrichment: email + project_id."""
|
||||
from src.core.provider_oauth_utils import fetch_google_email
|
||||
from src.services.provider.adapters.gemini_cli.client import (
|
||||
enrich_project_id,
|
||||
extract_plan_type,
|
||||
load_code_assist,
|
||||
)
|
||||
|
||||
if not auth_config.get("email"):
|
||||
email = await fetch_google_email(
|
||||
access_token,
|
||||
proxy_config=proxy_config,
|
||||
timeout_seconds=10.0,
|
||||
)
|
||||
if email:
|
||||
auth_config["email"] = email
|
||||
|
||||
try:
|
||||
code_assist = await load_code_assist(access_token, proxy_config=proxy_config)
|
||||
except Exception as exc:
|
||||
code_assist = None
|
||||
logger.warning("[enrich] GeminiCLI loadCodeAssist failed: {}", exc)
|
||||
|
||||
if code_assist and not auth_config.get("tier"):
|
||||
plan_type = extract_plan_type(code_assist)
|
||||
if plan_type:
|
||||
auth_config["tier"] = plan_type
|
||||
|
||||
if not auth_config.get("project_id"):
|
||||
try:
|
||||
project_id = (code_assist and code_assist.get("cloudaicompanionProject")) or None
|
||||
if isinstance(project_id, dict):
|
||||
project_id = project_id.get("id")
|
||||
if isinstance(project_id, str) and project_id.strip():
|
||||
auth_config["project_id"] = project_id.strip()
|
||||
else:
|
||||
project_id = await enrich_project_id(access_token, proxy_config=proxy_config)
|
||||
if project_id:
|
||||
auth_config["project_id"] = project_id
|
||||
logger.info("[enrich] GeminiCLI project_id: {}", project_id[:8] + "...")
|
||||
except Exception as exc:
|
||||
logger.warning("[enrich] GeminiCLI project_id enrichment failed: {}", exc)
|
||||
|
||||
return auth_config
|
||||
|
||||
|
||||
def register_all() -> None:
|
||||
"""Register all GeminiCLI hooks into shared registries."""
|
||||
from src.core.provider_oauth_utils import register_auth_enricher
|
||||
from src.services.model.upstream_fetcher import UpstreamModelsFetcherRegistry
|
||||
from src.services.provider.adapters.gemini_cli.envelope import gemini_cli_v1internal_envelope
|
||||
from src.services.provider.envelope import register_envelope
|
||||
from src.services.provider.transport import register_transport_hook
|
||||
|
||||
register_envelope("gemini_cli", "gemini:cli", gemini_cli_v1internal_envelope)
|
||||
register_envelope("gemini_cli", "gemini:chat", gemini_cli_v1internal_envelope)
|
||||
register_envelope("gemini_cli", "", gemini_cli_v1internal_envelope)
|
||||
|
||||
register_transport_hook("gemini_cli", "gemini:cli", build_gemini_cli_url)
|
||||
register_transport_hook("gemini_cli", "gemini:chat", build_gemini_cli_url)
|
||||
|
||||
register_auth_enricher("gemini_cli", enrich_gemini_cli)
|
||||
|
||||
UpstreamModelsFetcherRegistry.register(
|
||||
provider_types=["gemini_cli"],
|
||||
fetcher=fetch_models_gemini_cli,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"build_gemini_cli_url",
|
||||
"enrich_gemini_cli",
|
||||
"fetch_models_gemini_cli",
|
||||
"register_all",
|
||||
]
|
||||
@@ -0,0 +1,275 @@
|
||||
"""Gemini CLI quota / RESOURCE_EXHAUSTED helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
_DURATION_TOKEN_RE = re.compile(r"(\d+(?:\.\d+)?)([dhms])")
|
||||
_RESET_AFTER_RE = re.compile(r"reset after\s+([^.,;]+)", re.IGNORECASE)
|
||||
|
||||
|
||||
def _parse_json(error_text: str | None) -> dict[str, Any] | None:
|
||||
if not isinstance(error_text, str) or not error_text.strip():
|
||||
return None
|
||||
try:
|
||||
data = json.loads(error_text)
|
||||
except Exception:
|
||||
return None
|
||||
return data if isinstance(data, dict) else None
|
||||
|
||||
|
||||
def _parse_duration_seconds(raw: Any) -> int | None:
|
||||
if isinstance(raw, (int, float)):
|
||||
return max(1, int(raw))
|
||||
if not isinstance(raw, str):
|
||||
return None
|
||||
text = raw.strip().lower()
|
||||
if not text:
|
||||
return None
|
||||
|
||||
total_seconds = 0.0
|
||||
matched = False
|
||||
for amount_text, unit in _DURATION_TOKEN_RE.findall(text):
|
||||
matched = True
|
||||
amount = float(amount_text)
|
||||
if unit == "d":
|
||||
total_seconds += amount * 86400
|
||||
elif unit == "h":
|
||||
total_seconds += amount * 3600
|
||||
elif unit == "m":
|
||||
total_seconds += amount * 60
|
||||
elif unit == "s":
|
||||
total_seconds += amount
|
||||
if not matched:
|
||||
return None
|
||||
return max(1, int(total_seconds))
|
||||
|
||||
|
||||
def _iter_error_details(payload: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
error_obj = payload.get("error")
|
||||
if not isinstance(error_obj, dict):
|
||||
return []
|
||||
details = error_obj.get("details")
|
||||
if not isinstance(details, list):
|
||||
return []
|
||||
return [item for item in details if isinstance(item, dict)]
|
||||
|
||||
|
||||
def _error_status(payload: dict[str, Any]) -> str:
|
||||
error_obj = payload.get("error")
|
||||
if not isinstance(error_obj, dict):
|
||||
return ""
|
||||
status = error_obj.get("status")
|
||||
return status.strip() if isinstance(status, str) else ""
|
||||
|
||||
|
||||
def _error_message(payload: dict[str, Any]) -> str:
|
||||
error_obj = payload.get("error")
|
||||
if not isinstance(error_obj, dict):
|
||||
return ""
|
||||
message = error_obj.get("message")
|
||||
return message.strip() if isinstance(message, str) else ""
|
||||
|
||||
|
||||
def _error_reason(payload: dict[str, Any]) -> str:
|
||||
for detail in _iter_error_details(payload):
|
||||
reason = detail.get("reason")
|
||||
if isinstance(reason, str) and reason.strip():
|
||||
return reason.strip()
|
||||
return ""
|
||||
|
||||
|
||||
def _looks_like_uuid(value: str | None) -> bool:
|
||||
if not isinstance(value, str):
|
||||
return False
|
||||
text = value.strip()
|
||||
if not text:
|
||||
return False
|
||||
try:
|
||||
UUID(text)
|
||||
except Exception:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def is_resource_exhausted_error(error_text: str | None) -> bool:
|
||||
payload = _parse_json(error_text)
|
||||
if payload is None:
|
||||
return False
|
||||
|
||||
status = _error_status(payload).upper()
|
||||
reason = _error_reason(payload).upper()
|
||||
if status == "RESOURCE_EXHAUSTED" or reason == "QUOTA_EXHAUSTED":
|
||||
return True
|
||||
|
||||
message = _error_message(payload).lower()
|
||||
return ("exhausted your capacity" in message) or ("quota" in message and "exhaust" in message)
|
||||
|
||||
|
||||
def parse_quota_reset_timestamp(error_text: str | None) -> int | None:
|
||||
payload = _parse_json(error_text)
|
||||
if payload is None:
|
||||
return None
|
||||
|
||||
for detail in _iter_error_details(payload):
|
||||
metadata = detail.get("metadata")
|
||||
if not isinstance(metadata, dict):
|
||||
continue
|
||||
raw = metadata.get("quotaResetTimeStamp") or metadata.get("quotaResetTimestamp")
|
||||
if not isinstance(raw, str) or not raw.strip():
|
||||
continue
|
||||
text = raw.strip()
|
||||
try:
|
||||
if text.endswith("Z"):
|
||||
text = text[:-1] + "+00:00"
|
||||
parsed = datetime.fromisoformat(text)
|
||||
if parsed.tzinfo is None:
|
||||
parsed = parsed.replace(tzinfo=timezone.utc)
|
||||
return int(parsed.timestamp())
|
||||
except Exception:
|
||||
continue
|
||||
return None
|
||||
|
||||
|
||||
def parse_quota_reset_delay_seconds(error_text: str | None) -> int | None:
|
||||
payload = _parse_json(error_text)
|
||||
if payload is None:
|
||||
return None
|
||||
|
||||
for detail in _iter_error_details(payload):
|
||||
metadata = detail.get("metadata")
|
||||
if not isinstance(metadata, dict):
|
||||
continue
|
||||
delay = metadata.get("quotaResetDelay")
|
||||
parsed = _parse_duration_seconds(delay)
|
||||
if parsed is not None:
|
||||
return parsed
|
||||
return None
|
||||
|
||||
|
||||
def parse_quota_reset_message_seconds(error_text: str | None) -> int | None:
|
||||
payload = _parse_json(error_text)
|
||||
if payload is None:
|
||||
return None
|
||||
|
||||
message = _error_message(payload)
|
||||
if not message:
|
||||
return None
|
||||
|
||||
matched = _RESET_AFTER_RE.search(message)
|
||||
if not matched:
|
||||
return None
|
||||
|
||||
return _parse_duration_seconds(matched.group(1))
|
||||
|
||||
|
||||
def extract_error_model_name(error_text: str | None, *, fallback: str | None = None) -> str | None:
|
||||
payload = _parse_json(error_text)
|
||||
if payload is not None:
|
||||
for detail in _iter_error_details(payload):
|
||||
metadata = detail.get("metadata")
|
||||
if not isinstance(metadata, dict):
|
||||
continue
|
||||
model = metadata.get("model")
|
||||
if isinstance(model, str) and model.strip():
|
||||
return model.strip()
|
||||
|
||||
fallback_text = str(fallback or "").strip()
|
||||
if fallback_text and not _looks_like_uuid(fallback_text):
|
||||
return fallback_text
|
||||
return None
|
||||
|
||||
|
||||
def extract_quota_cooldown_seconds(
|
||||
error_text: str | None, *, now_ts: int | None = None
|
||||
) -> int | None:
|
||||
now = int(now_ts or time.time())
|
||||
|
||||
reset_at = parse_quota_reset_timestamp(error_text)
|
||||
if reset_at is not None:
|
||||
return max(1, reset_at - now)
|
||||
|
||||
delay = parse_quota_reset_delay_seconds(error_text)
|
||||
if delay is not None:
|
||||
return max(1, delay)
|
||||
|
||||
message_delay = parse_quota_reset_message_seconds(error_text)
|
||||
if message_delay is not None:
|
||||
return max(1, message_delay)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def build_quota_exhausted_metadata(
|
||||
*,
|
||||
model_name: str,
|
||||
error_text: str | None,
|
||||
current_namespace: dict[str, Any] | None = None,
|
||||
now_ts: int | None = None,
|
||||
) -> dict[str, Any] | None:
|
||||
normalized_model = str(model_name or "").strip()
|
||||
if not normalized_model:
|
||||
return None
|
||||
if not is_resource_exhausted_error(error_text):
|
||||
return None
|
||||
|
||||
now = int(now_ts or time.time())
|
||||
reset_at = parse_quota_reset_timestamp(error_text)
|
||||
if reset_at is None:
|
||||
delay = parse_quota_reset_delay_seconds(error_text)
|
||||
if delay is not None:
|
||||
reset_at = now + delay
|
||||
if reset_at is None:
|
||||
message_delay = parse_quota_reset_message_seconds(error_text)
|
||||
if message_delay is not None:
|
||||
reset_at = now + message_delay
|
||||
if reset_at is None:
|
||||
return None
|
||||
|
||||
payload = _parse_json(error_text) or {}
|
||||
namespace = dict(current_namespace) if isinstance(current_namespace, dict) else {}
|
||||
quota_by_model_raw = namespace.get("quota_by_model")
|
||||
quota_by_model = dict(quota_by_model_raw) if isinstance(quota_by_model_raw, dict) else {}
|
||||
model_entry_raw = quota_by_model.get(normalized_model)
|
||||
model_entry = dict(model_entry_raw) if isinstance(model_entry_raw, dict) else {}
|
||||
|
||||
model_entry["is_exhausted"] = True
|
||||
model_entry["remaining_fraction"] = 0.0
|
||||
model_entry["used_percent"] = 100.0
|
||||
model_entry["updated_at"] = now
|
||||
|
||||
model_entry["reset_at"] = reset_at
|
||||
model_entry["reset_time"] = datetime.fromtimestamp(reset_at, timezone.utc).isoformat()
|
||||
model_entry["reset_seconds"] = max(0, reset_at - now)
|
||||
|
||||
reason = _error_reason(payload) or _error_status(payload) or _error_message(payload)
|
||||
if reason:
|
||||
model_entry["reason"] = reason
|
||||
|
||||
quota_by_model[normalized_model] = model_entry
|
||||
namespace["quota_by_model"] = quota_by_model
|
||||
namespace["updated_at"] = now
|
||||
|
||||
status = _error_status(payload)
|
||||
if status:
|
||||
namespace["last_error_status"] = status
|
||||
if reason:
|
||||
namespace["last_error_reason"] = reason
|
||||
|
||||
return {"gemini_cli": namespace}
|
||||
|
||||
|
||||
__all__ = [
|
||||
"extract_error_model_name",
|
||||
"build_quota_exhausted_metadata",
|
||||
"extract_quota_cooldown_seconds",
|
||||
"is_resource_exhausted_error",
|
||||
"parse_quota_reset_message_seconds",
|
||||
"parse_quota_reset_delay_seconds",
|
||||
"parse_quota_reset_timestamp",
|
||||
]
|
||||
@@ -0,0 +1,108 @@
|
||||
"""Shared Rust executor HTTP helper for Gemini CLI side calls."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from src.config.settings import config
|
||||
from src.core.logger import logger
|
||||
|
||||
|
||||
async def execute_gemini_cli_rust_http_request(
|
||||
*,
|
||||
method: str,
|
||||
url: str,
|
||||
headers: dict[str, str],
|
||||
body: Any,
|
||||
proxy_config: dict[str, Any] | None,
|
||||
request_id: str,
|
||||
provider_api_format: str,
|
||||
timeout_seconds: float,
|
||||
content_type: str | None = None,
|
||||
) -> httpx.Response | None:
|
||||
from src.services.request.execution_runtime_plan import (
|
||||
ExecutionPlan,
|
||||
ExecutionPlanBody,
|
||||
ExecutionPlanTimeouts,
|
||||
build_execution_plan_body,
|
||||
build_proxy_snapshot,
|
||||
)
|
||||
from src.services.request.execution_runtime_client import (
|
||||
ExecutionRuntimeClient,
|
||||
ExecutionRuntimeClientError,
|
||||
)
|
||||
|
||||
if config.execution_runtime_backend != "rust":
|
||||
return None
|
||||
|
||||
final_headers = dict(headers)
|
||||
if (
|
||||
body is not None
|
||||
and content_type
|
||||
and not any(str(key).lower() == "content-type" for key in final_headers)
|
||||
):
|
||||
final_headers["content-type"] = content_type
|
||||
|
||||
timeout_ms = max(int(timeout_seconds * 1000), 1_000)
|
||||
|
||||
try:
|
||||
proxy_snapshot = await build_proxy_snapshot(proxy_config, label="GeminiCLI")
|
||||
result = await ExecutionRuntimeClient().execute_sync_json(
|
||||
ExecutionPlan(
|
||||
request_id=request_id,
|
||||
candidate_id=None,
|
||||
provider_name="gemini_cli",
|
||||
provider_id="",
|
||||
endpoint_id="",
|
||||
key_id="",
|
||||
method=method,
|
||||
url=url,
|
||||
headers=final_headers,
|
||||
body=(
|
||||
build_execution_plan_body(body, content_type=content_type)
|
||||
if body is not None
|
||||
else ExecutionPlanBody()
|
||||
),
|
||||
stream=False,
|
||||
provider_api_format=provider_api_format,
|
||||
client_api_format=provider_api_format,
|
||||
model_name="gemini_cli",
|
||||
content_type=content_type,
|
||||
proxy=proxy_snapshot,
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=timeout_ms,
|
||||
read_ms=timeout_ms,
|
||||
write_ms=timeout_ms,
|
||||
pool_ms=timeout_ms,
|
||||
total_ms=timeout_ms,
|
||||
),
|
||||
)
|
||||
)
|
||||
except (ExecutionRuntimeClientError, httpx.HTTPError, json.JSONDecodeError) as exc:
|
||||
logger.warning("GeminiCLI Rust HTTP fallback {} {}: {}", method, url, exc)
|
||||
return None
|
||||
except Exception as exc:
|
||||
logger.warning("GeminiCLI Rust HTTP unexpected fallback {} {}: {}", method, 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(method, url, headers=final_headers),
|
||||
headers=response_headers,
|
||||
content=response_body,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["execute_gemini_cli_rust_http_request"]
|
||||
@@ -0,0 +1,3 @@
|
||||
"""Kiro provider adapter."""
|
||||
|
||||
__all__ = []
|
||||
@@ -0,0 +1,95 @@
|
||||
"""Kiro adapter constants.
|
||||
|
||||
Kiro upstream uses AWS Event Stream (binary frames) for streaming responses.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import platform
|
||||
|
||||
AWS_EVENTSTREAM_CONTENT_TYPE = "application/vnd.amazon.eventstream"
|
||||
|
||||
# Kiro API endpoints
|
||||
KIRO_GENERATE_ASSISTANT_PATH = "/generateAssistantResponse"
|
||||
KIRO_USAGE_LIMITS_PATH = "/getUsageLimits"
|
||||
|
||||
# Default AWS region when not specified in credentials
|
||||
DEFAULT_REGION = "us-east-1"
|
||||
|
||||
# Default client fingerprints used in headers (best-effort)
|
||||
DEFAULT_KIRO_VERSION = "0.8.0"
|
||||
DEFAULT_NODE_VERSION = "22.21.1"
|
||||
|
||||
|
||||
def _detect_system_version() -> str:
|
||||
system = platform.system().lower() or "other"
|
||||
release = platform.release() or "unknown"
|
||||
# Match KiroIDE style: darwin#24.6.0, windows#10.0.22631, linux#6.8.0-...
|
||||
return f"{system}#{release}"
|
||||
|
||||
|
||||
DEFAULT_SYSTEM_VERSION = _detect_system_version()
|
||||
|
||||
# Header constants
|
||||
KIRO_AGENT_MODE = "vibe"
|
||||
CODEWHISPERER_OPTOUT = "true"
|
||||
|
||||
# aws-sdk-js versions observed in kiro.rs
|
||||
AWS_SDK_JS_MAIN_VERSION = "1.0.27"
|
||||
AWS_SDK_JS_USAGE_VERSION = "1.0.0"
|
||||
|
||||
# Claude model context window used by kiro.rs to convert contextUsage percentage -> tokens
|
||||
CONTEXT_WINDOW_TOKENS = 200_000
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Chunked-write policy injected into tool descriptions and system prompt
|
||||
# ---------------------------------------------------------------------------
|
||||
# Kiro upstream has lower per-message size limits than standard Claude.
|
||||
# We inject instructions for Write/Edit tools and a system-level policy
|
||||
# so the model splits large writes into smaller chunks automatically.
|
||||
|
||||
WRITE_TOOL_DESCRIPTION_SUFFIX = (
|
||||
"- IMPORTANT: If the content to write exceeds 150 lines, you MUST only write "
|
||||
"the first 50 lines using this tool, then use `Edit` tool to append the "
|
||||
"remaining content in chunks of no more than 50 lines each. If needed, leave "
|
||||
"a unique placeholder to help append content. Do NOT attempt to write all "
|
||||
"content at once."
|
||||
)
|
||||
|
||||
EDIT_TOOL_DESCRIPTION_SUFFIX = (
|
||||
"- IMPORTANT: If the `new_string` content exceeds 50 lines, you MUST split "
|
||||
"it into multiple Edit calls, each replacing no more than 50 lines at a time. "
|
||||
"If used to append content, leave a unique placeholder to help append content. "
|
||||
"On the final chunk, do NOT include the placeholder."
|
||||
)
|
||||
|
||||
TOOL_DESCRIPTION_SUFFIXES: dict[str, str] = {
|
||||
"Write": WRITE_TOOL_DESCRIPTION_SUFFIX,
|
||||
"Edit": EDIT_TOOL_DESCRIPTION_SUFFIX,
|
||||
}
|
||||
|
||||
SYSTEM_CHUNKED_POLICY = (
|
||||
"When the Write or Edit tool has content size limits, always comply silently. "
|
||||
"Never suggest bypassing these limits via alternative tools. "
|
||||
"Never ask the user whether to switch approaches. "
|
||||
"Complete all chunked operations without commentary."
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AWS_EVENTSTREAM_CONTENT_TYPE",
|
||||
"AWS_SDK_JS_MAIN_VERSION",
|
||||
"AWS_SDK_JS_USAGE_VERSION",
|
||||
"CODEWHISPERER_OPTOUT",
|
||||
"CONTEXT_WINDOW_TOKENS",
|
||||
"DEFAULT_KIRO_VERSION",
|
||||
"DEFAULT_NODE_VERSION",
|
||||
"DEFAULT_REGION",
|
||||
"DEFAULT_SYSTEM_VERSION",
|
||||
"EDIT_TOOL_DESCRIPTION_SUFFIX",
|
||||
"KIRO_AGENT_MODE",
|
||||
"KIRO_GENERATE_ASSISTANT_PATH",
|
||||
"KIRO_USAGE_LIMITS_PATH",
|
||||
"SYSTEM_CHUNKED_POLICY",
|
||||
"TOOL_DESCRIPTION_SUFFIXES",
|
||||
"WRITE_TOOL_DESCRIPTION_SUFFIX",
|
||||
]
|
||||
@@ -0,0 +1,82 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import contextvars
|
||||
from dataclasses import dataclass, replace
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class KiroRequestContext:
|
||||
"""Per-request context for the Kiro adapter.
|
||||
|
||||
This bridges data from `KiroEnvelope.wrap_request()` (which receives the
|
||||
decrypted auth_config + original request body) to other layers that only
|
||||
expose parameterless hooks (extra_headers) or transport hooks.
|
||||
"""
|
||||
|
||||
region: str
|
||||
machine_id: str
|
||||
kiro_version: str | None = None
|
||||
system_version: str | None = None
|
||||
node_version: str | None = None
|
||||
thinking_enabled: bool = False
|
||||
last_http_status: int | None = None
|
||||
last_http_error_category: str | None = None
|
||||
last_connection_error_category: str | None = None
|
||||
last_connection_error_summary: str | None = None
|
||||
|
||||
|
||||
_kiro_request_context: contextvars.ContextVar[KiroRequestContext | None] = contextvars.ContextVar(
|
||||
"kiro_request_context",
|
||||
default=None,
|
||||
)
|
||||
|
||||
|
||||
def set_kiro_request_context(ctx: KiroRequestContext | None) -> None:
|
||||
_kiro_request_context.set(ctx)
|
||||
|
||||
|
||||
def get_kiro_request_context() -> KiroRequestContext | None:
|
||||
return _kiro_request_context.get()
|
||||
|
||||
|
||||
def update_kiro_http_status(
|
||||
*,
|
||||
status_code: int,
|
||||
category: str,
|
||||
) -> None:
|
||||
ctx = get_kiro_request_context()
|
||||
if ctx is None:
|
||||
return
|
||||
set_kiro_request_context(
|
||||
replace(
|
||||
ctx,
|
||||
last_http_status=int(status_code),
|
||||
last_http_error_category=str(category),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def update_kiro_connection_error(
|
||||
*,
|
||||
category: str,
|
||||
summary: str,
|
||||
) -> None:
|
||||
ctx = get_kiro_request_context()
|
||||
if ctx is None:
|
||||
return
|
||||
set_kiro_request_context(
|
||||
replace(
|
||||
ctx,
|
||||
last_connection_error_category=str(category),
|
||||
last_connection_error_summary=str(summary),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"KiroRequestContext",
|
||||
"get_kiro_request_context",
|
||||
"set_kiro_request_context",
|
||||
"update_kiro_connection_error",
|
||||
"update_kiro_http_status",
|
||||
]
|
||||
613
_deprecated_py_src/services/provider/adapters/kiro/converter.py
Normal file
613
_deprecated_py_src/services/provider/adapters/kiro/converter.py
Normal file
@@ -0,0 +1,613 @@
|
||||
"""Claude Messages -> Kiro ConversationState converter (best-effort).
|
||||
|
||||
This mirrors `kiro.rs/src/anthropic/converter.rs` but focuses on the fields
|
||||
needed by generateAssistantResponse.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.services.provider.adapters.kiro.constants import (
|
||||
SYSTEM_CHUNKED_POLICY as _SYSTEM_CHUNKED_POLICY,
|
||||
)
|
||||
from src.services.provider.adapters.kiro.constants import (
|
||||
TOOL_DESCRIPTION_SUFFIXES as _TOOL_DESCRIPTION_SUFFIXES,
|
||||
)
|
||||
|
||||
|
||||
def map_model(model: str) -> str | None:
|
||||
"""Pass through the model name as-is to Kiro upstream."""
|
||||
raw = str(model or "").strip()
|
||||
return raw or None
|
||||
|
||||
|
||||
def _extract_session_id(user_id: str) -> str | None:
|
||||
text = str(user_id or "")
|
||||
pos = text.find("session_")
|
||||
if pos < 0:
|
||||
return None
|
||||
session_part = text[pos + 8 :]
|
||||
if len(session_part) < 36:
|
||||
return None
|
||||
candidate = session_part[:36]
|
||||
if candidate.count("-") != 4:
|
||||
return None
|
||||
return candidate
|
||||
|
||||
|
||||
def _generate_thinking_prefix(request_body: dict[str, Any]) -> str | None:
|
||||
thinking = request_body.get("thinking")
|
||||
if not isinstance(thinking, dict):
|
||||
return None
|
||||
|
||||
thinking_type = str(thinking.get("type") or "").strip()
|
||||
if thinking_type == "enabled":
|
||||
budget = thinking.get("budget_tokens")
|
||||
try:
|
||||
budget_i = int(budget) if budget is not None else 0
|
||||
except Exception:
|
||||
budget_i = 0
|
||||
return (
|
||||
f"<thinking_mode>enabled</thinking_mode>"
|
||||
f"<max_thinking_length>{budget_i}</max_thinking_length>"
|
||||
)
|
||||
|
||||
if thinking_type == "adaptive":
|
||||
output_cfg = request_body.get("output_config")
|
||||
effort = "high"
|
||||
if isinstance(output_cfg, dict):
|
||||
eff = output_cfg.get("effort")
|
||||
if isinstance(eff, str) and eff.strip():
|
||||
effort = eff.strip()
|
||||
return (
|
||||
f"<thinking_mode>adaptive</thinking_mode>"
|
||||
f"<thinking_effort>{effort}</thinking_effort>"
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _has_thinking_tags(content: str) -> bool:
|
||||
return "<thinking_mode>" in content or "<max_thinking_length>" in content
|
||||
|
||||
|
||||
def _system_to_text(system: Any) -> str:
|
||||
if system is None:
|
||||
return ""
|
||||
if isinstance(system, str):
|
||||
return system
|
||||
if isinstance(system, list):
|
||||
parts: list[str] = []
|
||||
for item in system:
|
||||
if isinstance(item, dict):
|
||||
t = item.get("text")
|
||||
if isinstance(t, str) and t:
|
||||
parts.append(t)
|
||||
else:
|
||||
# best-effort
|
||||
try:
|
||||
parts.append(str(item))
|
||||
except Exception:
|
||||
pass
|
||||
return "\n".join([p for p in parts if p])
|
||||
return ""
|
||||
|
||||
|
||||
def _get_image_format(media_type: str | None) -> str | None:
|
||||
if not isinstance(media_type, str) or "/" not in media_type:
|
||||
return None
|
||||
prefix, suffix = media_type.split("/", 1)
|
||||
if prefix != "image":
|
||||
return None
|
||||
suffix = suffix.strip().lower()
|
||||
if suffix in {"jpeg", "png", "gif", "webp"}:
|
||||
return suffix
|
||||
if suffix == "jpg":
|
||||
return "jpeg"
|
||||
return None
|
||||
|
||||
|
||||
def _process_message_content(
|
||||
content: Any,
|
||||
) -> tuple[str, list[dict[str, Any]], list[dict[str, Any]]]:
|
||||
"""Extract text/images/tool_results from a Claude content field."""
|
||||
text_parts: list[str] = []
|
||||
images: list[dict[str, Any]] = []
|
||||
tool_results: list[dict[str, Any]] = []
|
||||
|
||||
if isinstance(content, str):
|
||||
if content:
|
||||
text_parts.append(content)
|
||||
return "".join(text_parts), images, tool_results
|
||||
|
||||
if not isinstance(content, list):
|
||||
return "".join(text_parts), images, tool_results
|
||||
|
||||
for block in content:
|
||||
if not isinstance(block, dict):
|
||||
continue
|
||||
|
||||
btype = str(block.get("type") or "").strip()
|
||||
|
||||
if btype == "text":
|
||||
text = block.get("text")
|
||||
if isinstance(text, str) and text:
|
||||
text_parts.append(text)
|
||||
continue
|
||||
|
||||
if btype == "image":
|
||||
source = block.get("source")
|
||||
if not isinstance(source, dict):
|
||||
continue
|
||||
media_type = source.get("media_type") or source.get("mediaType")
|
||||
fmt = _get_image_format(media_type if isinstance(media_type, str) else None)
|
||||
data = source.get("data")
|
||||
if fmt and isinstance(data, str) and data:
|
||||
images.append({"format": fmt, "source": {"bytes": data}})
|
||||
continue
|
||||
|
||||
if btype == "tool_result":
|
||||
tool_use_id = block.get("tool_use_id") or block.get("toolUseId")
|
||||
if not isinstance(tool_use_id, str) or not tool_use_id.strip():
|
||||
continue
|
||||
|
||||
raw_content = block.get("content")
|
||||
if isinstance(raw_content, str):
|
||||
text = raw_content
|
||||
elif isinstance(raw_content, list):
|
||||
# Claude tool_result content blocks; keep only text parts.
|
||||
parts: list[str] = []
|
||||
for item in raw_content:
|
||||
if isinstance(item, dict) and item.get("type") == "text":
|
||||
t = item.get("text")
|
||||
if isinstance(t, str) and t:
|
||||
parts.append(t)
|
||||
text = "\n".join(parts)
|
||||
else:
|
||||
try:
|
||||
text = json.dumps(raw_content, ensure_ascii=False)
|
||||
except Exception:
|
||||
text = str(raw_content)
|
||||
|
||||
is_error = bool(block.get("is_error") or block.get("isError") or False)
|
||||
status = "error" if is_error else "success"
|
||||
|
||||
tool_results.append(
|
||||
{
|
||||
"toolUseId": tool_use_id.strip(),
|
||||
"content": [{"text": text or ""}],
|
||||
"status": status,
|
||||
"isError": bool(is_error),
|
||||
}
|
||||
)
|
||||
continue
|
||||
|
||||
return "".join(text_parts), images, tool_results
|
||||
|
||||
|
||||
def _clean_tool_schema(schema: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Recursively remove fields that Kiro API rejects.
|
||||
|
||||
Kiro returns 400 "Improperly formed request" when tool schemas contain
|
||||
``additionalProperties`` (any value) or empty ``required: []`` arrays.
|
||||
"""
|
||||
if not isinstance(schema, dict):
|
||||
return schema # type: ignore[return-value]
|
||||
result: dict[str, Any] = {}
|
||||
for key, value in schema.items():
|
||||
if key == "additionalProperties":
|
||||
continue
|
||||
if key == "required" and isinstance(value, list) and not value:
|
||||
continue
|
||||
if isinstance(value, dict):
|
||||
result[key] = _clean_tool_schema(value)
|
||||
elif isinstance(value, list):
|
||||
result[key] = [_clean_tool_schema(v) if isinstance(v, dict) else v for v in value]
|
||||
else:
|
||||
result[key] = value
|
||||
return result
|
||||
|
||||
|
||||
def _convert_tools(tools: Any) -> list[dict[str, Any]]:
|
||||
if not isinstance(tools, list):
|
||||
return []
|
||||
|
||||
out: list[dict[str, Any]] = []
|
||||
for t in tools:
|
||||
if not isinstance(t, dict):
|
||||
continue
|
||||
name = t.get("name")
|
||||
if not isinstance(name, str) or not name.strip():
|
||||
continue
|
||||
|
||||
description = t.get("description")
|
||||
description_str = description if isinstance(description, str) else ""
|
||||
|
||||
# Inject chunked-write instructions for Write/Edit tools.
|
||||
suffix = _TOOL_DESCRIPTION_SUFFIXES.get(name.strip())
|
||||
if suffix:
|
||||
description_str = f"{description_str}\n{suffix}" if description_str else suffix
|
||||
|
||||
if len(description_str) > 10000:
|
||||
description_str = description_str[:10000]
|
||||
|
||||
input_schema = t.get("input_schema") or t.get("inputSchema") or {}
|
||||
if not isinstance(input_schema, dict):
|
||||
input_schema = {}
|
||||
|
||||
input_schema = _clean_tool_schema(input_schema)
|
||||
|
||||
out.append(
|
||||
{
|
||||
"toolSpecification": {
|
||||
"name": name.strip(),
|
||||
"description": description_str,
|
||||
"inputSchema": {"json": input_schema},
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def _create_placeholder_tool(name: str) -> dict[str, Any]:
|
||||
return {
|
||||
"toolSpecification": {
|
||||
"name": name,
|
||||
"description": "Tool used in conversation history",
|
||||
"inputSchema": {
|
||||
"json": {
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def _convert_assistant_message(message: dict[str, Any]) -> dict[str, Any] | None:
|
||||
content = message.get("content")
|
||||
|
||||
tool_uses: list[dict[str, Any]] = []
|
||||
thinking_parts: list[str] = []
|
||||
text_parts: list[str] = []
|
||||
|
||||
if isinstance(content, str):
|
||||
if content:
|
||||
text_parts.append(content)
|
||||
elif isinstance(content, list):
|
||||
for block in content:
|
||||
if not isinstance(block, dict):
|
||||
continue
|
||||
btype = str(block.get("type") or "")
|
||||
if btype == "thinking":
|
||||
# Preserve thinking content so multi-turn context is not lost.
|
||||
t = block.get("thinking")
|
||||
if isinstance(t, str) and t:
|
||||
thinking_parts.append(t)
|
||||
elif btype == "text":
|
||||
t = block.get("text")
|
||||
if isinstance(t, str) and t:
|
||||
text_parts.append(t)
|
||||
elif btype == "tool_use":
|
||||
tool_use_id = block.get("id")
|
||||
name = block.get("name")
|
||||
if not isinstance(tool_use_id, str) or not tool_use_id.strip():
|
||||
continue
|
||||
if not isinstance(name, str) or not name.strip():
|
||||
continue
|
||||
inp = block.get("input")
|
||||
if not isinstance(inp, dict):
|
||||
inp = {}
|
||||
tool_uses.append(
|
||||
{
|
||||
"toolUseId": tool_use_id.strip(),
|
||||
"name": name.strip(),
|
||||
"input": inp,
|
||||
}
|
||||
)
|
||||
|
||||
# Combine thinking + text into final content.
|
||||
# Format: <thinking>...</thinking>\n\ntext
|
||||
thinking_str = "".join(thinking_parts)
|
||||
text_str = "".join(text_parts)
|
||||
|
||||
if thinking_str:
|
||||
if text_str:
|
||||
content_str = f"<thinking>{thinking_str}</thinking>\n\n{text_str}"
|
||||
else:
|
||||
content_str = f"<thinking>{thinking_str}</thinking>"
|
||||
else:
|
||||
content_str = text_str
|
||||
|
||||
if not content_str and tool_uses:
|
||||
content_str = " " # Kiro API requires non-empty content.
|
||||
|
||||
if not content_str and not tool_uses:
|
||||
return None
|
||||
|
||||
out: dict[str, Any] = {"content": content_str}
|
||||
if tool_uses:
|
||||
out["toolUses"] = tool_uses
|
||||
return out
|
||||
|
||||
|
||||
def convert_claude_messages_to_conversation_state(
|
||||
request_body: dict[str, Any],
|
||||
*,
|
||||
model: str,
|
||||
) -> dict[str, Any]:
|
||||
model_id = map_model(model)
|
||||
if not model_id:
|
||||
raise ValueError(f"kiro: model is required (got {model!r})")
|
||||
|
||||
messages = request_body.get("messages")
|
||||
if not isinstance(messages, list) or not messages:
|
||||
raise ValueError("kiro: empty messages")
|
||||
|
||||
conversation_id = None
|
||||
metadata = request_body.get("metadata")
|
||||
if isinstance(metadata, dict):
|
||||
user_id = metadata.get("user_id") or metadata.get("userId")
|
||||
if isinstance(user_id, str) and user_id:
|
||||
conversation_id = _extract_session_id(user_id)
|
||||
|
||||
if not conversation_id:
|
||||
conversation_id = str(uuid.uuid4())
|
||||
|
||||
agent_continuation_id = str(uuid.uuid4())
|
||||
|
||||
thinking_prefix = _generate_thinking_prefix(request_body)
|
||||
|
||||
history: list[dict[str, Any]] = []
|
||||
|
||||
# System injection: add as (user, assistant) pair.
|
||||
system_text = _system_to_text(request_body.get("system"))
|
||||
if system_text:
|
||||
# Append chunked-write policy so the model silently obeys tool limits.
|
||||
final_system = f"{system_text}\n{_SYSTEM_CHUNKED_POLICY}"
|
||||
history.append(
|
||||
{
|
||||
"userInputMessage": {
|
||||
"content": final_system,
|
||||
"modelId": model_id,
|
||||
"origin": "AI_EDITOR",
|
||||
}
|
||||
}
|
||||
)
|
||||
history.append(
|
||||
{"assistantResponseMessage": {"content": "I will follow these instructions."}}
|
||||
)
|
||||
|
||||
# Build history from messages.
|
||||
# If the last message is assistant, include it in history (Kiro currentMessage
|
||||
# must be user; we synthesise one). Otherwise the last user message becomes
|
||||
# currentMessage and everything before it goes into history.
|
||||
last_msg = messages[-1]
|
||||
last_is_assistant = (
|
||||
isinstance(last_msg, dict) and str(last_msg.get("role") or "") == "assistant"
|
||||
)
|
||||
|
||||
if last_is_assistant:
|
||||
# All messages go into history; we'll synthesise a currentMessage later.
|
||||
history_end_index = len(messages)
|
||||
else:
|
||||
history_end_index = max(len(messages) - 1, 0)
|
||||
|
||||
user_buffer: list[dict[str, Any]] = []
|
||||
|
||||
def _flush_user_buffer() -> dict[str, Any] | None:
|
||||
nonlocal user_buffer
|
||||
if not user_buffer:
|
||||
return None
|
||||
|
||||
parts: list[str] = []
|
||||
images: list[dict[str, Any]] = []
|
||||
tool_results: list[dict[str, Any]] = []
|
||||
|
||||
for msg in user_buffer:
|
||||
text, imgs, results = _process_message_content(msg.get("content"))
|
||||
if text:
|
||||
parts.append(text)
|
||||
images.extend(imgs)
|
||||
tool_results.extend(results)
|
||||
|
||||
user_buffer = []
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"content": "\n".join(parts),
|
||||
"modelId": model_id,
|
||||
"origin": "AI_EDITOR",
|
||||
}
|
||||
|
||||
if images:
|
||||
payload["images"] = images
|
||||
|
||||
if tool_results:
|
||||
payload["userInputMessageContext"] = {"toolResults": tool_results}
|
||||
|
||||
return {"userInputMessage": payload}
|
||||
|
||||
for i in range(history_end_index):
|
||||
msg = messages[i]
|
||||
if not isinstance(msg, dict):
|
||||
continue
|
||||
role = str(msg.get("role") or "")
|
||||
if role == "user":
|
||||
user_buffer.append(msg)
|
||||
continue
|
||||
if role == "assistant":
|
||||
user_item = _flush_user_buffer()
|
||||
if user_item is not None:
|
||||
history.append(user_item)
|
||||
elif not history or "assistantResponseMessage" in history[-1]:
|
||||
# No preceding user message: insert synthetic user message
|
||||
# to maintain alternating roles required by Kiro API.
|
||||
history.append(
|
||||
{
|
||||
"userInputMessage": {
|
||||
"content": "Continue.",
|
||||
"modelId": model_id,
|
||||
"origin": "AI_EDITOR",
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
assistant_item = _convert_assistant_message(msg)
|
||||
if assistant_item is not None:
|
||||
history.append({"assistantResponseMessage": assistant_item})
|
||||
continue
|
||||
|
||||
# trailing unpaired user messages in history
|
||||
tail_user = _flush_user_buffer()
|
||||
if tail_user is not None:
|
||||
history.append(tail_user)
|
||||
history.append({"assistantResponseMessage": {"content": "OK"}})
|
||||
|
||||
# Current message: last message as user input.
|
||||
if last_is_assistant:
|
||||
# Synthesise a minimal user continuation message.
|
||||
text_content = "Continue."
|
||||
images: list[dict[str, Any]] = []
|
||||
tool_results: list[dict[str, Any]] = []
|
||||
else:
|
||||
last = messages[-1]
|
||||
if not isinstance(last, dict) or str(last.get("role") or "") != "user":
|
||||
raise ValueError("kiro: last message must be user")
|
||||
text_content, images, tool_results = _process_message_content(last.get("content"))
|
||||
|
||||
tools = _convert_tools(request_body.get("tools"))
|
||||
|
||||
# Ensure tools referenced in history assistant toolUses are defined.
|
||||
# Also collect ids for tool_use / tool_result pairing validation.
|
||||
history_tool_names: set[str] = set()
|
||||
history_tool_results_ids: set[str] = set()
|
||||
history_tool_use_ids: set[str] = set()
|
||||
|
||||
for item in history:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
u = item.get("userInputMessage")
|
||||
if isinstance(u, dict):
|
||||
ctx = u.get("userInputMessageContext")
|
||||
if isinstance(ctx, dict):
|
||||
results = ctx.get("toolResults")
|
||||
if isinstance(results, list):
|
||||
for r in results:
|
||||
if isinstance(r, dict):
|
||||
tid = r.get("toolUseId")
|
||||
if isinstance(tid, str) and tid:
|
||||
history_tool_results_ids.add(tid)
|
||||
a = item.get("assistantResponseMessage")
|
||||
if isinstance(a, dict):
|
||||
uses = a.get("toolUses")
|
||||
if isinstance(uses, list):
|
||||
for tu in uses:
|
||||
if not isinstance(tu, dict):
|
||||
continue
|
||||
nm = tu.get("name")
|
||||
if isinstance(nm, str) and nm:
|
||||
history_tool_names.add(nm)
|
||||
tid = tu.get("toolUseId")
|
||||
if isinstance(tid, str) and tid:
|
||||
history_tool_use_ids.add(tid)
|
||||
|
||||
existing_tool_names = {
|
||||
str(t.get("toolSpecification", {}).get("name", "")).lower() for t in tools
|
||||
}
|
||||
|
||||
for tool_name in sorted(history_tool_names):
|
||||
if tool_name.lower() not in existing_tool_names:
|
||||
tools.append(_create_placeholder_tool(tool_name))
|
||||
|
||||
# Filter tool_results: only keep those with matching tool_use in history, and not duplicated.
|
||||
validated_tool_results: list[dict[str, Any]] = []
|
||||
current_tool_result_ids: set[str] = set()
|
||||
for r in tool_results:
|
||||
if not isinstance(r, dict):
|
||||
continue
|
||||
tid = r.get("toolUseId")
|
||||
if not isinstance(tid, str) or not tid:
|
||||
continue
|
||||
if tid not in history_tool_use_ids:
|
||||
continue
|
||||
if tid in history_tool_results_ids:
|
||||
continue
|
||||
validated_tool_results.append(r)
|
||||
current_tool_result_ids.add(tid)
|
||||
|
||||
# Remove orphaned tool_uses from history.
|
||||
# Kiro API requires every tool_use to have a matching tool_result; otherwise
|
||||
# it returns 400 Bad Request.
|
||||
orphaned_tool_use_ids = (
|
||||
history_tool_use_ids - history_tool_results_ids - current_tool_result_ids
|
||||
)
|
||||
if orphaned_tool_use_ids:
|
||||
logger.warning(
|
||||
"kiro: removing {} orphaned tool_use(s) from history: {}",
|
||||
len(orphaned_tool_use_ids),
|
||||
orphaned_tool_use_ids,
|
||||
)
|
||||
for item in history:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
a = item.get("assistantResponseMessage")
|
||||
if not isinstance(a, dict):
|
||||
continue
|
||||
uses = a.get("toolUses")
|
||||
if not isinstance(uses, list):
|
||||
continue
|
||||
filtered = [
|
||||
u
|
||||
for u in uses
|
||||
if not (
|
||||
isinstance(u, dict)
|
||||
and isinstance(u.get("toolUseId"), str)
|
||||
and u["toolUseId"] in orphaned_tool_use_ids
|
||||
)
|
||||
]
|
||||
if not filtered:
|
||||
a.pop("toolUses", None)
|
||||
elif len(filtered) != len(uses):
|
||||
a["toolUses"] = filtered
|
||||
|
||||
user_ctx: dict[str, Any] = {}
|
||||
if tools:
|
||||
user_ctx["tools"] = tools
|
||||
if validated_tool_results:
|
||||
user_ctx["toolResults"] = validated_tool_results
|
||||
|
||||
# Inject thinking tags into currentMessage (not history) so the
|
||||
# instruction applies to the current turn only.
|
||||
if thinking_prefix and not _has_thinking_tags(text_content):
|
||||
text_content = f"{thinking_prefix}\n{text_content}"
|
||||
|
||||
user_input: dict[str, Any] = {
|
||||
"userInputMessageContext": user_ctx,
|
||||
"content": text_content,
|
||||
"modelId": model_id,
|
||||
"origin": "AI_EDITOR",
|
||||
}
|
||||
if images:
|
||||
user_input["images"] = images
|
||||
|
||||
conversation_state = {
|
||||
"agentContinuationId": agent_continuation_id,
|
||||
"agentTaskType": "vibe",
|
||||
"chatTriggerType": "MANUAL",
|
||||
"currentMessage": {"userInputMessage": user_input},
|
||||
"conversationId": conversation_id,
|
||||
"history": history,
|
||||
}
|
||||
|
||||
return conversation_state
|
||||
|
||||
|
||||
__all__ = [
|
||||
"convert_claude_messages_to_conversation_state",
|
||||
"map_model",
|
||||
]
|
||||
121
_deprecated_py_src/services/provider/adapters/kiro/envelope.py
Normal file
121
_deprecated_py_src/services/provider/adapters/kiro/envelope.py
Normal file
@@ -0,0 +1,121 @@
|
||||
"""Kiro provider envelope.
|
||||
|
||||
Kiro upstream is not Claude wire-compatible:
|
||||
- Request: wrap Claude Messages body into Kiro `conversationState` request.
|
||||
- Stream response: handled by StreamProcessor via binary EventStream rewrite.
|
||||
|
||||
We use contextvars to pass request-scoped values (region, machine_id, thinking)
|
||||
from wrap_request() to extra_headers() and transport hook.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.services.provider.adapters.kiro.context import KiroRequestContext, set_kiro_request_context
|
||||
from src.services.provider.adapters.kiro.error_enhancer import (
|
||||
classify_kiro_connection_error,
|
||||
classify_kiro_http_status,
|
||||
extract_kiro_http_error_text,
|
||||
summarize_kiro_connection_error,
|
||||
)
|
||||
from src.services.provider.adapters.kiro.headers import build_generate_assistant_headers
|
||||
from src.services.provider.adapters.kiro.models.credentials import KiroAuthConfig
|
||||
from src.services.provider.adapters.kiro.request import (
|
||||
build_kiro_request_context,
|
||||
build_kiro_request_payload,
|
||||
)
|
||||
from src.services.provider.request_context import get_selected_base_url
|
||||
|
||||
|
||||
class KiroEnvelope:
|
||||
name = "kiro:generateAssistantResponse"
|
||||
|
||||
def extra_headers(self) -> dict[str, str] | None:
|
||||
# Called after wrap_request(); relies on KiroRequestContext.
|
||||
from src.services.provider.adapters.kiro.context import get_kiro_request_context
|
||||
|
||||
ctx = get_kiro_request_context()
|
||||
if ctx is None:
|
||||
return None
|
||||
|
||||
host = f"q.{ctx.region}.amazonaws.com"
|
||||
return build_generate_assistant_headers(
|
||||
host=host,
|
||||
machine_id=ctx.machine_id,
|
||||
kiro_version=ctx.kiro_version,
|
||||
system_version=ctx.system_version,
|
||||
node_version=ctx.node_version,
|
||||
)
|
||||
|
||||
def wrap_request(
|
||||
self,
|
||||
request_body: dict[str, Any],
|
||||
*,
|
||||
model: str,
|
||||
url_model: str | None,
|
||||
decrypted_auth_config: dict[str, Any] | None,
|
||||
) -> tuple[dict[str, Any], str | None]:
|
||||
cfg = KiroAuthConfig.from_dict(decrypted_auth_config or {})
|
||||
set_kiro_request_context(build_kiro_request_context(request_body, cfg=cfg))
|
||||
wrapped = build_kiro_request_payload(
|
||||
request_body,
|
||||
model=model,
|
||||
cfg=cfg,
|
||||
)
|
||||
|
||||
return wrapped, url_model
|
||||
|
||||
def unwrap_response(self, data: Any) -> Any:
|
||||
return data
|
||||
|
||||
def postprocess_unwrapped_response(self, *, model: str, data: Any) -> None: # noqa: ARG002
|
||||
return
|
||||
|
||||
def capture_selected_base_url(self) -> str | None:
|
||||
return get_selected_base_url()
|
||||
|
||||
def on_http_status(self, *, base_url: str | None, status_code: int) -> None:
|
||||
from src.services.provider.adapters.kiro.context import update_kiro_http_status
|
||||
|
||||
category = classify_kiro_http_status(status_code)
|
||||
update_kiro_http_status(status_code=status_code, category=category)
|
||||
if status_code >= 400:
|
||||
logger.warning(
|
||||
"kiro upstream http status: status={}, category={}, base_url={}",
|
||||
status_code,
|
||||
category,
|
||||
base_url or "-",
|
||||
)
|
||||
|
||||
def on_connection_error(self, *, base_url: str | None, exc: Exception) -> None:
|
||||
from src.services.provider.adapters.kiro.context import update_kiro_connection_error
|
||||
|
||||
category = classify_kiro_connection_error(exc)
|
||||
summary = summarize_kiro_connection_error(exc)
|
||||
update_kiro_connection_error(category=category, summary=summary)
|
||||
logger.warning(
|
||||
"kiro upstream connection error: category={}, base_url={}, error={}",
|
||||
category,
|
||||
base_url or "-",
|
||||
summary,
|
||||
)
|
||||
|
||||
def force_stream_rewrite(self) -> bool:
|
||||
# Kiro streaming is binary AWS Event Stream and must be rewritten.
|
||||
return True
|
||||
|
||||
async def extract_error_text(
|
||||
self,
|
||||
source: Any,
|
||||
*,
|
||||
limit: int = 4000,
|
||||
) -> str:
|
||||
return await extract_kiro_http_error_text(source, limit=limit)
|
||||
|
||||
|
||||
kiro_envelope = KiroEnvelope()
|
||||
|
||||
|
||||
__all__ = ["KiroEnvelope", "kiro_envelope"]
|
||||
@@ -0,0 +1,175 @@
|
||||
"""Kiro HTTP/network error classification helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import httpx
|
||||
|
||||
_KNOWN_REASON_MESSAGES: dict[str, str] = {
|
||||
"CONTENT_LENGTH_EXCEEDS_THRESHOLD": "输入超过模型上下文限制",
|
||||
"MONTHLY_REQUEST_COUNT": "账户已达到月度请求配额",
|
||||
}
|
||||
|
||||
|
||||
def classify_kiro_http_status(status_code: int) -> str:
|
||||
"""Classify upstream HTTP status into stable buckets."""
|
||||
if 200 <= status_code < 300:
|
||||
return "ok"
|
||||
if status_code in {401, 403}:
|
||||
return "auth_error"
|
||||
if status_code == 429:
|
||||
return "rate_limited"
|
||||
if status_code in {408, 504}:
|
||||
return "timeout"
|
||||
if 500 <= status_code < 600:
|
||||
return "upstream_server_error"
|
||||
if 400 <= status_code < 500:
|
||||
return "upstream_client_error"
|
||||
return "unexpected_status"
|
||||
|
||||
|
||||
def classify_kiro_connection_error(exc: Exception) -> str:
|
||||
"""Classify transport exceptions raised by httpx."""
|
||||
if isinstance(exc, httpx.ConnectTimeout):
|
||||
return "connect_timeout"
|
||||
if isinstance(exc, httpx.ReadTimeout):
|
||||
return "read_timeout"
|
||||
if isinstance(exc, httpx.WriteTimeout):
|
||||
return "write_timeout"
|
||||
if isinstance(exc, httpx.PoolTimeout):
|
||||
return "pool_timeout"
|
||||
if isinstance(exc, httpx.TimeoutException):
|
||||
return "timeout"
|
||||
if isinstance(exc, httpx.ConnectError):
|
||||
return "connect_error"
|
||||
return "network_error"
|
||||
|
||||
|
||||
def summarize_kiro_connection_error(exc: Exception) -> str:
|
||||
"""Build a compact diagnostic string safe for logs/errors."""
|
||||
category = classify_kiro_connection_error(exc)
|
||||
detail = str(exc).strip()
|
||||
if len(detail) > 200:
|
||||
detail = detail[:200]
|
||||
if detail:
|
||||
return f"{category}: {type(exc).__name__}: {detail}"
|
||||
return f"{category}: {type(exc).__name__}"
|
||||
|
||||
|
||||
def build_kiro_network_diagnostic(
|
||||
*,
|
||||
http_status: int | None,
|
||||
http_category: str | None,
|
||||
connection_summary: str | None,
|
||||
) -> str | None:
|
||||
"""Build short supplemental diagnostic text for user-facing error paths."""
|
||||
if connection_summary:
|
||||
return f"network={connection_summary}"
|
||||
if http_status is None:
|
||||
return None
|
||||
category = str(http_category or "unknown").strip() or "unknown"
|
||||
return f"http_status={http_status} ({category})"
|
||||
|
||||
|
||||
def parse_kiro_error_text(raw_text: str | None) -> dict[str, str]:
|
||||
result = {
|
||||
"type": "",
|
||||
"reason": "",
|
||||
"message": "",
|
||||
"raw": str(raw_text or "").strip(),
|
||||
}
|
||||
if not result["raw"]:
|
||||
return result
|
||||
|
||||
try:
|
||||
data = json.loads(result["raw"])
|
||||
except Exception:
|
||||
result["message"] = result["raw"]
|
||||
return result
|
||||
|
||||
error_obj = data.get("error")
|
||||
if isinstance(error_obj, dict):
|
||||
result["type"] = str(error_obj.get("type") or error_obj.get("__type") or "").strip()
|
||||
result["reason"] = str(error_obj.get("reason") or error_obj.get("code") or "").strip()
|
||||
message = error_obj.get("message")
|
||||
if isinstance(message, str) and message.strip():
|
||||
result["message"] = message.strip()
|
||||
|
||||
if not result["message"]:
|
||||
message = data.get("message")
|
||||
if isinstance(message, str) and message.strip():
|
||||
result["message"] = message.strip()
|
||||
|
||||
if not result["reason"]:
|
||||
reason = data.get("reason") or data.get("code")
|
||||
if isinstance(reason, str) and reason.strip():
|
||||
result["reason"] = reason.strip()
|
||||
|
||||
if not result["message"]:
|
||||
result["message"] = result["raw"]
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def enhance_kiro_http_error_text(
|
||||
raw_text: str | None,
|
||||
*,
|
||||
status_code: int | None = None,
|
||||
) -> str:
|
||||
parsed = parse_kiro_error_text(raw_text)
|
||||
reason = parsed["reason"].upper()
|
||||
type_name = parsed["type"]
|
||||
message = parsed["message"]
|
||||
|
||||
friendly_message = _KNOWN_REASON_MESSAGES.get(reason)
|
||||
if friendly_message:
|
||||
message = friendly_message
|
||||
elif status_code == 403 and "access denied" in message.lower():
|
||||
message = "Kiro 账户权限被拒绝"
|
||||
elif status_code == 429 and not reason:
|
||||
message = "Kiro 请求过于频繁,请稍后重试"
|
||||
|
||||
parts: list[str] = []
|
||||
if type_name:
|
||||
parts.append(type_name)
|
||||
if reason:
|
||||
parts.append(f"[{reason}]")
|
||||
if message:
|
||||
parts.append(message)
|
||||
|
||||
return ": ".join(parts) if parts else parsed["raw"]
|
||||
|
||||
|
||||
async def extract_kiro_http_error_text(
|
||||
source: httpx.Response | httpx.HTTPStatusError,
|
||||
*,
|
||||
limit: int = 4000,
|
||||
) -> str:
|
||||
response = source.response if isinstance(source, httpx.HTTPStatusError) else source
|
||||
|
||||
raw_text = ""
|
||||
try:
|
||||
if hasattr(response, "is_stream_consumed") and not response.is_stream_consumed:
|
||||
error_bytes = await response.aread()
|
||||
raw_text = error_bytes.decode("utf-8", errors="replace")
|
||||
else:
|
||||
raw_text = response.text if hasattr(response, "_content") else ""
|
||||
except Exception as exc:
|
||||
return f"Unable to read Kiro error response: {exc}"
|
||||
|
||||
raw_text = (raw_text or "")[:limit]
|
||||
if not raw_text:
|
||||
return ""
|
||||
return enhance_kiro_http_error_text(raw_text, status_code=response.status_code)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"build_kiro_network_diagnostic",
|
||||
"classify_kiro_connection_error",
|
||||
"classify_kiro_http_status",
|
||||
"enhance_kiro_http_error_text",
|
||||
"extract_kiro_http_error_text",
|
||||
"parse_kiro_error_text",
|
||||
"summarize_kiro_connection_error",
|
||||
]
|
||||
@@ -0,0 +1,763 @@
|
||||
"""AWS Event Stream -> Claude SSE rewriter for Kiro.
|
||||
|
||||
Kiro streaming responses are returned as `application/vnd.amazon.eventstream`
|
||||
(binary framed). This module decodes frames and emits Claude-style streaming
|
||||
SSE events (as UTF-8 bytes).
|
||||
|
||||
The output format uses ``event: {type}\\ndata: {...}\\n\\n`` for typed events and
|
||||
plain ``data: {...}\\n\\n`` for untyped events, matching how Aether parses Claude
|
||||
streams.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import AsyncGenerator
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.services.provider.adapters.kiro.constants import CONTEXT_WINDOW_TOKENS
|
||||
from src.services.provider.adapters.kiro.error_enhancer import build_kiro_network_diagnostic
|
||||
from src.services.provider.adapters.kiro.parser.decoder import EventStreamDecoder
|
||||
|
||||
# Safety limit for thinking_buffer to prevent memory exhaustion from
|
||||
# pathological upstream responses that never close the thinking tag.
|
||||
_MAX_THINKING_BUFFER = 1024 * 1024 # 1 MiB
|
||||
|
||||
_QUOTE_CHARS: frozenset[str] = frozenset("`\"'\\#!@$%^&*()-_=+[]{};:<>,.?/")
|
||||
|
||||
|
||||
def _is_quote_char(buffer: str, pos: int) -> bool:
|
||||
if pos < 0 or pos >= len(buffer):
|
||||
return False
|
||||
return buffer[pos] in _QUOTE_CHARS
|
||||
|
||||
|
||||
def _find_real_thinking_start_tag(buffer: str) -> int | None:
|
||||
tag = "<thinking>"
|
||||
search = 0
|
||||
while True:
|
||||
pos = buffer.find(tag, search)
|
||||
if pos < 0:
|
||||
return None
|
||||
has_before = pos > 0 and _is_quote_char(buffer, pos - 1)
|
||||
after_pos = pos + len(tag)
|
||||
has_after = _is_quote_char(buffer, after_pos)
|
||||
if not has_before and not has_after:
|
||||
return pos
|
||||
search = pos + 1
|
||||
|
||||
|
||||
def _find_real_thinking_end_tag(buffer: str) -> int | None:
|
||||
tag = "</thinking>"
|
||||
search = 0
|
||||
while True:
|
||||
pos = buffer.find(tag, search)
|
||||
if pos < 0:
|
||||
return None
|
||||
|
||||
has_before = pos > 0 and _is_quote_char(buffer, pos - 1)
|
||||
after_pos = pos + len(tag)
|
||||
has_after = _is_quote_char(buffer, after_pos)
|
||||
if has_before or has_after:
|
||||
search = pos + 1
|
||||
continue
|
||||
|
||||
after = buffer[after_pos:]
|
||||
if len(after) < 2:
|
||||
return None
|
||||
if after.startswith("\n\n"):
|
||||
return pos
|
||||
|
||||
search = pos + 1
|
||||
|
||||
|
||||
def _find_real_thinking_end_tag_at_buffer_end(buffer: str) -> int | None:
|
||||
tag = "</thinking>"
|
||||
search = 0
|
||||
while True:
|
||||
pos = buffer.find(tag, search)
|
||||
if pos < 0:
|
||||
return None
|
||||
|
||||
has_before = pos > 0 and _is_quote_char(buffer, pos - 1)
|
||||
after_pos = pos + len(tag)
|
||||
has_after = _is_quote_char(buffer, after_pos)
|
||||
if has_before or has_after:
|
||||
search = pos + 1
|
||||
continue
|
||||
|
||||
if buffer[after_pos:].strip() == "":
|
||||
return pos
|
||||
|
||||
search = pos + 1
|
||||
|
||||
|
||||
def _estimate_tokens(text: str) -> int:
|
||||
if not text:
|
||||
return 0
|
||||
chinese = 0
|
||||
other = 0
|
||||
for c in text:
|
||||
if "\u4e00" <= c <= "\u9fff":
|
||||
chinese += 1
|
||||
else:
|
||||
other += 1
|
||||
chinese_tokens = (chinese * 2 + 2) // 3
|
||||
other_tokens = (other + 3) // 4
|
||||
return max(chinese_tokens + other_tokens, 1)
|
||||
|
||||
|
||||
def _sse_data_bytes(obj: dict[str, Any]) -> bytes:
|
||||
data = json.dumps(obj, ensure_ascii=False)
|
||||
event_type = obj.get("type", "")
|
||||
if event_type:
|
||||
return f"event: {event_type}\ndata: {data}\n\n".encode("utf-8")
|
||||
return f"data: {data}\n\n".encode("utf-8")
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _KiroStreamState:
|
||||
model: str
|
||||
thinking_enabled: bool
|
||||
estimated_input_tokens: int = 0
|
||||
|
||||
message_id: str = field(default_factory=lambda: f"msg_{uuid.uuid4().hex}")
|
||||
output_tokens: int = 0
|
||||
context_input_tokens: int | None = None
|
||||
|
||||
next_block_index: int = 0
|
||||
open_blocks: dict[int, str] = field(default_factory=dict)
|
||||
|
||||
text_block_index: int | None = None
|
||||
thinking_block_index: int | None = None
|
||||
tool_block_indices: dict[str, int] = field(default_factory=dict)
|
||||
|
||||
thinking_buffer: str = ""
|
||||
in_thinking_block: bool = False
|
||||
thinking_extracted: bool = False
|
||||
strip_thinking_leading_newline: bool = False
|
||||
|
||||
has_tool_use: bool = False
|
||||
stop_reason_override: str | None = None
|
||||
had_error: bool = False
|
||||
_last_content: str = ""
|
||||
|
||||
def generate_initial_events(self) -> list[dict[str, Any]]:
|
||||
events: list[dict[str, Any]] = []
|
||||
|
||||
# message_start
|
||||
events.append(
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": self.message_id,
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [],
|
||||
"model": self.model,
|
||||
"stop_reason": None,
|
||||
"stop_sequence": None,
|
||||
# Claude CLI clients expect usage to exist.
|
||||
"usage": {
|
||||
"input_tokens": int(self.estimated_input_tokens or 0),
|
||||
"output_tokens": 1,
|
||||
},
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
if not self.thinking_enabled:
|
||||
events.extend(self._ensure_text_block_open())
|
||||
|
||||
return events
|
||||
|
||||
def _ensure_text_block_open(self) -> list[dict[str, Any]]:
|
||||
if self.text_block_index is not None:
|
||||
if (
|
||||
self.text_block_index in self.open_blocks
|
||||
and self.open_blocks[self.text_block_index] == "text"
|
||||
):
|
||||
return []
|
||||
self.text_block_index = None
|
||||
|
||||
idx = self.next_block_index
|
||||
self.next_block_index += 1
|
||||
self.text_block_index = idx
|
||||
self.open_blocks[idx] = "text"
|
||||
|
||||
return [
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": idx,
|
||||
"content_block": {"type": "text", "text": ""},
|
||||
}
|
||||
]
|
||||
|
||||
def _close_block(self, idx: int) -> list[dict[str, Any]]:
|
||||
if idx not in self.open_blocks:
|
||||
return []
|
||||
self.open_blocks.pop(idx, None)
|
||||
return [{"type": "content_block_stop", "index": idx}]
|
||||
|
||||
def _ensure_thinking_block_open(self) -> list[dict[str, Any]]:
|
||||
if self.thinking_block_index is not None:
|
||||
if (
|
||||
self.thinking_block_index in self.open_blocks
|
||||
and self.open_blocks[self.thinking_block_index] == "thinking"
|
||||
):
|
||||
return []
|
||||
|
||||
idx = self.next_block_index
|
||||
self.next_block_index += 1
|
||||
self.thinking_block_index = idx
|
||||
self.open_blocks[idx] = "thinking"
|
||||
|
||||
return [
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": idx,
|
||||
"content_block": {"type": "thinking", "thinking": ""},
|
||||
}
|
||||
]
|
||||
|
||||
def _emit_text_delta(self, text: str) -> list[dict[str, Any]]:
|
||||
if not text:
|
||||
return []
|
||||
events: list[dict[str, Any]] = []
|
||||
events.extend(self._ensure_text_block_open())
|
||||
idx = int(self.text_block_index or 0)
|
||||
events.append(
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": idx,
|
||||
"delta": {"type": "text_delta", "text": text},
|
||||
}
|
||||
)
|
||||
return events
|
||||
|
||||
def _emit_thinking_delta(self, thinking: str) -> list[dict[str, Any]]:
|
||||
if not thinking:
|
||||
return []
|
||||
events: list[dict[str, Any]] = []
|
||||
events.extend(self._ensure_thinking_block_open())
|
||||
idx = int(self.thinking_block_index or 0)
|
||||
events.append(
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": idx,
|
||||
"delta": {"type": "thinking_delta", "thinking": thinking},
|
||||
}
|
||||
)
|
||||
return events
|
||||
|
||||
def _close_thinking_block(self) -> list[dict[str, Any]]:
|
||||
"""Send an empty thinking_delta sentinel and close the thinking block."""
|
||||
if self.thinking_block_index is None:
|
||||
return []
|
||||
idx = int(self.thinking_block_index)
|
||||
events: list[dict[str, Any]] = [
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": idx,
|
||||
"delta": {"type": "thinking_delta", "thinking": ""},
|
||||
}
|
||||
]
|
||||
events.extend(self._close_block(idx))
|
||||
return events
|
||||
|
||||
def process_context_usage(self, percentage: float) -> None:
|
||||
try:
|
||||
pct = float(percentage)
|
||||
except Exception:
|
||||
return
|
||||
# percentage * CONTEXT_WINDOW_TOKENS / 100
|
||||
self.context_input_tokens = int(pct * float(CONTEXT_WINDOW_TOKENS) / 100.0)
|
||||
|
||||
def process_exception(self, exception_type: str) -> None:
|
||||
if exception_type == "ContentLengthExceededException":
|
||||
# ContentLengthExceededException is a normal completion signal (output
|
||||
# exceeded size limit), not a fatal error. We record the stop_reason
|
||||
# but do NOT set had_error so that finalize() still emits message_delta
|
||||
# with stop_reason="max_tokens" and message_stop.
|
||||
self.stop_reason_override = "max_tokens"
|
||||
return
|
||||
|
||||
def process_assistant_response(self, content: str) -> list[dict[str, Any]]:
|
||||
if not content:
|
||||
return []
|
||||
|
||||
# Kiro may send duplicate content events; skip exact repeats.
|
||||
if content == self._last_content:
|
||||
return []
|
||||
self._last_content = content
|
||||
|
||||
self.output_tokens += _estimate_tokens(content)
|
||||
|
||||
if not self.thinking_enabled:
|
||||
return self._emit_text_delta(content)
|
||||
|
||||
self.thinking_buffer += content
|
||||
|
||||
# Safety: flush as text if thinking_buffer grows too large without closing tag
|
||||
if len(self.thinking_buffer) > _MAX_THINKING_BUFFER:
|
||||
logger.warning(
|
||||
"kiro thinking_buffer exceeded {} bytes, force-flushing as text",
|
||||
_MAX_THINKING_BUFFER,
|
||||
)
|
||||
overflow = self.thinking_buffer
|
||||
self.thinking_buffer = ""
|
||||
if self.in_thinking_block:
|
||||
result = self._emit_thinking_delta(overflow)
|
||||
result.extend(self._close_thinking_block())
|
||||
self.in_thinking_block = False
|
||||
self.thinking_extracted = True
|
||||
return result
|
||||
return self._emit_text_delta(overflow)
|
||||
|
||||
events: list[dict[str, Any]] = []
|
||||
|
||||
while True:
|
||||
if not self.in_thinking_block and not self.thinking_extracted:
|
||||
start_pos = _find_real_thinking_start_tag(self.thinking_buffer)
|
||||
if start_pos is not None:
|
||||
before = self.thinking_buffer[:start_pos]
|
||||
if before and before.strip():
|
||||
events.extend(self._emit_text_delta(before))
|
||||
|
||||
self.in_thinking_block = True
|
||||
self.strip_thinking_leading_newline = True
|
||||
self.thinking_buffer = self.thinking_buffer[start_pos + len("<thinking>") :]
|
||||
events.extend(self._ensure_thinking_block_open())
|
||||
continue
|
||||
|
||||
# Keep a short suffix in buffer for partial tag detection.
|
||||
keep = len("<thinking>")
|
||||
if len(self.thinking_buffer) > keep:
|
||||
safe = self.thinking_buffer[:-keep]
|
||||
if safe and safe.strip():
|
||||
events.extend(self._emit_text_delta(safe))
|
||||
self.thinking_buffer = self.thinking_buffer[-keep:]
|
||||
break
|
||||
|
||||
if self.in_thinking_block:
|
||||
# Strip a single leading \n after <thinking> tag.
|
||||
# The model outputs `<thinking>\n` and the \n may arrive in the
|
||||
# same chunk or the next one; we drop it for cleaner output.
|
||||
if self.strip_thinking_leading_newline:
|
||||
if self.thinking_buffer.startswith("\n"):
|
||||
self.thinking_buffer = self.thinking_buffer[1:]
|
||||
self.strip_thinking_leading_newline = False
|
||||
elif self.thinking_buffer:
|
||||
# Buffer is non-empty but doesn't start with \n; stop waiting.
|
||||
self.strip_thinking_leading_newline = False
|
||||
# else: buffer is empty, keep the flag for the next chunk.
|
||||
|
||||
end_pos = _find_real_thinking_end_tag(self.thinking_buffer)
|
||||
if end_pos is not None:
|
||||
thinking_text = self.thinking_buffer[:end_pos]
|
||||
if thinking_text:
|
||||
events.extend(self._emit_thinking_delta(thinking_text))
|
||||
|
||||
events.extend(self._close_thinking_block())
|
||||
|
||||
self.in_thinking_block = False
|
||||
self.thinking_extracted = True
|
||||
self.thinking_buffer = self.thinking_buffer[end_pos + len("</thinking>") :]
|
||||
continue
|
||||
|
||||
keep = len("</thinking>")
|
||||
if len(self.thinking_buffer) > keep:
|
||||
safe = self.thinking_buffer[:-keep]
|
||||
if safe:
|
||||
events.extend(self._emit_thinking_delta(safe))
|
||||
self.thinking_buffer = self.thinking_buffer[-keep:]
|
||||
break
|
||||
|
||||
# thinking extracted: remaining buffer is text
|
||||
if self.thinking_buffer:
|
||||
remaining = self.thinking_buffer
|
||||
self.thinking_buffer = ""
|
||||
events.extend(self._emit_text_delta(remaining))
|
||||
break
|
||||
|
||||
return events
|
||||
|
||||
def process_tool_use(
|
||||
self,
|
||||
*,
|
||||
name: str,
|
||||
tool_use_id: str,
|
||||
input_json: str,
|
||||
stop: bool,
|
||||
) -> list[dict[str, Any]]:
|
||||
if not tool_use_id:
|
||||
return []
|
||||
|
||||
self.has_tool_use = True
|
||||
|
||||
events: list[dict[str, Any]] = []
|
||||
|
||||
# Boundary: close thinking block if needed, filtering a dangling </thinking>.
|
||||
if self.thinking_enabled and self.in_thinking_block and self.thinking_buffer:
|
||||
end_pos = _find_real_thinking_end_tag_at_buffer_end(self.thinking_buffer)
|
||||
if end_pos is not None:
|
||||
thinking_text = self.thinking_buffer[:end_pos]
|
||||
if thinking_text:
|
||||
events.extend(self._emit_thinking_delta(thinking_text))
|
||||
|
||||
events.extend(self._close_thinking_block())
|
||||
|
||||
after_pos = end_pos + len("</thinking>")
|
||||
remaining = self.thinking_buffer[after_pos:]
|
||||
self.thinking_buffer = ""
|
||||
self.in_thinking_block = False
|
||||
self.thinking_extracted = True
|
||||
if remaining:
|
||||
events.extend(self._emit_text_delta(remaining))
|
||||
else:
|
||||
# Best-effort flush all as thinking
|
||||
events.extend(self._emit_thinking_delta(self.thinking_buffer))
|
||||
events.extend(self._close_thinking_block())
|
||||
self.thinking_buffer = ""
|
||||
self.in_thinking_block = False
|
||||
self.thinking_extracted = True
|
||||
|
||||
# Flush any buffered pre-thinking tail so tool_use doesn't swallow it.
|
||||
if (
|
||||
self.thinking_enabled
|
||||
and not self.in_thinking_block
|
||||
and not self.thinking_extracted
|
||||
and self.thinking_buffer
|
||||
):
|
||||
buffered = self.thinking_buffer
|
||||
self.thinking_buffer = ""
|
||||
events.extend(self._emit_text_delta(buffered))
|
||||
|
||||
# Close current text block before tool_use.
|
||||
if self.text_block_index is not None:
|
||||
idx = int(self.text_block_index)
|
||||
events.extend(self._close_block(idx))
|
||||
|
||||
block_index = self.tool_block_indices.get(tool_use_id)
|
||||
if block_index is None:
|
||||
block_index = self.next_block_index
|
||||
self.next_block_index += 1
|
||||
self.tool_block_indices[tool_use_id] = block_index
|
||||
|
||||
# Start tool block if not open.
|
||||
if block_index not in self.open_blocks:
|
||||
self.open_blocks[block_index] = "tool_use"
|
||||
events.append(
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": block_index,
|
||||
"content_block": {
|
||||
"type": "tool_use",
|
||||
"id": tool_use_id,
|
||||
"name": name,
|
||||
"input": {},
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
if input_json:
|
||||
self.output_tokens += _estimate_tokens(input_json)
|
||||
events.append(
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": block_index,
|
||||
"delta": {"type": "input_json_delta", "partial_json": input_json},
|
||||
}
|
||||
)
|
||||
|
||||
if stop:
|
||||
events.extend(self._close_block(block_index))
|
||||
|
||||
return events
|
||||
|
||||
def finalize(self) -> list[dict[str, Any]]:
|
||||
events: list[dict[str, Any]] = []
|
||||
|
||||
# Flush remaining thinking/text buffer.
|
||||
if self.thinking_enabled and self.thinking_buffer:
|
||||
if self.in_thinking_block:
|
||||
end_pos = _find_real_thinking_end_tag_at_buffer_end(self.thinking_buffer)
|
||||
if end_pos is not None:
|
||||
thinking_text = self.thinking_buffer[:end_pos]
|
||||
if thinking_text:
|
||||
events.extend(self._emit_thinking_delta(thinking_text))
|
||||
|
||||
events.extend(self._close_thinking_block())
|
||||
|
||||
after_pos = end_pos + len("</thinking>")
|
||||
remaining = self.thinking_buffer[after_pos:]
|
||||
if remaining:
|
||||
events.extend(self._emit_text_delta(remaining))
|
||||
else:
|
||||
events.extend(self._emit_thinking_delta(self.thinking_buffer))
|
||||
events.extend(self._close_thinking_block())
|
||||
|
||||
else:
|
||||
events.extend(self._emit_text_delta(self.thinking_buffer))
|
||||
|
||||
self.thinking_buffer = ""
|
||||
self.in_thinking_block = False
|
||||
self.thinking_extracted = True
|
||||
|
||||
# Close any open blocks (best-effort).
|
||||
for idx in sorted(list(self.open_blocks.keys()), reverse=True):
|
||||
events.extend(self._close_block(idx))
|
||||
|
||||
stop_reason = self.stop_reason_override
|
||||
if not stop_reason:
|
||||
stop_reason = "tool_use" if self.has_tool_use else "end_turn"
|
||||
|
||||
input_tokens = (
|
||||
int(self.context_input_tokens)
|
||||
if self.context_input_tokens is not None
|
||||
else int(self.estimated_input_tokens or 0)
|
||||
)
|
||||
|
||||
events.append(
|
||||
{
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": stop_reason, "stop_sequence": None},
|
||||
"usage": {"input_tokens": input_tokens, "output_tokens": int(self.output_tokens)},
|
||||
}
|
||||
)
|
||||
events.append({"type": "message_stop"})
|
||||
|
||||
return events
|
||||
|
||||
|
||||
async def rewrite_eventstream_to_sse(
|
||||
byte_iterator: Any,
|
||||
*,
|
||||
model: str,
|
||||
thinking_enabled: bool,
|
||||
estimated_input_tokens: int = 0,
|
||||
) -> AsyncGenerator[bytes]:
|
||||
"""Rewrite Kiro AWS Event Stream bytes to Claude SSE bytes."""
|
||||
decoder = EventStreamDecoder()
|
||||
state = _KiroStreamState(
|
||||
model=str(model or ""),
|
||||
thinking_enabled=bool(thinking_enabled),
|
||||
estimated_input_tokens=int(estimated_input_tokens or 0),
|
||||
)
|
||||
|
||||
# 收集原始字节用于错误诊断
|
||||
raw_bytes_buffer = b""
|
||||
|
||||
# Initial events
|
||||
for evt in state.generate_initial_events():
|
||||
yield _sse_data_bytes(evt)
|
||||
|
||||
async for chunk in byte_iterator:
|
||||
if not chunk:
|
||||
continue
|
||||
|
||||
# 保留原始字节用于错误诊断(限制大小)
|
||||
if len(raw_bytes_buffer) < 4096:
|
||||
raw_bytes_buffer += chunk
|
||||
|
||||
try:
|
||||
decoder.feed(chunk)
|
||||
frames = decoder.decode_available()
|
||||
except Exception as e:
|
||||
logger.warning("kiro eventstream decode error: {}", e)
|
||||
# 尝试解析原始响应为 JSON 错误
|
||||
error_message = f"kiro eventstream decode failed: {type(e).__name__}"
|
||||
try:
|
||||
raw_text = raw_bytes_buffer.decode("utf-8", errors="replace")
|
||||
# 尝试解析为 JSON
|
||||
error_json = json.loads(raw_text)
|
||||
if isinstance(error_json, dict):
|
||||
# 提取上游错误信息
|
||||
upstream_msg = error_json.get("message") or error_json.get("error", {}).get(
|
||||
"message"
|
||||
)
|
||||
if upstream_msg:
|
||||
error_message = f"Kiro API error: {upstream_msg}"
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
from src.services.provider.adapters.kiro.context import get_kiro_request_context
|
||||
|
||||
kiro_ctx = get_kiro_request_context()
|
||||
diag = build_kiro_network_diagnostic(
|
||||
http_status=kiro_ctx.last_http_status if kiro_ctx else None,
|
||||
http_category=kiro_ctx.last_http_error_category if kiro_ctx else None,
|
||||
connection_summary=(
|
||||
kiro_ctx.last_connection_error_summary if kiro_ctx else None
|
||||
),
|
||||
)
|
||||
if diag:
|
||||
error_message = f"{error_message} | {diag}"
|
||||
except Exception:
|
||||
pass
|
||||
yield _sse_data_bytes(
|
||||
{
|
||||
"type": "error",
|
||||
"error": {
|
||||
"type": "upstream_stream_error",
|
||||
"message": error_message,
|
||||
},
|
||||
}
|
||||
)
|
||||
break
|
||||
|
||||
for frame in frames:
|
||||
mtype = (frame.message_type() or "event").strip().lower()
|
||||
etype = (frame.event_type() or "").strip()
|
||||
payload_text = frame.payload_as_text()
|
||||
|
||||
if mtype == "event":
|
||||
try:
|
||||
payload = json.loads(payload_text) if payload_text else {}
|
||||
except Exception:
|
||||
payload = {}
|
||||
|
||||
if etype == "assistantResponseEvent":
|
||||
content = payload.get("content") if isinstance(payload, dict) else None
|
||||
if isinstance(content, str) and content:
|
||||
for evt in state.process_assistant_response(content):
|
||||
yield _sse_data_bytes(evt)
|
||||
continue
|
||||
|
||||
if etype == "toolUseEvent":
|
||||
if isinstance(payload, dict):
|
||||
name = str(payload.get("name") or "")
|
||||
tool_use_id = payload.get("toolUseId") or payload.get("tool_use_id")
|
||||
tool_use_id = str(tool_use_id or "")
|
||||
raw_input = payload.get("input")
|
||||
if raw_input is None:
|
||||
input_json = ""
|
||||
elif isinstance(raw_input, str):
|
||||
input_json = raw_input
|
||||
else:
|
||||
try:
|
||||
input_json = json.dumps(raw_input, ensure_ascii=False)
|
||||
except Exception:
|
||||
input_json = str(raw_input)
|
||||
stop = bool(payload.get("stop", False))
|
||||
for evt in state.process_tool_use(
|
||||
name=name,
|
||||
tool_use_id=tool_use_id,
|
||||
input_json=input_json,
|
||||
stop=stop,
|
||||
):
|
||||
yield _sse_data_bytes(evt)
|
||||
continue
|
||||
|
||||
if etype == "contextUsageEvent":
|
||||
if isinstance(payload, dict):
|
||||
pct = payload.get("contextUsagePercentage")
|
||||
if pct is not None:
|
||||
try:
|
||||
state.process_context_usage(float(pct))
|
||||
except (ValueError, TypeError):
|
||||
logger.debug(
|
||||
"kiro: failed to parse contextUsagePercentage: {!r}", pct
|
||||
)
|
||||
continue
|
||||
|
||||
# meteringEvent / unknown: ignore
|
||||
continue
|
||||
|
||||
if mtype == "exception":
|
||||
ex_type = frame.headers.exception_type() or "UnknownException"
|
||||
state.process_exception(ex_type)
|
||||
# ContentLengthExceededException is handled by process_exception
|
||||
# (sets stop_reason_override) and should NOT prevent finalize().
|
||||
if not state.stop_reason_override:
|
||||
state.had_error = True
|
||||
logger.debug("kiro upstream exception: {} | {}", ex_type, payload_text[:200])
|
||||
if state.had_error:
|
||||
yield _sse_data_bytes(
|
||||
{
|
||||
"type": "error",
|
||||
"error": {
|
||||
"type": "upstream_exception",
|
||||
"message": ex_type,
|
||||
},
|
||||
}
|
||||
)
|
||||
continue
|
||||
|
||||
if mtype == "error":
|
||||
err_code = frame.headers.error_code() or "UnknownError"
|
||||
state.had_error = True
|
||||
logger.debug("kiro upstream error: {} | {}", err_code, payload_text[:200])
|
||||
yield _sse_data_bytes(
|
||||
{
|
||||
"type": "error",
|
||||
"error": {
|
||||
"type": "upstream_error",
|
||||
"message": err_code,
|
||||
},
|
||||
}
|
||||
)
|
||||
continue
|
||||
|
||||
if not state.had_error:
|
||||
for evt in state.finalize():
|
||||
yield _sse_data_bytes(evt)
|
||||
|
||||
|
||||
def apply_kiro_stream_rewrite(
|
||||
byte_iter: Any,
|
||||
*,
|
||||
model: str = "",
|
||||
input_tokens: int = 0,
|
||||
prefetched_chunks: list[bytes] | None = None,
|
||||
) -> AsyncGenerator[bytes]:
|
||||
"""Apply Kiro EventStream->SSE rewrite if context is available.
|
||||
|
||||
Consolidates the repeated import-context-rewrite pattern used across
|
||||
``chat_handler_base``, ``cli_handler_base``, and ``stream_processor``.
|
||||
|
||||
Args:
|
||||
byte_iter: Upstream byte iterator (raw AWS Event Stream).
|
||||
model: Model name for SSE events.
|
||||
input_tokens: Estimated input token count.
|
||||
prefetched_chunks: Optional pre-fetched bytes to prepend.
|
||||
|
||||
Returns:
|
||||
An async generator of Claude-compatible SSE bytes.
|
||||
"""
|
||||
from src.services.provider.adapters.kiro.context import get_kiro_request_context
|
||||
|
||||
kiro_ctx = get_kiro_request_context()
|
||||
thinking_enabled = bool(getattr(kiro_ctx, "thinking_enabled", False)) if kiro_ctx else False
|
||||
|
||||
if prefetched_chunks:
|
||||
upstream = byte_iter
|
||||
prefix = list(prefetched_chunks)
|
||||
|
||||
async def _combined() -> AsyncGenerator[bytes, None]:
|
||||
for c in prefix:
|
||||
if c:
|
||||
yield c
|
||||
async for c in upstream:
|
||||
if c:
|
||||
yield c
|
||||
|
||||
source: Any = _combined()
|
||||
else:
|
||||
source = byte_iter
|
||||
|
||||
return rewrite_eventstream_to_sse(
|
||||
source,
|
||||
model=str(model or ""),
|
||||
thinking_enabled=thinking_enabled,
|
||||
estimated_input_tokens=int(input_tokens or 0),
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"apply_kiro_stream_rewrite",
|
||||
"rewrite_eventstream_to_sse",
|
||||
]
|
||||
102
_deprecated_py_src/services/provider/adapters/kiro/headers.py
Normal file
102
_deprecated_py_src/services/provider/adapters/kiro/headers.py
Normal file
@@ -0,0 +1,102 @@
|
||||
"""Kiro header builders."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
|
||||
from src.services.provider.adapters.kiro.constants import (
|
||||
AWS_EVENTSTREAM_CONTENT_TYPE,
|
||||
AWS_SDK_JS_MAIN_VERSION,
|
||||
AWS_SDK_JS_USAGE_VERSION,
|
||||
CODEWHISPERER_OPTOUT,
|
||||
DEFAULT_KIRO_VERSION,
|
||||
DEFAULT_NODE_VERSION,
|
||||
DEFAULT_SYSTEM_VERSION,
|
||||
KIRO_AGENT_MODE,
|
||||
)
|
||||
|
||||
|
||||
def build_kiro_ide_tag(*, kiro_version: str, machine_id: str) -> str:
|
||||
version = (kiro_version or DEFAULT_KIRO_VERSION).strip() or DEFAULT_KIRO_VERSION
|
||||
mid = (machine_id or "").strip()
|
||||
return f"KiroIDE-{version}-{mid}" if mid else f"KiroIDE-{version}"
|
||||
|
||||
|
||||
def build_x_amz_user_agent_main(*, kiro_version: str, machine_id: str) -> str:
|
||||
return f"aws-sdk-js/{AWS_SDK_JS_MAIN_VERSION} {build_kiro_ide_tag(kiro_version=kiro_version, machine_id=machine_id)}"
|
||||
|
||||
|
||||
def build_user_agent_main(
|
||||
*, system_version: str, node_version: str, kiro_version: str, machine_id: str
|
||||
) -> str:
|
||||
os_tag = (system_version or DEFAULT_SYSTEM_VERSION).strip() or DEFAULT_SYSTEM_VERSION
|
||||
node_tag = (node_version or DEFAULT_NODE_VERSION).strip() or DEFAULT_NODE_VERSION
|
||||
ide = build_kiro_ide_tag(kiro_version=kiro_version, machine_id=machine_id)
|
||||
return (
|
||||
f"aws-sdk-js/{AWS_SDK_JS_MAIN_VERSION} ua/2.1 os/{os_tag} lang/js "
|
||||
f"md/nodejs#{node_tag} api/codewhispererstreaming#{AWS_SDK_JS_MAIN_VERSION} m/E {ide}"
|
||||
)
|
||||
|
||||
|
||||
def build_x_amz_user_agent_usage(*, kiro_version: str, machine_id: str) -> str:
|
||||
ide = build_kiro_ide_tag(kiro_version=kiro_version, machine_id=machine_id)
|
||||
return f"aws-sdk-js/{AWS_SDK_JS_USAGE_VERSION} {ide}"
|
||||
|
||||
|
||||
def build_user_agent_usage(*, kiro_version: str, machine_id: str) -> str:
|
||||
ide = build_kiro_ide_tag(kiro_version=kiro_version, machine_id=machine_id)
|
||||
os_tag = DEFAULT_SYSTEM_VERSION
|
||||
node_tag = DEFAULT_NODE_VERSION
|
||||
return (
|
||||
f"aws-sdk-js/{AWS_SDK_JS_USAGE_VERSION} ua/2.1 os/{os_tag} lang/js "
|
||||
f"md/nodejs#{node_tag} api/codewhispererruntime#1.0.0 m/N,E {ide}"
|
||||
)
|
||||
|
||||
|
||||
def build_generate_assistant_headers(
|
||||
*,
|
||||
host: str,
|
||||
access_token: str | None = None,
|
||||
machine_id: str,
|
||||
kiro_version: str | None = None,
|
||||
system_version: str | None = None,
|
||||
node_version: str | None = None,
|
||||
) -> dict[str, str]:
|
||||
version = (kiro_version or DEFAULT_KIRO_VERSION).strip() or DEFAULT_KIRO_VERSION
|
||||
sys_ver = (system_version or DEFAULT_SYSTEM_VERSION).strip() or DEFAULT_SYSTEM_VERSION
|
||||
node_ver = (node_version or DEFAULT_NODE_VERSION).strip() or DEFAULT_NODE_VERSION
|
||||
|
||||
headers: dict[str, str] = {
|
||||
"Content-Type": "application/json",
|
||||
"Accept": AWS_EVENTSTREAM_CONTENT_TYPE,
|
||||
"host": host,
|
||||
"Connection": "close",
|
||||
"x-amzn-codewhisperer-optout": CODEWHISPERER_OPTOUT,
|
||||
"x-amzn-kiro-agent-mode": KIRO_AGENT_MODE,
|
||||
"x-amz-user-agent": build_x_amz_user_agent_main(
|
||||
kiro_version=version, machine_id=machine_id
|
||||
),
|
||||
"User-Agent": build_user_agent_main(
|
||||
system_version=sys_ver,
|
||||
node_version=node_ver,
|
||||
kiro_version=version,
|
||||
machine_id=machine_id,
|
||||
),
|
||||
"amz-sdk-invocation-id": str(uuid.uuid4()),
|
||||
"amz-sdk-request": "attempt=1; max=3",
|
||||
}
|
||||
|
||||
if access_token:
|
||||
headers["Authorization"] = f"Bearer {access_token}"
|
||||
|
||||
return headers
|
||||
|
||||
|
||||
__all__ = [
|
||||
"build_generate_assistant_headers",
|
||||
"build_kiro_ide_tag",
|
||||
"build_user_agent_main",
|
||||
"build_user_agent_usage",
|
||||
"build_x_amz_user_agent_main",
|
||||
"build_x_amz_user_agent_usage",
|
||||
]
|
||||
@@ -0,0 +1,21 @@
|
||||
from .credentials import KiroAuthConfig
|
||||
from .usage_limits import (
|
||||
Bonus,
|
||||
FreeTrialInfo,
|
||||
SubscriptionInfo,
|
||||
UsageBreakdown,
|
||||
UsageLimitsResponse,
|
||||
calculate_current_usage,
|
||||
calculate_total_usage_limit,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"Bonus",
|
||||
"FreeTrialInfo",
|
||||
"KiroAuthConfig",
|
||||
"SubscriptionInfo",
|
||||
"UsageBreakdown",
|
||||
"UsageLimitsResponse",
|
||||
"calculate_current_usage",
|
||||
"calculate_total_usage_limit",
|
||||
]
|
||||
@@ -0,0 +1,247 @@
|
||||
"""Internal Kiro credential schema (stored in ProviderAPIKey.auth_config)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
|
||||
def _parse_epoch_seconds(value: object) -> int | None:
|
||||
if value is None:
|
||||
return None
|
||||
try:
|
||||
if isinstance(value, (int, float)):
|
||||
return int(value)
|
||||
if isinstance(value, str) and value.strip().isdigit():
|
||||
return int(value.strip())
|
||||
except Exception:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _get_str(raw: dict[str, Any], *keys: str) -> str | None:
|
||||
"""Return the first non-empty stripped string for *keys*, or ``None``."""
|
||||
for k in keys:
|
||||
v = raw.get(k)
|
||||
if isinstance(v, str) and v.strip():
|
||||
return v.strip()
|
||||
return None
|
||||
|
||||
|
||||
def _nonempty(s: str | None) -> str | None:
|
||||
"""Return *s* if it's a non-empty stripped string, else ``None``."""
|
||||
if isinstance(s, str) and s.strip():
|
||||
return s.strip()
|
||||
return None
|
||||
|
||||
|
||||
def _normalize_auth_method(value: str | None) -> str:
|
||||
method = (value or "").strip().lower()
|
||||
if not method:
|
||||
return "social"
|
||||
# 历史/别名兼容:统一映射到 idc
|
||||
if method in {
|
||||
"idc",
|
||||
"builder-id",
|
||||
"builder_id",
|
||||
"builderid",
|
||||
"identity-center",
|
||||
"identity_center",
|
||||
"identitycenter",
|
||||
"iam",
|
||||
"device",
|
||||
"device_authorization",
|
||||
"device-auth",
|
||||
}:
|
||||
return "idc"
|
||||
return method
|
||||
|
||||
|
||||
def _parse_iso_to_epoch_seconds(value: object) -> int | None:
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
return None
|
||||
text = value.strip()
|
||||
# Support RFC3339 with Z suffix.
|
||||
if text.endswith("Z"):
|
||||
text = text[:-1] + "+00:00"
|
||||
try:
|
||||
dt = datetime.fromisoformat(text)
|
||||
except Exception:
|
||||
return None
|
||||
if dt.tzinfo is None:
|
||||
dt = dt.replace(tzinfo=timezone.utc)
|
||||
return int(dt.timestamp())
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class KiroAuthConfig:
|
||||
provider_type: str = "kiro"
|
||||
|
||||
auth_method: str = "social" # social | idc
|
||||
refresh_token: str = ""
|
||||
expires_at: int = 0
|
||||
|
||||
profile_arn: str | None = None
|
||||
region: str | None = None # OIDC region(IdC token 刷新用)
|
||||
|
||||
# 独立的 auth / api region(与 kiro.rs 对齐)
|
||||
# auth_region: token 刷新端点,未设置时回退到 region
|
||||
# api_region: q.{region} 服务端点,未设置时回退到 DEFAULT_REGION
|
||||
auth_region: str | None = None
|
||||
api_region: str | None = None
|
||||
|
||||
client_id: str | None = None
|
||||
client_secret: str | None = None
|
||||
|
||||
machine_id: str | None = None
|
||||
kiro_version: str | None = None
|
||||
system_version: str | None = None
|
||||
node_version: str | None = None
|
||||
|
||||
email: str | None = None # 账号邮箱
|
||||
|
||||
# 缓存的 access_token(可选,用于避免频繁刷新)
|
||||
access_token: str | None = None
|
||||
|
||||
def effective_auth_region(self) -> str:
|
||||
"""Token 刷新用的 region。
|
||||
|
||||
优先级: auth_region > region > DEFAULT_REGION
|
||||
"""
|
||||
from src.services.provider.adapters.kiro.constants import DEFAULT_REGION
|
||||
|
||||
return _nonempty(self.auth_region) or _nonempty(self.region) or DEFAULT_REGION
|
||||
|
||||
def effective_api_region(self) -> str:
|
||||
"""API 服务端点(q.{region})用的 region。
|
||||
|
||||
优先级: api_region > DEFAULT_REGION
|
||||
注意: 不从 region 继承,因为 region 通常是 OIDC region(如 eu-north-1),
|
||||
而 q.{region} 端点目前仅 us-east-1 可用。
|
||||
"""
|
||||
from src.services.provider.adapters.kiro.constants import DEFAULT_REGION
|
||||
|
||||
return _nonempty(self.api_region) or DEFAULT_REGION
|
||||
|
||||
@staticmethod
|
||||
def infer_auth_method(raw: dict[str, Any]) -> str:
|
||||
"""
|
||||
根据凭据字段自动推断认证类型。
|
||||
|
||||
规则:
|
||||
- 包含 clientId + clientSecret -> IdC
|
||||
- 仅含 refreshToken -> Social
|
||||
"""
|
||||
explicit_method = _get_str(raw, "auth_method", "authMethod", "auth_type", "authType")
|
||||
normalized_explicit = _normalize_auth_method(explicit_method)
|
||||
if normalized_explicit != "social":
|
||||
return normalized_explicit
|
||||
|
||||
client_id = raw.get("client_id") or raw.get("clientId")
|
||||
client_secret = raw.get("client_secret") or raw.get("clientSecret")
|
||||
|
||||
if client_id and client_secret:
|
||||
return "idc"
|
||||
return "social"
|
||||
|
||||
@staticmethod
|
||||
def validate_required_fields(raw: dict[str, Any]) -> tuple[bool, str]:
|
||||
"""
|
||||
验证凭据是否包含必需字段。
|
||||
|
||||
返回: (is_valid, error_message)
|
||||
"""
|
||||
refresh_token = raw.get("refresh_token") or raw.get("refreshToken") or ""
|
||||
refresh_token = str(refresh_token).strip()
|
||||
|
||||
if not refresh_token:
|
||||
return False, "refreshToken 为必填字段"
|
||||
|
||||
# refreshToken 不能含有 ...(表示被截断)
|
||||
if "..." in refresh_token:
|
||||
return False, "refreshToken 不完整(含有 ...),请导出完整的 Token"
|
||||
|
||||
# IdC 类型需要 clientId 和 clientSecret
|
||||
explicit_method = _get_str(raw, "auth_method", "authMethod", "auth_type", "authType")
|
||||
auth_method = (
|
||||
_normalize_auth_method(explicit_method)
|
||||
if explicit_method
|
||||
else KiroAuthConfig.infer_auth_method(raw)
|
||||
)
|
||||
if auth_method == "idc":
|
||||
client_id = raw.get("client_id") or raw.get("clientId")
|
||||
client_secret = raw.get("client_secret") or raw.get("clientSecret")
|
||||
if not client_id:
|
||||
return False, "IdC 类型需要 clientId"
|
||||
if not client_secret:
|
||||
return False, "IdC 类型需要 clientSecret"
|
||||
|
||||
return True, ""
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, raw: dict[str, Any]) -> "KiroAuthConfig":
|
||||
if not isinstance(raw, dict):
|
||||
raw = {}
|
||||
|
||||
provider_type = _get_str(raw, "provider_type", "providerType") or "kiro"
|
||||
|
||||
# 自动推断 auth_method(如果未显式指定)
|
||||
explicit_method = _get_str(raw, "auth_method", "authMethod", "auth_type", "authType")
|
||||
auth_method = (
|
||||
_normalize_auth_method(explicit_method)
|
||||
if explicit_method
|
||||
else cls.infer_auth_method(raw)
|
||||
)
|
||||
|
||||
refresh_token = (_get_str(raw, "refresh_token", "refreshToken") or "").strip()
|
||||
|
||||
expires_at = _parse_epoch_seconds(raw.get("expires_at"))
|
||||
if expires_at is None:
|
||||
expires_at = _parse_iso_to_epoch_seconds(raw.get("expiresAt"))
|
||||
if expires_at is None:
|
||||
expires_at = 0
|
||||
|
||||
cfg = cls(
|
||||
provider_type=provider_type,
|
||||
auth_method=_normalize_auth_method(auth_method),
|
||||
refresh_token=refresh_token,
|
||||
expires_at=int(expires_at),
|
||||
profile_arn=_get_str(raw, "profile_arn", "profileArn"),
|
||||
region=_get_str(raw, "region"),
|
||||
auth_region=_get_str(raw, "auth_region", "authRegion"),
|
||||
api_region=_get_str(raw, "api_region", "apiRegion"),
|
||||
client_id=_get_str(raw, "client_id", "clientId"),
|
||||
client_secret=_get_str(raw, "client_secret", "clientSecret"),
|
||||
machine_id=_get_str(raw, "machine_id", "machineId"),
|
||||
kiro_version=_get_str(raw, "kiro_version", "kiroVersion"),
|
||||
system_version=_get_str(raw, "system_version", "systemVersion"),
|
||||
node_version=_get_str(raw, "node_version", "nodeVersion"),
|
||||
email=_get_str(raw, "email"),
|
||||
access_token=_get_str(raw, "access_token", "accessToken"),
|
||||
)
|
||||
|
||||
return cfg
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"provider_type": self.provider_type,
|
||||
"auth_method": self.auth_method,
|
||||
"refresh_token": self.refresh_token,
|
||||
"expires_at": self.expires_at,
|
||||
"profile_arn": self.profile_arn,
|
||||
"region": self.region,
|
||||
"auth_region": self.auth_region,
|
||||
"api_region": self.api_region,
|
||||
"client_id": self.client_id,
|
||||
"client_secret": self.client_secret,
|
||||
"machine_id": self.machine_id,
|
||||
"kiro_version": self.kiro_version,
|
||||
"system_version": self.system_version,
|
||||
"node_version": self.node_version,
|
||||
"email": self.email,
|
||||
"access_token": self.access_token,
|
||||
}
|
||||
|
||||
|
||||
__all__ = ["KiroAuthConfig"]
|
||||
@@ -0,0 +1,222 @@
|
||||
"""Kiro getUsageLimits response models (best-effort).
|
||||
|
||||
The AWS API uses camelCase fields; we parse defensively.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class SubscriptionInfo:
|
||||
subscription_title: str | None = None
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, raw: Any) -> "SubscriptionInfo | None":
|
||||
if not isinstance(raw, dict):
|
||||
return None
|
||||
title = raw.get("subscriptionTitle")
|
||||
if isinstance(title, str) and title.strip():
|
||||
return cls(subscription_title=title.strip())
|
||||
return cls(subscription_title=None)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class Bonus:
|
||||
current_usage: float = 0.0
|
||||
usage_limit: float = 0.0
|
||||
status: str | None = None
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, raw: Any) -> "Bonus | None":
|
||||
if not isinstance(raw, dict):
|
||||
return None
|
||||
status = raw.get("status")
|
||||
status_val = status.strip() if isinstance(status, str) and status.strip() else None
|
||||
cu = raw.get("currentUsage")
|
||||
ul = raw.get("usageLimit")
|
||||
try:
|
||||
cu_f = float(cu) if cu is not None else 0.0
|
||||
except Exception:
|
||||
cu_f = 0.0
|
||||
try:
|
||||
ul_f = float(ul) if ul is not None else 0.0
|
||||
except Exception:
|
||||
ul_f = 0.0
|
||||
return cls(current_usage=cu_f, usage_limit=ul_f, status=status_val)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class FreeTrialInfo:
|
||||
current_usage: int = 0
|
||||
current_usage_with_precision: float = 0.0
|
||||
usage_limit: int = 0
|
||||
usage_limit_with_precision: float = 0.0
|
||||
free_trial_expiry: float | None = None
|
||||
free_trial_status: str | None = None
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, raw: Any) -> "FreeTrialInfo | None":
|
||||
if not isinstance(raw, dict):
|
||||
return None
|
||||
|
||||
def _int(v: Any) -> int:
|
||||
try:
|
||||
return int(v)
|
||||
except Exception:
|
||||
return 0
|
||||
|
||||
def _float(v: Any) -> float:
|
||||
try:
|
||||
return float(v)
|
||||
except Exception:
|
||||
return 0.0
|
||||
|
||||
expiry = raw.get("freeTrialExpiry")
|
||||
try:
|
||||
expiry_f = float(expiry) if expiry is not None else None
|
||||
except Exception:
|
||||
expiry_f = None
|
||||
|
||||
status = raw.get("freeTrialStatus")
|
||||
status_val = status.strip() if isinstance(status, str) and status.strip() else None
|
||||
|
||||
return cls(
|
||||
current_usage=_int(raw.get("currentUsage")),
|
||||
current_usage_with_precision=_float(raw.get("currentUsageWithPrecision")),
|
||||
usage_limit=_int(raw.get("usageLimit")),
|
||||
usage_limit_with_precision=_float(raw.get("usageLimitWithPrecision")),
|
||||
free_trial_expiry=expiry_f,
|
||||
free_trial_status=status_val,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class UsageBreakdown:
|
||||
current_usage: int = 0
|
||||
current_usage_with_precision: float = 0.0
|
||||
usage_limit: int = 0
|
||||
usage_limit_with_precision: float = 0.0
|
||||
next_date_reset: float | None = None
|
||||
bonuses: list[Bonus] = field(default_factory=list)
|
||||
free_trial_info: FreeTrialInfo | None = None
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, raw: Any) -> "UsageBreakdown | None":
|
||||
if not isinstance(raw, dict):
|
||||
return None
|
||||
|
||||
def _int(v: Any) -> int:
|
||||
try:
|
||||
return int(v)
|
||||
except Exception:
|
||||
return 0
|
||||
|
||||
def _float(v: Any) -> float:
|
||||
try:
|
||||
return float(v)
|
||||
except Exception:
|
||||
return 0.0
|
||||
|
||||
next_reset = raw.get("nextDateReset")
|
||||
try:
|
||||
next_reset_f = float(next_reset) if next_reset is not None else None
|
||||
except Exception:
|
||||
next_reset_f = None
|
||||
|
||||
bonuses_raw = raw.get("bonuses")
|
||||
bonuses: list[Bonus] = []
|
||||
if isinstance(bonuses_raw, list):
|
||||
for b in bonuses_raw:
|
||||
parsed = Bonus.from_dict(b)
|
||||
if parsed is not None:
|
||||
bonuses.append(parsed)
|
||||
|
||||
return cls(
|
||||
current_usage=_int(raw.get("currentUsage")),
|
||||
current_usage_with_precision=_float(raw.get("currentUsageWithPrecision")),
|
||||
usage_limit=_int(raw.get("usageLimit")),
|
||||
usage_limit_with_precision=_float(raw.get("usageLimitWithPrecision")),
|
||||
next_date_reset=next_reset_f,
|
||||
bonuses=bonuses,
|
||||
free_trial_info=FreeTrialInfo.from_dict(raw.get("freeTrialInfo")),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class UsageLimitsResponse:
|
||||
next_date_reset: float | None = None
|
||||
subscription_info: SubscriptionInfo | None = None
|
||||
usage_breakdown_list: list[UsageBreakdown] = field(default_factory=list)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, raw: Any) -> "UsageLimitsResponse":
|
||||
if not isinstance(raw, dict):
|
||||
raw = {}
|
||||
|
||||
next_reset = raw.get("nextDateReset")
|
||||
try:
|
||||
next_reset_f = float(next_reset) if next_reset is not None else None
|
||||
except Exception:
|
||||
next_reset_f = None
|
||||
|
||||
breakdown_raw = raw.get("usageBreakdownList")
|
||||
breakdowns: list[UsageBreakdown] = []
|
||||
if isinstance(breakdown_raw, list):
|
||||
for b in breakdown_raw:
|
||||
parsed = UsageBreakdown.from_dict(b)
|
||||
if parsed is not None:
|
||||
breakdowns.append(parsed)
|
||||
|
||||
return cls(
|
||||
next_date_reset=next_reset_f,
|
||||
subscription_info=SubscriptionInfo.from_dict(raw.get("subscriptionInfo")),
|
||||
usage_breakdown_list=breakdowns,
|
||||
)
|
||||
|
||||
|
||||
def calculate_total_usage_limit(response: UsageLimitsResponse) -> float:
|
||||
if not response.usage_breakdown_list:
|
||||
return 0.0
|
||||
|
||||
breakdown = response.usage_breakdown_list[0]
|
||||
total = breakdown.usage_limit_with_precision
|
||||
|
||||
if breakdown.free_trial_info and breakdown.free_trial_info.free_trial_status == "ACTIVE":
|
||||
total += breakdown.free_trial_info.usage_limit_with_precision
|
||||
|
||||
for bonus in breakdown.bonuses:
|
||||
if bonus.status == "ACTIVE":
|
||||
total += bonus.usage_limit
|
||||
|
||||
return total
|
||||
|
||||
|
||||
def calculate_current_usage(response: UsageLimitsResponse) -> float:
|
||||
if not response.usage_breakdown_list:
|
||||
return 0.0
|
||||
|
||||
breakdown = response.usage_breakdown_list[0]
|
||||
total = breakdown.current_usage_with_precision
|
||||
|
||||
if breakdown.free_trial_info and breakdown.free_trial_info.free_trial_status == "ACTIVE":
|
||||
total += breakdown.free_trial_info.current_usage_with_precision
|
||||
|
||||
for bonus in breakdown.bonuses:
|
||||
if bonus.status == "ACTIVE":
|
||||
total += bonus.current_usage
|
||||
|
||||
return total
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Bonus",
|
||||
"FreeTrialInfo",
|
||||
"SubscriptionInfo",
|
||||
"UsageBreakdown",
|
||||
"UsageLimitsResponse",
|
||||
"calculate_current_usage",
|
||||
"calculate_total_usage_limit",
|
||||
]
|
||||
@@ -0,0 +1,6 @@
|
||||
"""AWS Event Stream parser for Kiro."""
|
||||
|
||||
from .decoder import EventStreamDecoder
|
||||
from .frame import Frame
|
||||
|
||||
__all__ = ["EventStreamDecoder", "Frame"]
|
||||
@@ -0,0 +1,13 @@
|
||||
"""CRC helpers for AWS Event Stream frames."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import binascii
|
||||
|
||||
|
||||
def crc32(data: bytes) -> int:
|
||||
"""Compute unsigned CRC32 (IEEE)."""
|
||||
return binascii.crc32(data) & 0xFFFFFFFF
|
||||
|
||||
|
||||
__all__ = ["crc32"]
|
||||
@@ -0,0 +1,91 @@
|
||||
"""Incremental AWS Event Stream decoder."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from .error import BufferOverflowError, EventStreamParseError
|
||||
from .frame import MAX_MESSAGE_SIZE, Frame, parse_frame
|
||||
|
||||
DEFAULT_MAX_BUFFER_SIZE = MAX_MESSAGE_SIZE
|
||||
DEFAULT_MAX_ERRORS = 5
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class DecoderStats:
|
||||
frames_decoded: int = 0
|
||||
bytes_skipped: int = 0
|
||||
error_count: int = 0
|
||||
|
||||
|
||||
class EventStreamDecoder:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
max_buffer_size: int = DEFAULT_MAX_BUFFER_SIZE,
|
||||
max_errors: int = DEFAULT_MAX_ERRORS,
|
||||
) -> None:
|
||||
self._buffer = bytearray()
|
||||
self._max_buffer_size = int(max_buffer_size)
|
||||
self._max_errors = int(max_errors)
|
||||
self._stopped = False
|
||||
self.stats = DecoderStats()
|
||||
|
||||
@property
|
||||
def stopped(self) -> bool:
|
||||
return self._stopped
|
||||
|
||||
def feed(self, data: bytes) -> None:
|
||||
if self._stopped:
|
||||
return
|
||||
if not data:
|
||||
return
|
||||
new_size = len(self._buffer) + len(data)
|
||||
if new_size > self._max_buffer_size:
|
||||
self._stopped = True
|
||||
raise BufferOverflowError(size=new_size, max_size=self._max_buffer_size)
|
||||
self._buffer.extend(data)
|
||||
|
||||
def decode_available(self) -> list[Frame]:
|
||||
"""Decode all complete frames currently in buffer."""
|
||||
out: list[Frame] = []
|
||||
if self._stopped:
|
||||
return out
|
||||
|
||||
while True:
|
||||
try:
|
||||
# Use memoryview to avoid full buffer copy on each iteration
|
||||
parsed = parse_frame(memoryview(self._buffer))
|
||||
except EventStreamParseError:
|
||||
self.stats.error_count += 1
|
||||
if self.stats.error_count >= self._max_errors:
|
||||
self._stopped = True
|
||||
raise
|
||||
|
||||
# Recovery: skip a byte and keep scanning.
|
||||
if self._buffer:
|
||||
del self._buffer[0]
|
||||
self.stats.bytes_skipped += 1
|
||||
else:
|
||||
break
|
||||
continue
|
||||
|
||||
if parsed is None:
|
||||
break
|
||||
|
||||
frame, consumed = parsed
|
||||
if consumed <= 0:
|
||||
break
|
||||
|
||||
out.append(frame)
|
||||
del self._buffer[:consumed]
|
||||
self.stats.frames_decoded += 1
|
||||
self.stats.error_count = 0
|
||||
|
||||
return out
|
||||
|
||||
|
||||
__all__ = [
|
||||
"DecoderStats",
|
||||
"EventStreamDecoder",
|
||||
]
|
||||
@@ -0,0 +1,72 @@
|
||||
"""AWS Event Stream parsing errors."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
class EventStreamParseError(Exception):
|
||||
"""Base error for AWS Event Stream decoding."""
|
||||
|
||||
|
||||
class IncompleteFrameError(EventStreamParseError):
|
||||
def __init__(self, *, needed: int, available: int) -> None:
|
||||
super().__init__(f"incomplete frame: needed={needed} available={available}")
|
||||
self.needed = needed
|
||||
self.available = available
|
||||
|
||||
|
||||
class MessageTooSmallError(EventStreamParseError):
|
||||
def __init__(self, *, length: int, min_length: int) -> None:
|
||||
super().__init__(f"message too small: length={length} min={min_length}")
|
||||
self.length = length
|
||||
self.min_length = min_length
|
||||
|
||||
|
||||
class MessageTooLargeError(EventStreamParseError):
|
||||
def __init__(self, *, length: int, max_length: int) -> None:
|
||||
super().__init__(f"message too large: length={length} max={max_length}")
|
||||
self.length = length
|
||||
self.max_length = max_length
|
||||
|
||||
|
||||
class PreludeCrcMismatchError(EventStreamParseError):
|
||||
def __init__(self, *, expected: int, actual: int) -> None:
|
||||
super().__init__(f"prelude crc mismatch: expected={expected} actual={actual}")
|
||||
self.expected = expected
|
||||
self.actual = actual
|
||||
|
||||
|
||||
class MessageCrcMismatchError(EventStreamParseError):
|
||||
def __init__(self, *, expected: int, actual: int) -> None:
|
||||
super().__init__(f"message crc mismatch: expected={expected} actual={actual}")
|
||||
self.expected = expected
|
||||
self.actual = actual
|
||||
|
||||
|
||||
class InvalidHeaderTypeError(EventStreamParseError):
|
||||
def __init__(self, type_id: int) -> None:
|
||||
super().__init__(f"invalid header type: {type_id}")
|
||||
self.type_id = type_id
|
||||
|
||||
|
||||
class HeaderParseError(EventStreamParseError):
|
||||
pass
|
||||
|
||||
|
||||
class BufferOverflowError(EventStreamParseError):
|
||||
def __init__(self, *, size: int, max_size: int) -> None:
|
||||
super().__init__(f"buffer overflow: size={size} max={max_size}")
|
||||
self.size = size
|
||||
self.max_size = max_size
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BufferOverflowError",
|
||||
"EventStreamParseError",
|
||||
"HeaderParseError",
|
||||
"IncompleteFrameError",
|
||||
"InvalidHeaderTypeError",
|
||||
"MessageCrcMismatchError",
|
||||
"MessageTooLargeError",
|
||||
"MessageTooSmallError",
|
||||
"PreludeCrcMismatchError",
|
||||
]
|
||||
@@ -0,0 +1,95 @@
|
||||
"""AWS Event Stream message frame parsing."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from .crc import crc32
|
||||
from .error import (
|
||||
HeaderParseError,
|
||||
IncompleteFrameError,
|
||||
MessageCrcMismatchError,
|
||||
MessageTooLargeError,
|
||||
MessageTooSmallError,
|
||||
PreludeCrcMismatchError,
|
||||
)
|
||||
from .header import Headers, parse_headers
|
||||
|
||||
PRELUDE_SIZE = 12
|
||||
MIN_MESSAGE_SIZE = PRELUDE_SIZE + 4
|
||||
MAX_MESSAGE_SIZE = 16 * 1024 * 1024
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class Frame:
|
||||
headers: Headers
|
||||
payload: bytes
|
||||
|
||||
def message_type(self) -> str | None:
|
||||
return self.headers.message_type()
|
||||
|
||||
def event_type(self) -> str | None:
|
||||
return self.headers.event_type()
|
||||
|
||||
def payload_as_text(self) -> str:
|
||||
return self.payload.decode("utf-8", errors="replace")
|
||||
|
||||
|
||||
def parse_frame(buffer: bytes | memoryview) -> tuple[Frame, int] | None:
|
||||
"""Parse a single frame from the front of buffer.
|
||||
|
||||
Returns:
|
||||
(frame, consumed_bytes) if a full frame is available, otherwise None.
|
||||
|
||||
Raises:
|
||||
EventStreamParseError subclasses on validation errors.
|
||||
"""
|
||||
if len(buffer) < PRELUDE_SIZE:
|
||||
return None
|
||||
|
||||
total_length = int.from_bytes(buffer[0:4], "big", signed=False)
|
||||
header_length = int.from_bytes(buffer[4:8], "big", signed=False)
|
||||
prelude_crc = int.from_bytes(buffer[8:12], "big", signed=False)
|
||||
|
||||
if total_length < MIN_MESSAGE_SIZE:
|
||||
raise MessageTooSmallError(length=total_length, min_length=MIN_MESSAGE_SIZE)
|
||||
if total_length > MAX_MESSAGE_SIZE:
|
||||
raise MessageTooLargeError(length=total_length, max_length=MAX_MESSAGE_SIZE)
|
||||
|
||||
if len(buffer) < total_length:
|
||||
return None
|
||||
|
||||
actual_prelude_crc = crc32(buffer[0:8])
|
||||
if actual_prelude_crc != prelude_crc:
|
||||
raise PreludeCrcMismatchError(expected=prelude_crc, actual=actual_prelude_crc)
|
||||
|
||||
message_crc = int.from_bytes(buffer[total_length - 4 : total_length], "big", signed=False)
|
||||
actual_message_crc = crc32(buffer[0 : total_length - 4])
|
||||
if actual_message_crc != message_crc:
|
||||
raise MessageCrcMismatchError(expected=message_crc, actual=actual_message_crc)
|
||||
|
||||
headers_start = PRELUDE_SIZE
|
||||
headers_end = headers_start + header_length
|
||||
|
||||
if headers_end > total_length - 4:
|
||||
raise HeaderParseError("header length exceeds frame boundary")
|
||||
|
||||
headers = parse_headers(bytes(buffer[headers_start:headers_end]), header_length)
|
||||
|
||||
payload_start = headers_end
|
||||
payload_end = total_length - 4
|
||||
if payload_end < payload_start:
|
||||
raise IncompleteFrameError(needed=payload_start, available=payload_end)
|
||||
|
||||
payload = bytes(buffer[payload_start:payload_end])
|
||||
|
||||
return Frame(headers=headers, payload=payload), total_length
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Frame",
|
||||
"MAX_MESSAGE_SIZE",
|
||||
"MIN_MESSAGE_SIZE",
|
||||
"PRELUDE_SIZE",
|
||||
"parse_frame",
|
||||
]
|
||||
@@ -0,0 +1,144 @@
|
||||
"""AWS Event Stream header parsing."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from enum import IntEnum
|
||||
|
||||
from .error import HeaderParseError, IncompleteFrameError, InvalidHeaderTypeError
|
||||
|
||||
|
||||
class HeaderValueType(IntEnum):
|
||||
BOOL_TRUE = 0
|
||||
BOOL_FALSE = 1
|
||||
BYTE = 2
|
||||
SHORT = 3
|
||||
INTEGER = 4
|
||||
LONG = 5
|
||||
BYTE_ARRAY = 6
|
||||
STRING = 7
|
||||
TIMESTAMP = 8
|
||||
UUID = 9
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class Headers:
|
||||
values: dict[str, object]
|
||||
|
||||
def get(self, name: str) -> object | None:
|
||||
return self.values.get(name)
|
||||
|
||||
def get_string(self, name: str) -> str | None:
|
||||
v = self.values.get(name)
|
||||
return v if isinstance(v, str) else None
|
||||
|
||||
def message_type(self) -> str | None:
|
||||
return self.get_string(":message-type")
|
||||
|
||||
def event_type(self) -> str | None:
|
||||
return self.get_string(":event-type")
|
||||
|
||||
def exception_type(self) -> str | None:
|
||||
return self.get_string(":exception-type")
|
||||
|
||||
def error_code(self) -> str | None:
|
||||
return self.get_string(":error-code")
|
||||
|
||||
|
||||
def _ensure_bytes(data: bytes, offset: int, needed: int) -> None:
|
||||
available = len(data) - offset
|
||||
if available < needed:
|
||||
raise IncompleteFrameError(needed=needed, available=available)
|
||||
|
||||
|
||||
def parse_headers(data: bytes, header_length: int) -> Headers:
|
||||
if len(data) < header_length:
|
||||
raise IncompleteFrameError(needed=header_length, available=len(data))
|
||||
|
||||
values: dict[str, object] = {}
|
||||
offset = 0
|
||||
|
||||
while offset < header_length:
|
||||
_ensure_bytes(data, offset, 1)
|
||||
name_len = data[offset]
|
||||
offset += 1
|
||||
if name_len == 0:
|
||||
raise HeaderParseError("header name length cannot be 0")
|
||||
|
||||
_ensure_bytes(data, offset, name_len)
|
||||
name = data[offset : offset + name_len].decode("utf-8", errors="replace")
|
||||
offset += name_len
|
||||
|
||||
_ensure_bytes(data, offset, 1)
|
||||
type_id = data[offset]
|
||||
offset += 1
|
||||
try:
|
||||
value_type = HeaderValueType(type_id)
|
||||
except ValueError as e:
|
||||
raise InvalidHeaderTypeError(type_id) from e
|
||||
|
||||
if value_type == HeaderValueType.BOOL_TRUE:
|
||||
values[name] = True
|
||||
continue
|
||||
if value_type == HeaderValueType.BOOL_FALSE:
|
||||
values[name] = False
|
||||
continue
|
||||
|
||||
if value_type == HeaderValueType.BYTE:
|
||||
_ensure_bytes(data, offset, 1)
|
||||
values[name] = int.from_bytes(data[offset : offset + 1], "big", signed=True)
|
||||
offset += 1
|
||||
continue
|
||||
|
||||
if value_type == HeaderValueType.SHORT:
|
||||
_ensure_bytes(data, offset, 2)
|
||||
values[name] = int.from_bytes(data[offset : offset + 2], "big", signed=True)
|
||||
offset += 2
|
||||
continue
|
||||
|
||||
if value_type == HeaderValueType.INTEGER:
|
||||
_ensure_bytes(data, offset, 4)
|
||||
values[name] = int.from_bytes(data[offset : offset + 4], "big", signed=True)
|
||||
offset += 4
|
||||
continue
|
||||
|
||||
if value_type in (HeaderValueType.LONG, HeaderValueType.TIMESTAMP):
|
||||
_ensure_bytes(data, offset, 8)
|
||||
values[name] = int.from_bytes(data[offset : offset + 8], "big", signed=True)
|
||||
offset += 8
|
||||
continue
|
||||
|
||||
if value_type == HeaderValueType.BYTE_ARRAY:
|
||||
_ensure_bytes(data, offset, 2)
|
||||
length = int.from_bytes(data[offset : offset + 2], "big", signed=False)
|
||||
offset += 2
|
||||
_ensure_bytes(data, offset, length)
|
||||
values[name] = data[offset : offset + length]
|
||||
offset += length
|
||||
continue
|
||||
|
||||
if value_type == HeaderValueType.STRING:
|
||||
_ensure_bytes(data, offset, 2)
|
||||
length = int.from_bytes(data[offset : offset + 2], "big", signed=False)
|
||||
offset += 2
|
||||
_ensure_bytes(data, offset, length)
|
||||
values[name] = data[offset : offset + length].decode("utf-8", errors="replace")
|
||||
offset += length
|
||||
continue
|
||||
|
||||
if value_type == HeaderValueType.UUID:
|
||||
_ensure_bytes(data, offset, 16)
|
||||
values[name] = bytes(data[offset : offset + 16])
|
||||
offset += 16
|
||||
continue
|
||||
|
||||
raise HeaderParseError(f"unhandled header type: {value_type}")
|
||||
|
||||
return Headers(values=values)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"HeaderValueType",
|
||||
"Headers",
|
||||
"parse_headers",
|
||||
]
|
||||
133
_deprecated_py_src/services/provider/adapters/kiro/plugin.py
Normal file
133
_deprecated_py_src/services/provider/adapters/kiro/plugin.py
Normal file
@@ -0,0 +1,133 @@
|
||||
"""Kiro provider plugin — unified registration entry.
|
||||
|
||||
Kiro upstream looks like Claude CLI (Bearer token) from the outside, but uses a
|
||||
custom wire protocol:
|
||||
- Request: Claude Messages API -> Kiro generateAssistantResponse envelope
|
||||
- Response (stream): AWS Event Stream (binary) -> Claude SSE events
|
||||
|
||||
This plugin registers:
|
||||
- Envelope
|
||||
- Transport hook (dynamic region base_url)
|
||||
- Model fetcher (fixed model catalog — Kiro has no /v1/models endpoint)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
from urllib.parse import urlencode
|
||||
|
||||
from src.services.provider.adapters.kiro.constants import DEFAULT_REGION
|
||||
from src.services.provider.adapters.kiro.context import get_kiro_request_context
|
||||
from src.services.provider.adapters.kiro.models.credentials import KiroAuthConfig
|
||||
from src.services.provider.adapters.kiro.request import (
|
||||
build_kiro_generate_assistant_url,
|
||||
resolve_kiro_base_url,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Preset model catalog
|
||||
# ---------------------------------------------------------------------------
|
||||
# Kiro upstream has no /v1/models endpoint. We use the unified preset models
|
||||
# registry from preset_models.py.
|
||||
from src.services.provider.preset_models import create_preset_models_fetcher
|
||||
from src.services.provider.request_context import set_selected_base_url
|
||||
|
||||
fetch_models_kiro = create_preset_models_fetcher("kiro")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Transport hook
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def build_kiro_url(
|
||||
endpoint: Any,
|
||||
*,
|
||||
is_stream: bool,
|
||||
effective_query_params: dict[str, Any],
|
||||
**_kwargs: Any,
|
||||
) -> str:
|
||||
"""Build Kiro generateAssistantResponse URL.
|
||||
|
||||
Endpoint base_url may contain a `{region}` placeholder. The actual region is
|
||||
resolved from per-request context (set by the envelope).
|
||||
"""
|
||||
_ = is_stream
|
||||
|
||||
ctx = get_kiro_request_context()
|
||||
region = (ctx.region if ctx else "") or DEFAULT_REGION
|
||||
raw_base = str(getattr(endpoint, "base_url", "") or "").rstrip("/")
|
||||
cfg = KiroAuthConfig(api_region=region)
|
||||
base = resolve_kiro_base_url(raw_base, cfg=cfg)
|
||||
set_selected_base_url(base)
|
||||
|
||||
url = build_kiro_generate_assistant_url(raw_base, cfg=cfg)
|
||||
|
||||
if effective_query_params:
|
||||
query_string = urlencode(effective_query_params, doseq=True)
|
||||
if query_string:
|
||||
url = f"{url}?{query_string}"
|
||||
|
||||
return url
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Export builder
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_KIRO_SKIP_KEYS = frozenset(
|
||||
{
|
||||
"access_token",
|
||||
"expires_at",
|
||||
"updated_at",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def kiro_export_builder(
|
||||
auth_config: dict[str, Any],
|
||||
upstream_metadata: dict[str, Any] | None,
|
||||
) -> dict[str, Any]:
|
||||
"""Kiro 导出:保留 auth_method / refresh_token / machine_id / profile_arn 等,
|
||||
IdC 模式额外保留 client_id / client_secret / region。"""
|
||||
data = {
|
||||
k: v
|
||||
for k, v in auth_config.items()
|
||||
if k not in _KIRO_SKIP_KEYS and v is not None and v != ""
|
||||
}
|
||||
# email 可能仅在 upstream_metadata.kiro 中
|
||||
if not data.get("email"):
|
||||
kiro_meta = (upstream_metadata or {}).get("kiro") or {}
|
||||
if kiro_meta.get("email"):
|
||||
data["email"] = kiro_meta["email"]
|
||||
return data
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Registration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def register_all() -> None:
|
||||
"""Register all Kiro hooks into shared registries."""
|
||||
|
||||
from src.services.model.upstream_fetcher import UpstreamModelsFetcherRegistry
|
||||
from src.services.provider.adapters.kiro.envelope import kiro_envelope
|
||||
from src.services.provider.envelope import register_envelope
|
||||
from src.services.provider.export import register_export_builder
|
||||
from src.services.provider.transport import register_transport_hook
|
||||
|
||||
register_envelope("kiro", "claude:cli", kiro_envelope)
|
||||
register_envelope("kiro", "", kiro_envelope)
|
||||
|
||||
register_transport_hook("kiro", "claude:cli", build_kiro_url)
|
||||
|
||||
register_export_builder("kiro", kiro_export_builder)
|
||||
|
||||
UpstreamModelsFetcherRegistry.register(
|
||||
provider_types=["kiro"],
|
||||
fetcher=fetch_models_kiro,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["build_kiro_url", "fetch_models_kiro", "kiro_export_builder", "register_all"]
|
||||
148
_deprecated_py_src/services/provider/adapters/kiro/request.py
Normal file
148
_deprecated_py_src/services/provider/adapters/kiro/request.py
Normal file
@@ -0,0 +1,148 @@
|
||||
"""Helpers for building Kiro generateAssistantResponse requests."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from src.services.provider.adapters.kiro.constants import KIRO_GENERATE_ASSISTANT_PATH
|
||||
from src.services.provider.adapters.kiro.context import KiroRequestContext
|
||||
from src.services.provider.adapters.kiro.converter import (
|
||||
convert_claude_messages_to_conversation_state,
|
||||
)
|
||||
from src.services.provider.adapters.kiro.headers import build_generate_assistant_headers
|
||||
from src.services.provider.adapters.kiro.models.credentials import KiroAuthConfig
|
||||
from src.services.provider.adapters.kiro.token_manager import generate_machine_id
|
||||
|
||||
|
||||
def is_kiro_thinking_enabled(request_body: dict[str, Any]) -> bool:
|
||||
thinking = request_body.get("thinking")
|
||||
if not isinstance(thinking, dict):
|
||||
return False
|
||||
ttype = str(thinking.get("type") or "").strip().lower()
|
||||
return ttype in {"enabled", "adaptive"}
|
||||
|
||||
|
||||
def build_kiro_request_context(
|
||||
request_body: dict[str, Any],
|
||||
*,
|
||||
cfg: KiroAuthConfig,
|
||||
) -> KiroRequestContext:
|
||||
return KiroRequestContext(
|
||||
region=cfg.effective_api_region(),
|
||||
machine_id=generate_machine_id(cfg),
|
||||
kiro_version=cfg.kiro_version,
|
||||
system_version=cfg.system_version,
|
||||
node_version=cfg.node_version,
|
||||
thinking_enabled=is_kiro_thinking_enabled(request_body),
|
||||
)
|
||||
|
||||
|
||||
def build_kiro_request_headers(
|
||||
cfg: KiroAuthConfig,
|
||||
*,
|
||||
access_token: str | None = None,
|
||||
) -> dict[str, str]:
|
||||
region = cfg.effective_api_region()
|
||||
host = f"q.{region}.amazonaws.com"
|
||||
return build_generate_assistant_headers(
|
||||
host=host,
|
||||
access_token=access_token,
|
||||
machine_id=generate_machine_id(cfg),
|
||||
kiro_version=cfg.kiro_version,
|
||||
system_version=cfg.system_version,
|
||||
node_version=cfg.node_version,
|
||||
)
|
||||
|
||||
|
||||
def resolve_kiro_base_url(base_url: str, *, cfg: KiroAuthConfig) -> str:
|
||||
resolved = str(base_url or "").rstrip("/")
|
||||
region = cfg.effective_api_region()
|
||||
if "{region}" in resolved:
|
||||
resolved = resolved.replace("{region}", region)
|
||||
return resolved
|
||||
|
||||
|
||||
def build_kiro_generate_assistant_url(base_url: str, *, cfg: KiroAuthConfig) -> str:
|
||||
resolved = resolve_kiro_base_url(base_url, cfg=cfg)
|
||||
if resolved.endswith(KIRO_GENERATE_ASSISTANT_PATH):
|
||||
return resolved
|
||||
return f"{resolved}{KIRO_GENERATE_ASSISTANT_PATH}"
|
||||
|
||||
|
||||
def build_kiro_inference_config(request_body: dict[str, Any]) -> dict[str, Any] | None:
|
||||
inference_config: dict[str, Any] = {}
|
||||
|
||||
max_tokens = request_body.get("max_tokens")
|
||||
try:
|
||||
max_tokens_i = int(max_tokens) if max_tokens is not None else 0
|
||||
except Exception:
|
||||
max_tokens_i = 0
|
||||
if max_tokens_i > 0:
|
||||
inference_config["maxTokens"] = max_tokens_i
|
||||
|
||||
temperature = request_body.get("temperature")
|
||||
try:
|
||||
temperature_f = float(temperature) if temperature is not None else None
|
||||
except Exception:
|
||||
temperature_f = None
|
||||
if temperature_f is not None and temperature_f >= 0:
|
||||
inference_config["temperature"] = temperature_f
|
||||
|
||||
top_p = request_body.get("top_p")
|
||||
try:
|
||||
top_p_f = float(top_p) if top_p is not None else None
|
||||
except Exception:
|
||||
top_p_f = None
|
||||
if top_p_f is not None and top_p_f > 0:
|
||||
inference_config["topP"] = top_p_f
|
||||
|
||||
return inference_config or None
|
||||
|
||||
|
||||
def get_profile_arn_for_payload(cfg: KiroAuthConfig) -> str | None:
|
||||
profile_arn = str(cfg.profile_arn or "").strip()
|
||||
if not profile_arn:
|
||||
return None
|
||||
|
||||
from src.services.provider.adapters.kiro.models.credentials import _normalize_auth_method
|
||||
|
||||
if _normalize_auth_method(cfg.auth_method) == "idc":
|
||||
return None
|
||||
|
||||
return profile_arn
|
||||
|
||||
|
||||
def build_kiro_request_payload(
|
||||
request_body: dict[str, Any],
|
||||
*,
|
||||
model: str,
|
||||
cfg: KiroAuthConfig,
|
||||
) -> dict[str, Any]:
|
||||
payload: dict[str, Any] = {
|
||||
"conversationState": convert_claude_messages_to_conversation_state(
|
||||
request_body,
|
||||
model=model,
|
||||
)
|
||||
}
|
||||
|
||||
inference_config = build_kiro_inference_config(request_body)
|
||||
if inference_config:
|
||||
payload["inferenceConfig"] = inference_config
|
||||
|
||||
profile_arn = get_profile_arn_for_payload(cfg)
|
||||
if profile_arn:
|
||||
payload["profileArn"] = profile_arn
|
||||
|
||||
return payload
|
||||
|
||||
|
||||
__all__ = [
|
||||
"build_kiro_generate_assistant_url",
|
||||
"build_kiro_inference_config",
|
||||
"build_kiro_request_context",
|
||||
"build_kiro_request_headers",
|
||||
"build_kiro_request_payload",
|
||||
"get_profile_arn_for_payload",
|
||||
"is_kiro_thinking_enabled",
|
||||
"resolve_kiro_base_url",
|
||||
]
|
||||
108
_deprecated_py_src/services/provider/adapters/kiro/rust_http.py
Normal file
108
_deprecated_py_src/services/provider/adapters/kiro/rust_http.py
Normal file
@@ -0,0 +1,108 @@
|
||||
"""Shared Rust executor HTTP helper for Kiro provider side calls."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from src.config.settings import config
|
||||
from src.core.logger import logger
|
||||
|
||||
|
||||
async def execute_kiro_rust_http_request(
|
||||
*,
|
||||
method: str,
|
||||
url: str,
|
||||
headers: dict[str, str],
|
||||
body: Any,
|
||||
proxy_config: dict[str, Any] | None,
|
||||
request_id: str,
|
||||
provider_api_format: str,
|
||||
content_type: str | None = None,
|
||||
timeout_seconds: float = 30.0,
|
||||
) -> httpx.Response | None:
|
||||
from src.services.request.execution_runtime_plan import (
|
||||
ExecutionPlan,
|
||||
ExecutionPlanBody,
|
||||
ExecutionPlanTimeouts,
|
||||
build_execution_plan_body,
|
||||
build_proxy_snapshot,
|
||||
)
|
||||
from src.services.request.execution_runtime_client import (
|
||||
ExecutionRuntimeClient,
|
||||
ExecutionRuntimeClientError,
|
||||
)
|
||||
|
||||
if config.execution_runtime_backend != "rust":
|
||||
return None
|
||||
|
||||
final_headers = dict(headers)
|
||||
if (
|
||||
body is not None
|
||||
and content_type
|
||||
and not any(str(key).lower() == "content-type" for key in final_headers)
|
||||
):
|
||||
final_headers["content-type"] = content_type
|
||||
|
||||
timeout_ms = max(int(timeout_seconds * 1000), 1_000)
|
||||
|
||||
try:
|
||||
proxy_snapshot = await build_proxy_snapshot(proxy_config, label="Kiro")
|
||||
result = await ExecutionRuntimeClient().execute_sync_json(
|
||||
ExecutionPlan(
|
||||
request_id=request_id,
|
||||
candidate_id=None,
|
||||
provider_name="kiro",
|
||||
provider_id="",
|
||||
endpoint_id="",
|
||||
key_id="",
|
||||
method=method,
|
||||
url=url,
|
||||
headers=final_headers,
|
||||
body=(
|
||||
build_execution_plan_body(body, content_type=content_type)
|
||||
if body is not None
|
||||
else ExecutionPlanBody()
|
||||
),
|
||||
stream=False,
|
||||
provider_api_format=provider_api_format,
|
||||
client_api_format=provider_api_format,
|
||||
model_name="kiro",
|
||||
content_type=content_type,
|
||||
proxy=proxy_snapshot,
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=timeout_ms,
|
||||
read_ms=timeout_ms,
|
||||
write_ms=timeout_ms,
|
||||
pool_ms=timeout_ms,
|
||||
total_ms=timeout_ms,
|
||||
),
|
||||
)
|
||||
)
|
||||
except (ExecutionRuntimeClientError, httpx.HTTPError, json.JSONDecodeError) as exc:
|
||||
logger.warning("Kiro Rust HTTP fallback {} {}: {}", method, url, exc)
|
||||
return None
|
||||
except Exception as exc:
|
||||
logger.warning("Kiro Rust HTTP unexpected fallback {} {}: {}", method, 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(method, url, headers=final_headers),
|
||||
headers=response_headers,
|
||||
content=response_body,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["execute_kiro_rust_http_request"]
|
||||
@@ -0,0 +1,346 @@
|
||||
"""Kiro token refresh helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from src.clients.http_client import HTTPClientPool
|
||||
from src.core.logger import logger
|
||||
from src.services.provider.adapters.kiro.headers import build_kiro_ide_tag
|
||||
from src.services.provider.adapters.kiro.models.credentials import KiroAuthConfig
|
||||
from src.services.provider.adapters.kiro.rust_http import execute_kiro_rust_http_request
|
||||
|
||||
IDC_AMZ_USER_AGENT = (
|
||||
"aws-sdk-js/3.738.0 ua/2.1 os/other lang/js md/browser#unknown_unknown "
|
||||
"api/sso-oidc#3.738.0 m/E KiroIDE"
|
||||
)
|
||||
|
||||
_REGION_RE = re.compile(r"^[a-z]{2}-[a-z0-9-]+-\d+$")
|
||||
_HEX64_RE = re.compile(r"^[0-9a-fA-F]{64}$")
|
||||
_UUID_RE = re.compile(
|
||||
r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$"
|
||||
)
|
||||
|
||||
|
||||
def validate_refresh_token(refresh_token: str) -> None:
|
||||
token = str(refresh_token or "").strip()
|
||||
if not token:
|
||||
raise ValueError("missing refresh_token")
|
||||
|
||||
# kiro.rs: length < 100 or contains "..." is considered truncated.
|
||||
if len(token) < 100 or token.endswith("...") or "..." in token:
|
||||
raise ValueError(
|
||||
"refresh_token appears truncated; please export the full token from Kiro IDE"
|
||||
)
|
||||
|
||||
|
||||
def normalize_machine_id(machine_id: str) -> str | None:
|
||||
raw = str(machine_id or "").strip()
|
||||
if not raw:
|
||||
return None
|
||||
|
||||
if _HEX64_RE.fullmatch(raw):
|
||||
return raw.lower()
|
||||
|
||||
if _UUID_RE.fullmatch(raw):
|
||||
without = raw.replace("-", "").lower()
|
||||
return without + without
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def generate_machine_id(cfg: KiroAuthConfig) -> str:
|
||||
normalized = normalize_machine_id(cfg.machine_id or "")
|
||||
if normalized:
|
||||
return normalized
|
||||
|
||||
validate_refresh_token(cfg.refresh_token)
|
||||
seed = f"KotlinNativeAPI/{cfg.refresh_token}".encode("utf-8")
|
||||
return hashlib.sha256(seed).hexdigest()
|
||||
|
||||
|
||||
def is_token_expired(expires_at: int | None, *, skew_seconds: int = 120) -> bool:
|
||||
try:
|
||||
ts = int(expires_at or 0)
|
||||
except Exception:
|
||||
ts = 0
|
||||
if ts <= 0:
|
||||
return True
|
||||
return int(time.time()) >= ts - int(skew_seconds)
|
||||
|
||||
|
||||
def _resolve_region(cfg: KiroAuthConfig) -> str:
|
||||
"""解析 token 刷新端点的 region。"""
|
||||
region = cfg.effective_auth_region()
|
||||
if _REGION_RE.fullmatch(region):
|
||||
return region
|
||||
from src.services.provider.adapters.kiro.constants import DEFAULT_REGION
|
||||
|
||||
return DEFAULT_REGION
|
||||
|
||||
|
||||
def _try_extract_email_from_jwt(token: str) -> str | None:
|
||||
"""尝试从 JWT access_token 中提取 email。
|
||||
|
||||
Kiro Social / IdC 返回的 accessToken 可能是 JWT 格式,
|
||||
payload 中可能包含 email 字段。仅做 base64 解码,不验证签名。
|
||||
失败时静默返回 None。
|
||||
"""
|
||||
try:
|
||||
parts = token.split(".")
|
||||
if len(parts) != 3:
|
||||
return None
|
||||
# base64url decode the payload (second segment)
|
||||
payload_b64 = parts[1]
|
||||
# Add padding
|
||||
padding = 4 - len(payload_b64) % 4
|
||||
if padding != 4:
|
||||
payload_b64 += "=" * padding
|
||||
payload_bytes = base64.urlsafe_b64decode(payload_b64)
|
||||
claims = json.loads(payload_bytes)
|
||||
# Try common email claim keys
|
||||
for key in ("email", "Email", "mail", "upn"):
|
||||
val = claims.get(key)
|
||||
if isinstance(val, str) and "@" in val:
|
||||
return val.strip()
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
async def refresh_social_token(
|
||||
cfg: KiroAuthConfig,
|
||||
*,
|
||||
proxy_config: dict[str, Any] | None,
|
||||
timeout_seconds: float = 30.0,
|
||||
) -> tuple[str, KiroAuthConfig]:
|
||||
"""Refresh access token via Kiro Social refresh endpoint."""
|
||||
validate_refresh_token(cfg.refresh_token)
|
||||
|
||||
region = _resolve_region(cfg)
|
||||
url = f"https://prod.{region}.auth.desktop.kiro.dev/refreshToken"
|
||||
host = f"prod.{region}.auth.desktop.kiro.dev"
|
||||
|
||||
machine_id = generate_machine_id(cfg)
|
||||
kiro_version = (cfg.kiro_version or "").strip() or "0.8.0"
|
||||
ua = build_kiro_ide_tag(kiro_version=kiro_version, machine_id=machine_id)
|
||||
|
||||
body = {"refreshToken": cfg.refresh_token}
|
||||
headers = {
|
||||
"User-Agent": ua,
|
||||
"Host": host,
|
||||
"Accept": "application/json, text/plain, */*",
|
||||
"Content-Type": "application/json",
|
||||
"Connection": "close",
|
||||
"Accept-Encoding": "gzip, compress, deflate, br",
|
||||
}
|
||||
|
||||
resp = await execute_kiro_rust_http_request(
|
||||
method="POST",
|
||||
url=url,
|
||||
headers=headers,
|
||||
body=body,
|
||||
proxy_config=proxy_config,
|
||||
request_id=f"kiro-social-refresh:{region}:{machine_id}",
|
||||
provider_api_format="kiro:social_refresh",
|
||||
content_type="application/json",
|
||||
)
|
||||
if resp is None:
|
||||
client = await HTTPClientPool.get_proxy_client(proxy_config=proxy_config)
|
||||
resp = await client.post(
|
||||
url,
|
||||
headers=headers,
|
||||
json=body,
|
||||
timeout=httpx.Timeout(timeout_seconds),
|
||||
)
|
||||
|
||||
if resp.status_code < 200 or resp.status_code >= 300:
|
||||
body_text = (resp.text or "").strip()[:500]
|
||||
logger.warning(
|
||||
"kiro social refresh error: HTTP {} | {}",
|
||||
resp.status_code,
|
||||
body_text,
|
||||
)
|
||||
raise RuntimeError(f"kiro social refresh failed: HTTP {resp.status_code} | {body_text}")
|
||||
|
||||
data: dict[str, Any]
|
||||
try:
|
||||
data = resp.json()
|
||||
except Exception as e:
|
||||
raise RuntimeError("kiro social refresh: invalid json response") from e
|
||||
|
||||
access_token = str(data.get("accessToken") or "").strip()
|
||||
if not access_token:
|
||||
raise RuntimeError("kiro social refresh returned empty accessToken")
|
||||
|
||||
new_cfg = KiroAuthConfig.from_dict(cfg.to_dict())
|
||||
|
||||
# refreshToken/profileArn may rotate
|
||||
rt = data.get("refreshToken")
|
||||
if isinstance(rt, str) and rt.strip():
|
||||
new_cfg.refresh_token = rt.strip()
|
||||
|
||||
profile_arn = data.get("profileArn")
|
||||
if isinstance(profile_arn, str) and profile_arn.strip():
|
||||
new_cfg.profile_arn = profile_arn.strip()
|
||||
|
||||
expires_in = data.get("expiresIn")
|
||||
try:
|
||||
if expires_in is not None:
|
||||
new_cfg.expires_at = int(time.time()) + int(expires_in)
|
||||
except Exception:
|
||||
new_cfg.expires_at = int(time.time()) + 3600
|
||||
|
||||
# Persist computed machine_id if user didn't provide one.
|
||||
if not (cfg.machine_id or "").strip():
|
||||
new_cfg.machine_id = machine_id
|
||||
|
||||
# 尝试从 accessToken 中提取 email(如果尚未设置)
|
||||
if not (new_cfg.email or "").strip():
|
||||
extracted_email = _try_extract_email_from_jwt(access_token)
|
||||
if extracted_email:
|
||||
new_cfg.email = extracted_email
|
||||
logger.debug("kiro social: extracted email from accessToken: {}", extracted_email)
|
||||
|
||||
# 缓存 access_token
|
||||
new_cfg.access_token = access_token
|
||||
|
||||
return access_token, new_cfg
|
||||
|
||||
|
||||
async def refresh_idc_token(
|
||||
cfg: KiroAuthConfig,
|
||||
*,
|
||||
proxy_config: dict[str, Any] | None,
|
||||
timeout_seconds: float = 30.0,
|
||||
) -> tuple[str, KiroAuthConfig]:
|
||||
"""Refresh access token via AWS SSO OIDC endpoint (IdC)."""
|
||||
validate_refresh_token(cfg.refresh_token)
|
||||
|
||||
if not (cfg.client_id or "").strip() or not (cfg.client_secret or "").strip():
|
||||
raise ValueError("idc refresh requires client_id and client_secret")
|
||||
|
||||
region = _resolve_region(cfg)
|
||||
url = f"https://oidc.{region}.amazonaws.com/token"
|
||||
host = f"oidc.{region}.amazonaws.com"
|
||||
|
||||
body = {
|
||||
"clientId": cfg.client_id,
|
||||
"clientSecret": cfg.client_secret,
|
||||
"refreshToken": cfg.refresh_token,
|
||||
"grantType": "refresh_token",
|
||||
}
|
||||
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Host": host,
|
||||
"x-amz-user-agent": IDC_AMZ_USER_AGENT,
|
||||
"User-Agent": "node",
|
||||
"Accept": "*/*",
|
||||
}
|
||||
|
||||
resp = await execute_kiro_rust_http_request(
|
||||
method="POST",
|
||||
url=url,
|
||||
headers=headers,
|
||||
body=body,
|
||||
proxy_config=proxy_config,
|
||||
request_id=f"kiro-idc-refresh:{region}",
|
||||
provider_api_format="kiro:idc_refresh",
|
||||
content_type="application/json",
|
||||
)
|
||||
if resp is None:
|
||||
client = await HTTPClientPool.get_proxy_client(proxy_config=proxy_config)
|
||||
resp = await client.post(
|
||||
url,
|
||||
headers=headers,
|
||||
json=body,
|
||||
timeout=httpx.Timeout(timeout_seconds),
|
||||
)
|
||||
|
||||
if resp.status_code < 200 or resp.status_code >= 300:
|
||||
body_text = (resp.text or "").strip()[:500]
|
||||
logger.warning(
|
||||
"kiro idc refresh error: HTTP {} | {}",
|
||||
resp.status_code,
|
||||
body_text,
|
||||
)
|
||||
raise RuntimeError(f"kiro idc refresh failed: HTTP {resp.status_code} | {body_text}")
|
||||
|
||||
data: dict[str, Any]
|
||||
try:
|
||||
data = resp.json()
|
||||
except Exception as e:
|
||||
raise RuntimeError("kiro idc refresh: invalid json response") from e
|
||||
|
||||
access_token = str(data.get("accessToken") or "").strip()
|
||||
if not access_token:
|
||||
raise RuntimeError("kiro idc refresh returned empty accessToken")
|
||||
|
||||
new_cfg = KiroAuthConfig.from_dict(cfg.to_dict())
|
||||
|
||||
rt = data.get("refreshToken")
|
||||
if isinstance(rt, str) and rt.strip():
|
||||
new_cfg.refresh_token = rt.strip()
|
||||
|
||||
expires_in = data.get("expiresIn")
|
||||
try:
|
||||
if expires_in is not None:
|
||||
new_cfg.expires_at = int(time.time()) + int(expires_in)
|
||||
except Exception:
|
||||
new_cfg.expires_at = int(time.time()) + 3600
|
||||
|
||||
# Persist computed machine_id if user didn't provide one.
|
||||
if not (cfg.machine_id or "").strip():
|
||||
new_cfg.machine_id = generate_machine_id(cfg)
|
||||
|
||||
# 尝试从 accessToken 中提取 email(如果尚未设置)
|
||||
if not (new_cfg.email or "").strip():
|
||||
extracted_email = _try_extract_email_from_jwt(access_token)
|
||||
if extracted_email:
|
||||
new_cfg.email = extracted_email
|
||||
logger.debug("kiro idc: extracted email from accessToken: {}", extracted_email)
|
||||
|
||||
# 缓存 access_token
|
||||
new_cfg.access_token = access_token
|
||||
|
||||
return access_token, new_cfg
|
||||
|
||||
|
||||
async def refresh_access_token(
|
||||
cfg: KiroAuthConfig,
|
||||
*,
|
||||
proxy_config: dict[str, Any] | None,
|
||||
timeout_seconds: float = 30.0,
|
||||
) -> tuple[str, KiroAuthConfig]:
|
||||
method = (cfg.auth_method or "social").strip().lower()
|
||||
if method == "idc":
|
||||
return await refresh_idc_token(
|
||||
cfg,
|
||||
proxy_config=proxy_config,
|
||||
timeout_seconds=timeout_seconds,
|
||||
)
|
||||
return await refresh_social_token(
|
||||
cfg,
|
||||
proxy_config=proxy_config,
|
||||
timeout_seconds=timeout_seconds,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"IDC_AMZ_USER_AGENT",
|
||||
"generate_machine_id",
|
||||
"is_token_expired",
|
||||
"normalize_machine_id",
|
||||
"refresh_access_token",
|
||||
"refresh_idc_token",
|
||||
"refresh_social_token",
|
||||
"validate_refresh_token",
|
||||
]
|
||||
253
_deprecated_py_src/services/provider/adapters/kiro/usage.py
Normal file
253
_deprecated_py_src/services/provider/adapters/kiro/usage.py
Normal file
@@ -0,0 +1,253 @@
|
||||
"""Kiro usage/quota fetching utilities."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from src.clients.http_client import HTTPClientPool
|
||||
from src.core.logger import logger
|
||||
from src.services.provider.adapters.kiro.headers import (
|
||||
build_user_agent_usage,
|
||||
build_x_amz_user_agent_usage,
|
||||
)
|
||||
from src.services.provider.adapters.kiro.models.credentials import KiroAuthConfig
|
||||
from src.services.provider.adapters.kiro.models.usage_limits import (
|
||||
UsageLimitsResponse,
|
||||
calculate_current_usage,
|
||||
calculate_total_usage_limit,
|
||||
)
|
||||
from src.services.provider.adapters.kiro.request import get_profile_arn_for_payload
|
||||
from src.services.provider.adapters.kiro.rust_http import execute_kiro_rust_http_request
|
||||
from src.services.provider.adapters.kiro.token_manager import (
|
||||
generate_machine_id,
|
||||
is_token_expired,
|
||||
refresh_access_token,
|
||||
)
|
||||
|
||||
|
||||
class KiroAccountBannedException(Exception):
|
||||
"""Kiro 账户被封禁异常"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
message: str = "账户已封禁",
|
||||
status_code: int = 403,
|
||||
reason: str | None = None,
|
||||
):
|
||||
super().__init__(message)
|
||||
self.message = message
|
||||
self.status_code = status_code
|
||||
self.reason = reason
|
||||
|
||||
|
||||
async def fetch_kiro_usage_limits(
|
||||
auth_config: dict[str, Any],
|
||||
proxy_config: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
调用 Kiro getUsageLimits API 获取使用额度信息
|
||||
|
||||
Args:
|
||||
auth_config: 解密后的 KiroAuthConfig 数据
|
||||
proxy_config: 代理配置(可选)
|
||||
|
||||
Returns:
|
||||
包含 usage_data 和 updated_auth_config 的字典
|
||||
|
||||
Raises:
|
||||
RuntimeError: 请求失败时抛出
|
||||
"""
|
||||
cfg = KiroAuthConfig.from_dict(auth_config)
|
||||
|
||||
# 检查是否有缓存的 access_token 且未过期
|
||||
access_token: str | None = None
|
||||
updated_cfg: KiroAuthConfig | None = None
|
||||
|
||||
if cfg.access_token and not is_token_expired(cfg.expires_at):
|
||||
# 使用缓存的 token
|
||||
access_token = cfg.access_token
|
||||
updated_cfg = cfg
|
||||
logger.debug("[KIRO_QUOTA] 使用缓存的 access_token")
|
||||
else:
|
||||
# token 过期或不存在,需要刷新
|
||||
logger.debug("[KIRO_QUOTA] Token 已过期或不存在,正在刷新...")
|
||||
access_token, updated_cfg = await refresh_access_token(cfg, proxy_config=proxy_config)
|
||||
|
||||
if not access_token:
|
||||
raise RuntimeError("无法获取 Kiro access_token")
|
||||
|
||||
# 构建请求
|
||||
effective_cfg = updated_cfg or cfg
|
||||
region = effective_cfg.effective_api_region()
|
||||
host = f"q.{region}.amazonaws.com"
|
||||
machine_id = generate_machine_id(effective_cfg)
|
||||
kiro_version = (effective_cfg.kiro_version or "0.8.0").strip() or "0.8.0"
|
||||
|
||||
# 构建 URL(添加 isEmailRequired=true 获取邮箱)
|
||||
url = f"https://{host}/getUsageLimits?origin=AI_EDITOR&resourceType=AGENTIC_REQUEST&isEmailRequired=true"
|
||||
|
||||
profile_arn = get_profile_arn_for_payload(effective_cfg)
|
||||
if profile_arn:
|
||||
from urllib.parse import quote
|
||||
|
||||
url += f"&profileArn={quote(profile_arn, safe='')}"
|
||||
|
||||
logger.debug("[KIRO_QUOTA] 请求 URL: {}", url)
|
||||
|
||||
# 构建 headers
|
||||
headers = {
|
||||
"x-amz-user-agent": build_x_amz_user_agent_usage(
|
||||
kiro_version=kiro_version, machine_id=machine_id
|
||||
),
|
||||
"User-Agent": build_user_agent_usage(kiro_version=kiro_version, machine_id=machine_id),
|
||||
"host": host,
|
||||
"amz-sdk-invocation-id": str(uuid.uuid4()),
|
||||
"amz-sdk-request": "attempt=1; max=1",
|
||||
"Authorization": f"Bearer {access_token}",
|
||||
"Connection": "close",
|
||||
}
|
||||
|
||||
response = await execute_kiro_rust_http_request(
|
||||
method="GET",
|
||||
url=url,
|
||||
headers=headers,
|
||||
body=None,
|
||||
proxy_config=proxy_config,
|
||||
request_id=f"kiro-usage:{region}:{machine_id}",
|
||||
provider_api_format="kiro:usage",
|
||||
)
|
||||
if response is None:
|
||||
client = await HTTPClientPool.get_proxy_client(proxy_config=proxy_config)
|
||||
response = await client.get(url, headers=headers, timeout=httpx.Timeout(30.0))
|
||||
|
||||
if response.status_code != 200:
|
||||
response_text = (response.text or "").strip()
|
||||
|
||||
# 检测账户封禁(403/423 均视为封禁)
|
||||
# 对于 Kiro,403 表示账户权限被拒绝(无论是封禁、权限不足还是其他原因),
|
||||
# 都应该标记为异常状态,以便管理员及时处理
|
||||
is_banned = False
|
||||
ban_reason = None
|
||||
if response.status_code in (403, 423):
|
||||
is_banned = True
|
||||
# 检测封禁相关错误(用于提供更详细的原因说明)
|
||||
banned_keywords = [
|
||||
"AccountSuspendedException",
|
||||
"account.*suspend",
|
||||
"account.*banned",
|
||||
"account.*disabled",
|
||||
"account.*access.*denied",
|
||||
]
|
||||
for keyword in banned_keywords:
|
||||
if re.search(keyword, response_text, re.IGNORECASE):
|
||||
ban_reason = (
|
||||
response_text[:200] if response_text else f"HTTP {response.status_code}"
|
||||
)
|
||||
break
|
||||
# 如果没有匹配到特定关键词,使用通用原因
|
||||
if not ban_reason:
|
||||
if response.status_code == 423:
|
||||
ban_reason = response_text[:200] if response_text else "HTTP 423 Locked"
|
||||
else:
|
||||
ban_reason = response_text[:200] if response_text else "HTTP 403 权限被拒绝"
|
||||
|
||||
error_msg = {
|
||||
401: "认证失败,Token 无效或已过期",
|
||||
403: "账户异常,权限被拒绝",
|
||||
423: "账户已封禁",
|
||||
429: "请求过于频繁,已被限流",
|
||||
}.get(response.status_code, "获取使用额度失败")
|
||||
if 500 <= response.status_code < 600:
|
||||
error_msg = "服务器错误,AWS 服务暂时不可用"
|
||||
logger.debug(
|
||||
"kiro usage API error: HTTP {} | {}",
|
||||
response.status_code,
|
||||
response_text[:200],
|
||||
)
|
||||
|
||||
# 如果检测到封禁,抛出带有封禁标记的异常
|
||||
if is_banned:
|
||||
raise KiroAccountBannedException(
|
||||
message=error_msg,
|
||||
status_code=response.status_code,
|
||||
reason=ban_reason,
|
||||
)
|
||||
|
||||
raise RuntimeError(f"{error_msg}: HTTP {response.status_code}")
|
||||
|
||||
try:
|
||||
data = response.json()
|
||||
except Exception as exc:
|
||||
raise RuntimeError(f"获取使用额度成功但响应解析失败: HTTP {response.status_code}") from exc
|
||||
|
||||
# 返回刷新后的配置(用于更新 auth_config)
|
||||
return {
|
||||
"usage_data": data,
|
||||
"updated_auth_config": updated_cfg.to_dict() if updated_cfg else None,
|
||||
}
|
||||
|
||||
|
||||
def parse_kiro_usage_response(data: dict) -> dict | None:
|
||||
"""
|
||||
解析 Kiro getUsageLimits API 响应,提取限额信息和用户邮箱
|
||||
|
||||
返回格式与 kiro.rs BalanceResponse 类似:
|
||||
- subscription_title: 订阅类型(如 "KIRO PRO+")
|
||||
- current_usage: 当前使用量
|
||||
- usage_limit: 使用限额
|
||||
- remaining: 剩余额度
|
||||
- usage_percentage: 使用百分比
|
||||
- next_reset_at: 下次重置时间(Unix 时间戳)
|
||||
- email: 用户邮箱(通过 isEmailRequired=true 获取)
|
||||
"""
|
||||
if not data:
|
||||
return None
|
||||
|
||||
usage_resp = UsageLimitsResponse.from_dict(data)
|
||||
|
||||
current_usage = calculate_current_usage(usage_resp)
|
||||
usage_limit = calculate_total_usage_limit(usage_resp)
|
||||
remaining = max(usage_limit - current_usage, 0.0)
|
||||
usage_percentage = (current_usage / usage_limit * 100.0) if usage_limit > 0 else 0.0
|
||||
usage_percentage = min(usage_percentage, 100.0)
|
||||
|
||||
result: dict[str, Any] = {
|
||||
"current_usage": current_usage,
|
||||
"usage_limit": usage_limit,
|
||||
"remaining": remaining,
|
||||
"usage_percentage": usage_percentage,
|
||||
}
|
||||
|
||||
# 订阅类型
|
||||
if usage_resp.subscription_info and usage_resp.subscription_info.subscription_title:
|
||||
result["subscription_title"] = usage_resp.subscription_info.subscription_title
|
||||
|
||||
# 下次重置时间
|
||||
if usage_resp.next_date_reset is not None:
|
||||
result["next_reset_at"] = usage_resp.next_date_reset
|
||||
elif usage_resp.usage_breakdown_list and usage_resp.usage_breakdown_list[0].next_date_reset:
|
||||
result["next_reset_at"] = usage_resp.usage_breakdown_list[0].next_date_reset
|
||||
|
||||
# 解析用户邮箱(从 desktopUserInfo 或 userInfo 中获取)
|
||||
user_info = data.get("desktopUserInfo") or data.get("userInfo") or {}
|
||||
if isinstance(user_info, dict):
|
||||
email = user_info.get("email")
|
||||
if isinstance(email, str) and email.strip():
|
||||
result["email"] = email.strip()
|
||||
|
||||
# 添加更新时间戳
|
||||
result["updated_at"] = int(time.time())
|
||||
|
||||
return result
|
||||
|
||||
|
||||
__all__ = [
|
||||
"KiroAccountBannedException",
|
||||
"fetch_kiro_usage_limits",
|
||||
"parse_kiro_usage_response",
|
||||
]
|
||||
@@ -0,0 +1,63 @@
|
||||
"""Vertex AI 认证处理。
|
||||
|
||||
- Service Account: GCP SA JSON → JWT → Access Token → Bearer header
|
||||
- API Key: 通过 URL ?key= 查询参数认证,auth 层返回 None(由 transport hook 处理)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
from src.core.provider_auth_types import ProviderAuthInfo
|
||||
|
||||
|
||||
async def _auth_service_account(key: Any, endpoint: Any | None = None) -> ProviderAuthInfo:
|
||||
"""Service Account 认证:SA JSON → JWT → Access Token。"""
|
||||
from src.core.crypto import crypto_service
|
||||
from src.core.exceptions import InvalidRequestException
|
||||
from src.core.vertex_auth import VertexAuthError, VertexAuthService
|
||||
|
||||
try:
|
||||
# 优先从 auth_config 读取,兼容从 api_key 读取(过渡期)
|
||||
encrypted_auth_config = getattr(key, "auth_config", None)
|
||||
if encrypted_auth_config:
|
||||
if isinstance(encrypted_auth_config, dict):
|
||||
sa_json = encrypted_auth_config
|
||||
else:
|
||||
decrypted_config = crypto_service.decrypt(encrypted_auth_config)
|
||||
sa_json = json.loads(decrypted_config)
|
||||
else:
|
||||
# 兼容旧数据:从 api_key 读取
|
||||
decrypted_key = crypto_service.decrypt(key.api_key)
|
||||
if decrypted_key == "__placeholder__":
|
||||
raise InvalidRequestException("认证配置丢失,请重新添加该密钥。")
|
||||
sa_json = json.loads(decrypted_key)
|
||||
|
||||
if not isinstance(sa_json, dict):
|
||||
raise InvalidRequestException("Service Account JSON 无效,请重新添加该密钥。")
|
||||
|
||||
# 获取 Access Token(注入代理配置)
|
||||
from src.services.provider.auth import _get_proxy_config
|
||||
from src.services.proxy_node.resolver import build_proxy_client_kwargs
|
||||
|
||||
effective_proxy = _get_proxy_config(key, endpoint)
|
||||
|
||||
service = VertexAuthService(sa_json)
|
||||
access_token = await service.get_access_token(
|
||||
httpx_client_kwargs=build_proxy_client_kwargs(effective_proxy, timeout=30),
|
||||
)
|
||||
|
||||
return ProviderAuthInfo(
|
||||
auth_header="Authorization",
|
||||
auth_value=f"Bearer {access_token}",
|
||||
decrypted_auth_config=sa_json,
|
||||
)
|
||||
except InvalidRequestException:
|
||||
raise
|
||||
except VertexAuthError as e:
|
||||
raise InvalidRequestException(f"Vertex AI 认证失败:{e}")
|
||||
except json.JSONDecodeError:
|
||||
raise InvalidRequestException("Service Account JSON 格式无效,请重新添加该密钥。")
|
||||
except Exception:
|
||||
raise InvalidRequestException("Vertex AI 认证失败,请检查 Key 的 auth_config")
|
||||
@@ -0,0 +1,45 @@
|
||||
"""Vertex AI 常量配置。
|
||||
|
||||
从 transport.py 迁移,集中管理 Vertex AI 模型格式映射和 region 配置。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
# Vertex AI 模型前缀到 API 格式的映射
|
||||
# 用于 provider_type=vertex_ai 时,根据模型名动态确定实际的请求/响应格式
|
||||
# 格式:前缀 -> endpoint signature(family:kind)
|
||||
MODEL_FORMAT_MAPPING: dict[str, str] = {
|
||||
"claude-": "claude:chat", # Anthropic Claude 模型
|
||||
"gemini-": "gemini:chat", # Google Gemini 模型
|
||||
"imagen-": "gemini:chat", # Google Imagen 模型(使用 Gemini chat 格式)
|
||||
}
|
||||
|
||||
# Vertex AI 默认 endpoint signature(当模型前缀不匹配时)
|
||||
DEFAULT_FORMAT: str = "gemini:chat"
|
||||
|
||||
# Vertex AI 模型默认 region 映射
|
||||
# 用户可以通过 auth_config.model_regions 覆盖
|
||||
DEFAULT_MODEL_REGIONS: dict[str, str] = {
|
||||
# Gemini 3 系列(使用 global)
|
||||
"gemini-3.1-pro-preview": "global",
|
||||
"gemini-3-pro-image-preview": "global",
|
||||
# Gemini 2.0 系列
|
||||
"gemini-2.0-flash": "us-central1",
|
||||
"gemini-2.0-flash-exp": "us-central1",
|
||||
"gemini-2.0-flash-001": "us-central1",
|
||||
"gemini-2.0-pro-exp": "us-central1",
|
||||
"gemini-2.0-flash-exp-image-generation": "us-central1",
|
||||
# Gemini 1.5 系列
|
||||
"gemini-1.5-pro": "us-central1",
|
||||
"gemini-1.5-pro-001": "us-central1",
|
||||
"gemini-1.5-pro-002": "us-central1",
|
||||
"gemini-1.5-flash": "us-central1",
|
||||
"gemini-1.5-flash-001": "us-central1",
|
||||
"gemini-1.5-flash-002": "us-central1",
|
||||
# Imagen 系列
|
||||
"imagen-3.0-generate-001": "us-central1",
|
||||
"imagen-3.0-fast-generate-001": "us-central1",
|
||||
}
|
||||
|
||||
# API Key 认证的全局端点
|
||||
API_KEY_BASE_URL = "https://aiplatform.googleapis.com"
|
||||
@@ -0,0 +1,451 @@
|
||||
"""Vertex AI provider plugin — 统一注册入口。
|
||||
|
||||
注册 Vertex AI 对各通用 registry / capability registry 的 hooks:
|
||||
- Transport Hook (URL 构建:Gemini 走 Express mode,Claude 走 Service Account)
|
||||
- Model Fetcher (专用上游模型获取链路)
|
||||
- Provider Format Capability(跨格式支持:同一 Provider 可配置 Gemini / Claude)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.core.vertex_auth import VertexAuthError, VertexAuthService
|
||||
from src.services.provider.adapters.vertex_ai.transport import get_effective_format
|
||||
|
||||
# Vertex AI 公共 API 根
|
||||
_VERTEX_API_BASE = "https://aiplatform.googleapis.com"
|
||||
|
||||
_MODEL_PAGE_SIZE = 100
|
||||
_MODEL_MAX_PAGES = 20
|
||||
|
||||
|
||||
def _normalize_extra_headers(raw: Any) -> dict[str, str]:
|
||||
if not isinstance(raw, dict):
|
||||
return {}
|
||||
return {str(k): str(v) for k, v in raw.items() if k and v is not None}
|
||||
|
||||
|
||||
def _looks_like_service_account(auth_config: dict[str, Any] | None) -> bool:
|
||||
if not isinstance(auth_config, dict):
|
||||
return False
|
||||
return all(
|
||||
isinstance(auth_config.get(k), str) and str(auth_config.get(k)).strip()
|
||||
for k in ("client_email", "private_key", "project_id")
|
||||
)
|
||||
|
||||
|
||||
def _extract_model_id(raw_name: str) -> str:
|
||||
name = str(raw_name or "").strip()
|
||||
if not name:
|
||||
return ""
|
||||
if "/models/" in name:
|
||||
return name.split("/models/", 1)[-1].strip()
|
||||
if name.startswith("models/"):
|
||||
return name.split("models/", 1)[-1].strip()
|
||||
return name
|
||||
|
||||
|
||||
def _extract_publisher(item: dict[str, Any], fallback: str | None = None) -> str | None:
|
||||
publisher = item.get("publisher")
|
||||
if isinstance(publisher, str) and publisher.strip():
|
||||
return publisher.strip()
|
||||
|
||||
raw_name = item.get("name")
|
||||
if isinstance(raw_name, str) and "/publishers/" in raw_name:
|
||||
try:
|
||||
after = raw_name.split("/publishers/", 1)[1]
|
||||
candidate = after.split("/", 1)[0].strip()
|
||||
if candidate:
|
||||
return candidate
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return fallback
|
||||
|
||||
|
||||
def _extract_items(data: Any) -> list[dict[str, Any]]:
|
||||
if isinstance(data, list):
|
||||
return [item for item in data if isinstance(item, dict)]
|
||||
if not isinstance(data, dict):
|
||||
return []
|
||||
|
||||
for key in ("publisherModels", "models", "data", "items"):
|
||||
value = data.get(key)
|
||||
if isinstance(value, list):
|
||||
return [item for item in value if isinstance(item, dict)]
|
||||
return []
|
||||
|
||||
|
||||
def _parse_models_payload(
|
||||
data: Any,
|
||||
*,
|
||||
auth_config: dict[str, Any] | None,
|
||||
fallback_publisher: str | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
models: list[dict[str, Any]] = []
|
||||
for item in _extract_items(data):
|
||||
raw_name = item.get("id") or item.get("name") or item.get("model")
|
||||
if not isinstance(raw_name, str):
|
||||
continue
|
||||
|
||||
model_id = _extract_model_id(raw_name)
|
||||
if not model_id:
|
||||
continue
|
||||
|
||||
display_name_raw = (
|
||||
item.get("displayName") or item.get("display_name") or item.get("title") or model_id
|
||||
)
|
||||
display_name = (
|
||||
str(display_name_raw).strip() if isinstance(display_name_raw, str) else model_id
|
||||
)
|
||||
if not display_name:
|
||||
display_name = model_id
|
||||
|
||||
models.append(
|
||||
{
|
||||
"id": model_id,
|
||||
"owned_by": _extract_publisher(item, fallback=fallback_publisher),
|
||||
"display_name": display_name,
|
||||
"api_format": get_effective_format(model_id, auth_config),
|
||||
}
|
||||
)
|
||||
|
||||
return models
|
||||
|
||||
|
||||
def _build_google_publisher_list_url(base_url: str) -> str:
|
||||
base = str(base_url or "").rstrip("/")
|
||||
if not base:
|
||||
base = _VERTEX_API_BASE
|
||||
|
||||
if base.endswith("/v1"):
|
||||
return f"{base}/publishers/google/models"
|
||||
if base.endswith("/v1beta"):
|
||||
return f"{base}/publishers/google/models"
|
||||
return f"{base}/v1/publishers/google/models"
|
||||
|
||||
|
||||
def _iter_endpoint_base_urls(ctx: Any) -> list[str]:
|
||||
seen: set[str] = set()
|
||||
urls: list[str] = []
|
||||
for cfg in (ctx.format_to_endpoint or {}).values():
|
||||
base_url = str(getattr(cfg, "base_url", "") or "").strip()
|
||||
if not base_url:
|
||||
continue
|
||||
norm = base_url.rstrip("/")
|
||||
if norm in seen:
|
||||
continue
|
||||
seen.add(norm)
|
||||
urls.append(norm)
|
||||
|
||||
if _VERTEX_API_BASE not in seen:
|
||||
urls.append(_VERTEX_API_BASE)
|
||||
return urls
|
||||
|
||||
|
||||
def _get_endpoint_headers(ctx: Any, api_format: str) -> dict[str, str]:
|
||||
cfg = (ctx.format_to_endpoint or {}).get(api_format)
|
||||
return _normalize_extra_headers(getattr(cfg, "extra_headers", None))
|
||||
|
||||
|
||||
def _dedupe_models(models: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
seen: set[str] = set()
|
||||
result: list[dict[str, Any]] = []
|
||||
for model in models:
|
||||
model_id = str(model.get("id", "")).strip()
|
||||
api_format = str(model.get("api_format", "")).strip()
|
||||
if not model_id:
|
||||
continue
|
||||
unique_key = f"{model_id}:{api_format}"
|
||||
if unique_key in seen:
|
||||
continue
|
||||
seen.add(unique_key)
|
||||
result.append(model)
|
||||
return result
|
||||
|
||||
|
||||
def _is_soft_not_found(error: str) -> bool:
|
||||
return str(error).strip().startswith("HTTP 404:")
|
||||
|
||||
|
||||
def _iter_regions(auth_config: dict[str, Any] | None) -> list[str]:
|
||||
seen: set[str] = set()
|
||||
regions: list[str] = []
|
||||
|
||||
def _add(raw: Any) -> None:
|
||||
if not isinstance(raw, str):
|
||||
return
|
||||
region = raw.strip()
|
||||
if not region or region in seen:
|
||||
return
|
||||
seen.add(region)
|
||||
regions.append(region)
|
||||
|
||||
if isinstance(auth_config, dict):
|
||||
_add(auth_config.get("region"))
|
||||
model_regions = auth_config.get("model_regions")
|
||||
if isinstance(model_regions, dict):
|
||||
for region in model_regions.values():
|
||||
_add(region)
|
||||
|
||||
_add("global")
|
||||
_add("us-central1")
|
||||
return regions
|
||||
|
||||
|
||||
async def _fetch_models_from_url(
|
||||
client: httpx.AsyncClient,
|
||||
*,
|
||||
url: str,
|
||||
headers: dict[str, str],
|
||||
params: dict[str, Any],
|
||||
auth_config: dict[str, Any] | None,
|
||||
fallback_publisher: str | None = None,
|
||||
) -> tuple[list[dict[str, Any]], str | None, bool]:
|
||||
all_models: list[dict[str, Any]] = []
|
||||
next_page_token: str | None = None
|
||||
has_success = False
|
||||
|
||||
for _ in range(_MODEL_MAX_PAGES):
|
||||
req_params = dict(params)
|
||||
if next_page_token:
|
||||
req_params["pageToken"] = next_page_token
|
||||
|
||||
try:
|
||||
resp = await client.get(url, headers=headers, params=req_params)
|
||||
except httpx.TimeoutException:
|
||||
return [], "timeout", has_success
|
||||
except Exception as exc:
|
||||
return [], f"request error: {exc}", has_success
|
||||
|
||||
if resp.status_code != 200:
|
||||
body = resp.text[:500] if resp.text else "(empty)"
|
||||
return [], f"HTTP {resp.status_code}: {body}", has_success
|
||||
|
||||
has_success = True
|
||||
|
||||
try:
|
||||
payload = resp.json()
|
||||
except Exception:
|
||||
body = resp.text[:500] if resp.text else "(empty)"
|
||||
return [], f"invalid json body: {body}", has_success
|
||||
|
||||
all_models.extend(
|
||||
_parse_models_payload(
|
||||
payload,
|
||||
auth_config=auth_config,
|
||||
fallback_publisher=fallback_publisher,
|
||||
)
|
||||
)
|
||||
|
||||
if not isinstance(payload, dict):
|
||||
break
|
||||
|
||||
token = payload.get("nextPageToken")
|
||||
next_page_token = str(token).strip() if isinstance(token, str) else None
|
||||
if not next_page_token:
|
||||
break
|
||||
|
||||
return all_models, None, has_success
|
||||
|
||||
|
||||
async def _fetch_models_vertex_api_key(
|
||||
client: httpx.AsyncClient,
|
||||
*,
|
||||
ctx: Any,
|
||||
auth_config: dict[str, Any] | None,
|
||||
) -> tuple[list[dict[str, Any]], list[str], bool]:
|
||||
"""API Key 仅抓取 Vertex AI Express mode 的 Google publisher models。"""
|
||||
api_key = str(ctx.api_key_value or "").strip()
|
||||
if not api_key or api_key == "__placeholder__":
|
||||
return [], ["vertex_ai(api_key): missing api key"], False
|
||||
|
||||
all_models: list[dict[str, Any]] = []
|
||||
hard_errors: list[str] = []
|
||||
soft_errors: list[str] = []
|
||||
has_success = False
|
||||
|
||||
endpoint_headers = _get_endpoint_headers(ctx, "gemini:chat")
|
||||
vertex_list_urls = [
|
||||
_build_google_publisher_list_url(base) for base in _iter_endpoint_base_urls(ctx)
|
||||
]
|
||||
|
||||
# Vertex Express mode list (publisher=google)
|
||||
for url in vertex_list_urls:
|
||||
headers = {"Accept": "application/json", **endpoint_headers}
|
||||
models, err, success = await _fetch_models_from_url(
|
||||
client,
|
||||
url=url,
|
||||
headers=headers,
|
||||
params={"key": api_key, "pageSize": _MODEL_PAGE_SIZE},
|
||||
auth_config=auth_config,
|
||||
fallback_publisher="google",
|
||||
)
|
||||
if success:
|
||||
has_success = True
|
||||
if err:
|
||||
labeled = f"{url}: {err}"
|
||||
if _is_soft_not_found(err):
|
||||
soft_errors.append(labeled)
|
||||
else:
|
||||
hard_errors.append(labeled)
|
||||
continue
|
||||
all_models.extend(models)
|
||||
|
||||
deduped = _dedupe_models(all_models)
|
||||
if deduped:
|
||||
return deduped, hard_errors, has_success or True
|
||||
|
||||
if hard_errors:
|
||||
return [], hard_errors, has_success
|
||||
if soft_errors:
|
||||
return [], [soft_errors[0]], has_success
|
||||
return [], [], has_success
|
||||
|
||||
|
||||
async def _fetch_models_vertex_service_account(
|
||||
client: httpx.AsyncClient,
|
||||
*,
|
||||
ctx: Any,
|
||||
auth_config: dict[str, Any] | None,
|
||||
client_kwargs: dict[str, Any],
|
||||
) -> tuple[list[dict[str, Any]], list[str], bool]:
|
||||
"""Service Account 抓取 Vertex AI Google + Anthropic publisher models。"""
|
||||
if not isinstance(auth_config, dict):
|
||||
return [], ["vertex_ai(service_account): missing auth_config"], False
|
||||
|
||||
try:
|
||||
auth_service = VertexAuthService(auth_config)
|
||||
access_token = await auth_service.get_access_token(httpx_client_kwargs=client_kwargs)
|
||||
except VertexAuthError as exc:
|
||||
return [], [f"vertex_ai(service_account): auth failed: {exc}"], False
|
||||
except Exception as exc:
|
||||
return [], [f"vertex_ai(service_account): auth failed: {exc}"], False
|
||||
|
||||
project_id = str(auth_config.get("project_id") or "").strip()
|
||||
if not project_id:
|
||||
return [], ["vertex_ai(service_account): missing project_id"], False
|
||||
|
||||
all_models: list[dict[str, Any]] = []
|
||||
hard_errors: list[str] = []
|
||||
soft_errors: list[str] = []
|
||||
has_success = False
|
||||
|
||||
gemini_headers = {"Accept": "application/json", **_get_endpoint_headers(ctx, "gemini:chat")}
|
||||
claude_headers = {"Accept": "application/json", **_get_endpoint_headers(ctx, "claude:chat")}
|
||||
gemini_headers["Authorization"] = f"Bearer {access_token}"
|
||||
claude_headers["Authorization"] = f"Bearer {access_token}"
|
||||
|
||||
for region in _iter_regions(auth_config):
|
||||
base = (
|
||||
_VERTEX_API_BASE
|
||||
if region == "global"
|
||||
else f"https://{region}-aiplatform.googleapis.com"
|
||||
)
|
||||
|
||||
requests = [
|
||||
(
|
||||
"google",
|
||||
f"{base}/v1/projects/{project_id}/locations/{region}/publishers/google/models",
|
||||
gemini_headers,
|
||||
),
|
||||
(
|
||||
"anthropic",
|
||||
f"{base}/v1/projects/{project_id}/locations/{region}/publishers/anthropic/models",
|
||||
claude_headers,
|
||||
),
|
||||
]
|
||||
|
||||
for publisher, url, headers in requests:
|
||||
models, err, success = await _fetch_models_from_url(
|
||||
client,
|
||||
url=url,
|
||||
headers=headers,
|
||||
params={"pageSize": _MODEL_PAGE_SIZE},
|
||||
auth_config=auth_config,
|
||||
fallback_publisher=publisher,
|
||||
)
|
||||
if success:
|
||||
has_success = True
|
||||
if err:
|
||||
labeled = f"{url}: {err}"
|
||||
if _is_soft_not_found(err):
|
||||
soft_errors.append(labeled)
|
||||
else:
|
||||
hard_errors.append(labeled)
|
||||
continue
|
||||
all_models.extend(models)
|
||||
|
||||
deduped = _dedupe_models(all_models)
|
||||
if deduped:
|
||||
return deduped, hard_errors, has_success or True
|
||||
|
||||
if hard_errors:
|
||||
return [], hard_errors, has_success
|
||||
if soft_errors:
|
||||
return [], [soft_errors[0]], has_success
|
||||
return [], [], has_success
|
||||
|
||||
|
||||
async def fetch_models_vertex_ai(
|
||||
ctx: Any,
|
||||
timeout_seconds: float,
|
||||
) -> tuple[list[dict], list[str], bool, dict[str, Any] | None]:
|
||||
"""Vertex AI 专用模型获取链路。
|
||||
|
||||
- API Key: 仅请求 Vertex AI Express mode 的 Gemini models
|
||||
- Service Account: 使用 SA 凭证换取 Bearer Token,按 region 查询 Gemini + Claude models
|
||||
"""
|
||||
from src.services.proxy_node.resolver import build_proxy_client_kwargs
|
||||
|
||||
auth_config = ctx.auth_config if isinstance(ctx.auth_config, dict) else None
|
||||
is_service_account = _looks_like_service_account(auth_config)
|
||||
|
||||
client_kwargs = build_proxy_client_kwargs(ctx.proxy_config, timeout=timeout_seconds)
|
||||
|
||||
async with httpx.AsyncClient(**client_kwargs) as client:
|
||||
if is_service_account:
|
||||
models, errors, has_success = await _fetch_models_vertex_service_account(
|
||||
client,
|
||||
ctx=ctx,
|
||||
auth_config=auth_config,
|
||||
client_kwargs=client_kwargs,
|
||||
)
|
||||
else:
|
||||
models, errors, has_success = await _fetch_models_vertex_api_key(
|
||||
client,
|
||||
ctx=ctx,
|
||||
auth_config=auth_config,
|
||||
)
|
||||
|
||||
if not models and errors:
|
||||
logger.warning("Vertex 模型获取失败: {}", "; ".join(errors))
|
||||
return models, errors, has_success, None
|
||||
|
||||
|
||||
def register_all() -> None:
|
||||
"""一次性注册 Vertex AI 的所有 hooks 到各通用 registry。"""
|
||||
from src.core.api_format.capabilities import register_provider_behavior_variant
|
||||
from src.services.model.upstream_fetcher import UpstreamModelsFetcherRegistry
|
||||
from src.services.provider.adapters.vertex_ai.transport import build_vertex_ai_url
|
||||
from src.services.provider.transport import register_transport_hook
|
||||
|
||||
# Transport: Vertex AI 同时支持 gemini:chat 和 claude:chat 格式
|
||||
register_transport_hook("vertex_ai", "gemini:chat", build_vertex_ai_url)
|
||||
register_transport_hook("vertex_ai", "claude:chat", build_vertex_ai_url)
|
||||
|
||||
# Model Fetcher: Vertex 走专用模型获取链路
|
||||
UpstreamModelsFetcherRegistry.register(
|
||||
provider_types=["vertex_ai"],
|
||||
fetcher=fetch_models_vertex_ai,
|
||||
)
|
||||
|
||||
# Provider Format Capability:跨格式支持(同一 Vertex AI Provider 可同时访问 Gemini 和 Claude 模型)
|
||||
register_provider_behavior_variant("vertex_ai", cross_format=True)
|
||||
|
||||
|
||||
__all__ = ["fetch_models_vertex_ai", "register_all"]
|
||||
@@ -0,0 +1,291 @@
|
||||
"""Vertex AI URL 构建(Transport Hook)。
|
||||
|
||||
Vertex AI Gemini / Imagen 支持两种认证路径:
|
||||
|
||||
- API Key + Gemini/Imagen (Express mode):
|
||||
https://aiplatform.googleapis.com/v1/publishers/google/models/{model}:{action}?key={API_KEY}
|
||||
- Service Account + Gemini/Imagen:
|
||||
https://{region}-aiplatform.googleapis.com/v1/projects/{project_id}/locations/{region}/publishers/google/models/{model}:{action}
|
||||
|
||||
Claude 仍走标准 Vertex AI Service Account 路径:
|
||||
|
||||
- Service Account + Claude:
|
||||
https://{region}-aiplatform.googleapis.com/v1/projects/{project_id}/locations/{region}/publishers/anthropic/models/{model}:{action}
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
from urllib.parse import urlencode
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.core.provider_types import ProviderType, normalize_provider_type
|
||||
from src.services.provider.adapters.vertex_ai.constants import (
|
||||
API_KEY_BASE_URL,
|
||||
DEFAULT_FORMAT,
|
||||
DEFAULT_MODEL_REGIONS,
|
||||
MODEL_FORMAT_MAPPING,
|
||||
)
|
||||
from src.services.provider.format import normalize_endpoint_signature
|
||||
from src.services.provider.transport import looks_like_vertex_ai_host, redact_url_for_log
|
||||
|
||||
|
||||
def is_vertex_ai_context(
|
||||
*,
|
||||
base_url: str | None = None,
|
||||
provider_type: Any = None,
|
||||
endpoint: Any = None,
|
||||
key: Any = None,
|
||||
) -> bool:
|
||||
"""Best-effort 判断当前测试/请求上下文是否应视为 Vertex AI。"""
|
||||
if normalize_provider_type(provider_type) == ProviderType.VERTEX_AI.value:
|
||||
return True
|
||||
|
||||
for obj in (endpoint, key):
|
||||
provider = getattr(obj, "provider", None) if obj is not None else None
|
||||
if normalize_provider_type(getattr(provider, "provider_type", None)) == (
|
||||
ProviderType.VERTEX_AI.value
|
||||
):
|
||||
return True
|
||||
|
||||
candidate_base_url = str(base_url or getattr(endpoint, "base_url", "") or "").strip()
|
||||
if not candidate_base_url:
|
||||
return False
|
||||
|
||||
endpoint_sig = str(getattr(endpoint, "api_format", "") or "").strip()
|
||||
auth_type = str(getattr(key, "auth_type", "") or "").strip()
|
||||
return looks_like_vertex_ai_host(candidate_base_url, endpoint_sig, auth_type)
|
||||
|
||||
|
||||
def get_effective_format(
|
||||
model: str,
|
||||
auth_config: dict[str, Any] | None = None,
|
||||
) -> str:
|
||||
"""获取 Vertex AI 模式下模型的实际 API 格式。
|
||||
|
||||
优先级:
|
||||
1. auth_config.model_format_mapping 中的精确匹配
|
||||
2. auth_config.model_format_mapping 中的前缀匹配
|
||||
3. 内置 MODEL_FORMAT_MAPPING 前缀匹配
|
||||
4. auth_config.default_format
|
||||
5. 内置 DEFAULT_FORMAT
|
||||
"""
|
||||
user_format_mapping: dict[str, str] = {}
|
||||
user_default_format: str | None = None
|
||||
|
||||
if auth_config:
|
||||
user_format_mapping = auth_config.get("model_format_mapping", {})
|
||||
user_default_format = auth_config.get("default_format")
|
||||
|
||||
# 1. 用户配置:精确匹配
|
||||
if model in user_format_mapping:
|
||||
try:
|
||||
return normalize_endpoint_signature(user_format_mapping[model])
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Invalid vertex_ai model_format_mapping value for model '{}': {!r}",
|
||||
model,
|
||||
user_format_mapping[model],
|
||||
)
|
||||
|
||||
# 2. 用户配置:前缀匹配
|
||||
for prefix, api_format in user_format_mapping.items():
|
||||
if prefix.endswith("-") and model.startswith(prefix):
|
||||
try:
|
||||
return normalize_endpoint_signature(api_format)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Invalid vertex_ai model_format_mapping value for prefix '{}': {!r}",
|
||||
prefix,
|
||||
api_format,
|
||||
)
|
||||
break
|
||||
|
||||
# 3. 内置配置:前缀匹配
|
||||
for prefix, api_format in MODEL_FORMAT_MAPPING.items():
|
||||
if model.startswith(prefix):
|
||||
return normalize_endpoint_signature(api_format)
|
||||
|
||||
# 4. 用户默认格式
|
||||
if user_default_format:
|
||||
try:
|
||||
return normalize_endpoint_signature(user_default_format)
|
||||
except Exception:
|
||||
logger.warning("Invalid vertex_ai default_format: {!r}", user_default_format)
|
||||
|
||||
# 5. 内置默认格式
|
||||
return DEFAULT_FORMAT
|
||||
|
||||
|
||||
def build_vertex_ai_url(
|
||||
endpoint: Any,
|
||||
*,
|
||||
is_stream: bool,
|
||||
effective_query_params: dict[str, Any],
|
||||
path_params: dict[str, Any] | None = None,
|
||||
key: Any = None,
|
||||
decrypted_auth_config: dict[str, Any] | None = None,
|
||||
) -> str:
|
||||
"""Vertex AI transport hook — 统一 URL 构建入口。"""
|
||||
from src.core.exceptions import InvalidRequestException
|
||||
|
||||
model = str((path_params or {}).get("model", "") or "").strip()
|
||||
if not model:
|
||||
raise InvalidRequestException("Vertex AI 请求缺少 model 参数")
|
||||
|
||||
auth_type = str(getattr(key, "auth_type", "api_key") or "api_key").strip().lower()
|
||||
is_claude_model = model.startswith("claude-")
|
||||
|
||||
if auth_type == "api_key":
|
||||
if is_claude_model:
|
||||
raise InvalidRequestException(
|
||||
"Vertex API Key 不支持 Claude 模型,请改用 Service Account 认证。"
|
||||
)
|
||||
return _build_api_key_url(
|
||||
key=key,
|
||||
path_params=path_params,
|
||||
query_params=effective_query_params,
|
||||
is_stream=is_stream,
|
||||
)
|
||||
|
||||
# service_account(以及向后兼容旧的 "vertex_ai" auth_type)
|
||||
return _build_service_account_url(
|
||||
key=key,
|
||||
path_params=path_params,
|
||||
query_params=effective_query_params,
|
||||
is_stream=is_stream,
|
||||
decrypted_auth_config=decrypted_auth_config,
|
||||
)
|
||||
|
||||
|
||||
def _build_api_key_url(
|
||||
key: Any,
|
||||
*,
|
||||
path_params: dict[str, Any] | None = None,
|
||||
query_params: dict[str, Any] | None = None,
|
||||
is_stream: bool = False,
|
||||
) -> str:
|
||||
"""构建 API Key 认证的全局端点 URL。
|
||||
|
||||
格式: https://aiplatform.googleapis.com/v1/publishers/google/models/{model}:{action}?key={API_KEY}
|
||||
"""
|
||||
from src.core.crypto import crypto_service
|
||||
from src.core.exceptions import InvalidRequestException
|
||||
|
||||
model = (path_params or {}).get("model", "")
|
||||
if not model:
|
||||
raise InvalidRequestException("Vertex AI 请求缺少 model 参数")
|
||||
|
||||
action = "streamGenerateContent" if is_stream else "generateContent"
|
||||
path = f"/v1/publishers/google/models/{model}:{action}"
|
||||
url = f"{API_KEY_BASE_URL}{path}"
|
||||
|
||||
# 构建查询参数
|
||||
params = dict(query_params) if query_params else {}
|
||||
|
||||
# 附加 API Key
|
||||
api_key_value = crypto_service.decrypt(key.api_key) if key else ""
|
||||
if api_key_value:
|
||||
params["key"] = api_key_value
|
||||
|
||||
# Gemini 流式请求使用 SSE
|
||||
if is_stream:
|
||||
params.setdefault("alt", "sse")
|
||||
|
||||
params.pop("beta", None)
|
||||
|
||||
if params:
|
||||
query_string = urlencode(params, doseq=True)
|
||||
if query_string:
|
||||
url = f"{url}?{query_string}"
|
||||
|
||||
logger.debug("Vertex AI (API Key) URL: {}", redact_url_for_log(url))
|
||||
return url
|
||||
|
||||
|
||||
def _build_service_account_url(
|
||||
key: Any,
|
||||
*,
|
||||
path_params: dict[str, Any] | None = None,
|
||||
query_params: dict[str, Any] | None = None,
|
||||
is_stream: bool = False,
|
||||
decrypted_auth_config: dict[str, Any] | None = None,
|
||||
) -> str:
|
||||
"""构建 Service Account 认证的 Vertex AI 区域端点 URL。"""
|
||||
from src.core.crypto import crypto_service
|
||||
from src.core.exceptions import InvalidRequestException
|
||||
|
||||
# 优先使用传入的已解密配置,避免重复解密
|
||||
auth_config: dict[str, Any] = {}
|
||||
if decrypted_auth_config:
|
||||
auth_config = decrypted_auth_config
|
||||
else:
|
||||
# 兜底:从 key.auth_config 解密(理论上不应走到这里)
|
||||
raw_auth_config = getattr(key, "auth_config", None) if key else None
|
||||
if raw_auth_config:
|
||||
try:
|
||||
if isinstance(raw_auth_config, dict):
|
||||
auth_config = raw_auth_config
|
||||
else:
|
||||
decrypted_config = crypto_service.decrypt(raw_auth_config)
|
||||
auth_config = json.loads(decrypted_config)
|
||||
except Exception as e:
|
||||
logger.error("解密 Vertex AI auth_config 失败: {}", e)
|
||||
auth_config = {}
|
||||
|
||||
# 获取必需的配置
|
||||
project_id = auth_config.get("project_id")
|
||||
if not project_id:
|
||||
raise InvalidRequestException(
|
||||
"Vertex AI 配置缺少 project_id(请在 Key 的 auth_config 中提供)"
|
||||
)
|
||||
|
||||
# 获取模型名
|
||||
model = (path_params or {}).get("model", "")
|
||||
if not model:
|
||||
raise InvalidRequestException("Vertex AI 请求缺少 model 参数")
|
||||
|
||||
# 确定 region(优先级:用户配置 > 内置默认 > 用户默认 > 兜底)
|
||||
user_model_regions = auth_config.get("model_regions", {})
|
||||
user_default_region = auth_config.get("region")
|
||||
|
||||
if model in user_model_regions:
|
||||
region = user_model_regions[model]
|
||||
elif model in DEFAULT_MODEL_REGIONS:
|
||||
region = DEFAULT_MODEL_REGIONS[model]
|
||||
elif user_default_region:
|
||||
region = user_default_region
|
||||
else:
|
||||
region = "global"
|
||||
|
||||
if model.startswith("claude-"):
|
||||
publisher = "anthropic"
|
||||
action = "streamRawPredict" if is_stream else "rawPredict"
|
||||
else:
|
||||
publisher = "google"
|
||||
action = "streamGenerateContent" if is_stream else "generateContent"
|
||||
|
||||
# 构建 URL(global region 使用不同的 URL 格式)
|
||||
if region == "global":
|
||||
base_url = "https://aiplatform.googleapis.com"
|
||||
else:
|
||||
base_url = f"https://{region}-aiplatform.googleapis.com"
|
||||
path = f"/v1/projects/{project_id}/locations/{region}/publishers/{publisher}/models/{model}:{action}"
|
||||
url = f"{base_url}{path}"
|
||||
|
||||
# 添加查询参数
|
||||
effective_query_params = dict(query_params) if query_params else {}
|
||||
# Gemini 流式请求使用 SSE 格式,Claude 不需要
|
||||
if is_stream and not model.startswith("claude-"):
|
||||
effective_query_params.setdefault("alt", "sse")
|
||||
# 移除不适用于 Vertex AI 的参数
|
||||
effective_query_params.pop("beta", None)
|
||||
|
||||
if effective_query_params:
|
||||
query_string = urlencode(effective_query_params, doseq=True)
|
||||
if query_string:
|
||||
url = f"{url}?{query_string}"
|
||||
|
||||
logger.debug("Vertex AI (SA) URL: {} (region={})", redact_url_for_log(url), region)
|
||||
return url
|
||||
575
_deprecated_py_src/services/provider/auth.py
Normal file
575
_deprecated_py_src/services/provider/auth.py
Normal file
@@ -0,0 +1,575 @@
|
||||
"""
|
||||
Provider 认证逻辑(OAuth / Service Account / Vertex AI)。
|
||||
|
||||
从 api/handlers/base/request_builder.py 迁移到 services 层,
|
||||
消除 services→api 的反向依赖。
|
||||
|
||||
注意:
|
||||
- AI request hot-path 逐步迁到 Rust 后,这里仍然保留 Python 侧的 OAuth
|
||||
refresh / invalidation 状态持久化 owner。
|
||||
- 不要把新的 decision/control 路径继续扩展到这个模块。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from sqlalchemy.orm import object_session
|
||||
|
||||
from src.clients.redis_client import get_redis_client
|
||||
from src.core.crypto import crypto_service
|
||||
from src.core.logger import logger
|
||||
from src.core.provider_auth_types import ProviderAuthInfo
|
||||
from src.core.provider_oauth_utils import enrich_auth_config, post_oauth_token
|
||||
from src.services.provider.provider_context import resolve_provider_proxy
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from src.models.database import ProviderAPIKey, ProviderEndpoint
|
||||
|
||||
|
||||
# ==============================================================================
|
||||
# OAuth Token Refresh helpers
|
||||
# ==============================================================================
|
||||
|
||||
|
||||
async def _acquire_refresh_lock(key_id: str) -> tuple[Any, bool]:
|
||||
"""尝试获取 OAuth refresh 分布式锁。
|
||||
|
||||
返回 ``(redis_client | None, got_lock)``。调用方在刷新完成后
|
||||
必须调用 :func:`_release_refresh_lock` 释放锁。
|
||||
"""
|
||||
redis = await get_redis_client(require_redis=False)
|
||||
lock_key = f"provider_oauth_refresh_lock:{key_id}"
|
||||
got_lock = False
|
||||
if redis is not None:
|
||||
try:
|
||||
got_lock = bool(await redis.set(lock_key, "1", ex=30, nx=True))
|
||||
except Exception:
|
||||
got_lock = False
|
||||
return redis, got_lock
|
||||
|
||||
|
||||
async def _release_refresh_lock(redis: Any, key_id: str) -> None:
|
||||
"""释放 OAuth refresh 分布式锁(best-effort)。"""
|
||||
if redis is not None:
|
||||
try:
|
||||
await redis.delete(f"provider_oauth_refresh_lock:{key_id}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _safe_object_session(key: Any) -> Any | None:
|
||||
try:
|
||||
return object_session(key)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _persist_detached_oauth_invalid_state(
|
||||
key: Any,
|
||||
*,
|
||||
invalid_at: Any,
|
||||
invalid_reason: str,
|
||||
) -> None:
|
||||
"""持久化 oauth_invalid 状态。
|
||||
|
||||
这是 Python 仍保留的 status owner 之一:即使请求执行热路径迁到 Rust,
|
||||
admin/manual repair 和 detached refresh 仍会依赖这条持久化路径。
|
||||
"""
|
||||
key.oauth_invalid_at = invalid_at
|
||||
key.oauth_invalid_reason = invalid_reason
|
||||
|
||||
sess = _safe_object_session(key)
|
||||
if sess is not None:
|
||||
sess.add(key)
|
||||
sess.commit()
|
||||
return
|
||||
|
||||
key_id = str(getattr(key, "id", "") or "").strip()
|
||||
if not key_id:
|
||||
raise ValueError("OAuth key missing id")
|
||||
|
||||
from src.database import create_session
|
||||
from src.models.database import ProviderAPIKey
|
||||
|
||||
with create_session() as db:
|
||||
row = db.query(ProviderAPIKey).filter(ProviderAPIKey.id == key_id).first()
|
||||
if row is None:
|
||||
raise ValueError(f"OAuth key not found: {key_id}")
|
||||
row.oauth_invalid_at = invalid_at
|
||||
row.oauth_invalid_reason = invalid_reason
|
||||
db.commit()
|
||||
|
||||
|
||||
def _persist_refreshed_token(
|
||||
key: Any,
|
||||
access_token: str,
|
||||
token_meta: dict[str, Any],
|
||||
) -> None:
|
||||
"""将刷新后的 access_token 和 auth_config 持久化到数据库。"""
|
||||
key.api_key = crypto_service.encrypt(access_token)
|
||||
key.auth_config = crypto_service.encrypt(json.dumps(token_meta))
|
||||
|
||||
# 刷新成功只清除可恢复的 token 类异常。
|
||||
# 账号级 block(如验证要求/工作区停用)不能靠 token refresh 自动恢复。
|
||||
from src.services.provider.oauth_token import is_account_level_block
|
||||
|
||||
current_reason = str(getattr(key, "oauth_invalid_reason", None) or "").strip()
|
||||
if getattr(key, "oauth_invalid_at", None) is not None and not is_account_level_block(
|
||||
current_reason
|
||||
):
|
||||
key.oauth_invalid_at = None
|
||||
key.oauth_invalid_reason = None
|
||||
|
||||
sess = _safe_object_session(key)
|
||||
if sess is not None:
|
||||
sess.add(key)
|
||||
sess.commit()
|
||||
else:
|
||||
logger.warning(
|
||||
"[OAUTH_REFRESH] key {} refreshed but cannot persist (no session); "
|
||||
"next request will refresh again",
|
||||
key.id,
|
||||
)
|
||||
|
||||
|
||||
def _extract_refresh_error_detail(error_body: str) -> str:
|
||||
"""Best-effort extraction of error detail from refresh token error response."""
|
||||
try:
|
||||
data = json.loads(error_body)
|
||||
if isinstance(data, dict):
|
||||
err = data.get("error")
|
||||
if isinstance(err, dict):
|
||||
code = err.get("code") or ""
|
||||
msg = err.get("message") or ""
|
||||
return f"{code}: {msg}".strip(": ") if (code or msg) else ""
|
||||
if isinstance(err, str):
|
||||
return err
|
||||
return str(data.get("error_description") or data.get("message") or "")
|
||||
except Exception:
|
||||
pass
|
||||
return error_body[:200] if error_body else ""
|
||||
|
||||
|
||||
def _mark_refresh_token_invalid(
|
||||
key: Any,
|
||||
status_code: int,
|
||||
error_body: str,
|
||||
) -> None:
|
||||
"""标记 refresh token 已失效(仅设置 oauth_invalid 标记,不停用 key)。
|
||||
|
||||
Access token 在过期前仍可正常使用。oauth_invalid_reason 使用 [REFRESH_FAILED]
|
||||
前缀。注意:如果上游错误体中包含账号封禁关键词(如 "deactivated"),
|
||||
该 reason 仍会被 account_state 的关键词匹配判定为 blocked,这是预期行为。
|
||||
"""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
detail = _extract_refresh_error_detail(error_body)
|
||||
reason = f"[REFRESH_FAILED] Token 续期失败 ({status_code})"
|
||||
if detail:
|
||||
reason = f"{reason}: {detail}"
|
||||
|
||||
try:
|
||||
_persist_detached_oauth_invalid_state(
|
||||
key,
|
||||
invalid_at=datetime.now(timezone.utc),
|
||||
invalid_reason=reason,
|
||||
)
|
||||
logger.info(
|
||||
"[OAUTH_REFRESH] key {} marked refresh_token invalid: {}",
|
||||
str(getattr(key, "id", "?"))[:8],
|
||||
reason[:120],
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"[OAUTH_REFRESH] failed to mark key {} refresh invalid: {}",
|
||||
str(getattr(key, "id", "?"))[:8],
|
||||
str(exc),
|
||||
)
|
||||
|
||||
|
||||
def _mark_oauth_token_expired(key: Any, expires_at: Any) -> None:
|
||||
"""标记 OAuth key 为 Token 已过期且无法续期,阻止后续调度。
|
||||
|
||||
当 refresh token 已失效且 access token 也已过期时调用。
|
||||
使用 [OAUTH_EXPIRED] 前缀,account_state 会将其判定为 blocked。
|
||||
不设置 is_active = False(管理员可通过重新导入凭据恢复)。
|
||||
"""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
# 如果已经有更严重的标记([ACCOUNT_BLOCK]),不降级
|
||||
existing = str(getattr(key, "oauth_invalid_reason", None) or "")
|
||||
if existing.startswith("[ACCOUNT_BLOCK]"):
|
||||
return
|
||||
|
||||
reason = f"[OAUTH_EXPIRED] Token 已过期且续期失败 (expired_at={expires_at})"
|
||||
|
||||
try:
|
||||
_persist_detached_oauth_invalid_state(
|
||||
key,
|
||||
invalid_at=datetime.now(timezone.utc),
|
||||
invalid_reason=reason,
|
||||
)
|
||||
logger.info(
|
||||
"[OAUTH_EXPIRED] key {} token expired and refresh failed, blocking scheduling",
|
||||
str(getattr(key, "id", "?"))[:8],
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"[OAUTH_EXPIRED] failed to mark key {} as expired: {}",
|
||||
str(getattr(key, "id", "?"))[:8],
|
||||
str(exc),
|
||||
)
|
||||
|
||||
|
||||
def _get_proxy_config(key: Any, endpoint: Any = None) -> Any:
|
||||
"""获取有效代理配置(Key 级别优先于 Provider 级别)。"""
|
||||
try:
|
||||
from src.services.proxy_node.resolver import resolve_effective_proxy
|
||||
|
||||
provider_proxy = resolve_provider_proxy(endpoint=endpoint, key=key)
|
||||
key_proxy = getattr(key, "proxy", None)
|
||||
return resolve_effective_proxy(provider_proxy, key_proxy)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
# ==============================================================================
|
||||
# Provider-specific refresh implementations
|
||||
# ==============================================================================
|
||||
|
||||
|
||||
async def _refresh_kiro_token(
|
||||
key: Any,
|
||||
endpoint: Any,
|
||||
token_meta: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
"""Kiro OAuth refresh: validate + call Kiro-specific refresh endpoint."""
|
||||
from src.core.exceptions import InvalidRequestException
|
||||
from src.services.provider.adapters.kiro.models.credentials import KiroAuthConfig
|
||||
from src.services.provider.adapters.kiro.token_manager import (
|
||||
refresh_access_token,
|
||||
validate_refresh_token,
|
||||
)
|
||||
|
||||
cfg = KiroAuthConfig.from_dict(token_meta or {})
|
||||
if not (cfg.refresh_token or "").strip():
|
||||
raise InvalidRequestException(
|
||||
"Kiro auth_config missing refresh_token; please re-import credentials."
|
||||
)
|
||||
|
||||
proxy_config = _get_proxy_config(key, endpoint)
|
||||
|
||||
validate_refresh_token(cfg.refresh_token)
|
||||
access_token, new_cfg = await refresh_access_token(
|
||||
cfg,
|
||||
proxy_config=proxy_config,
|
||||
)
|
||||
new_meta = new_cfg.to_dict()
|
||||
new_meta["updated_at"] = int(time.time())
|
||||
|
||||
_persist_refreshed_token(key, access_token, new_meta)
|
||||
return new_meta
|
||||
|
||||
|
||||
async def _refresh_generic_oauth_token(
|
||||
key: Any,
|
||||
endpoint: Any,
|
||||
template: Any,
|
||||
provider_type: str,
|
||||
refresh_token: str,
|
||||
token_meta: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
"""Generic OAuth refresh via template (Codex, Antigravity, ClaudeCode, etc.)."""
|
||||
token_url = template.oauth.token_url
|
||||
is_json = "anthropic.com" in token_url
|
||||
|
||||
scopes = getattr(template.oauth, "scopes", None) or []
|
||||
scope_str = " ".join(scopes) if scopes else ""
|
||||
|
||||
if is_json:
|
||||
body: dict[str, Any] = {
|
||||
"grant_type": "refresh_token",
|
||||
"client_id": template.oauth.client_id,
|
||||
"refresh_token": str(refresh_token),
|
||||
}
|
||||
if scope_str:
|
||||
body["scope"] = scope_str
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Accept": "application/json",
|
||||
}
|
||||
data = None
|
||||
json_body = body
|
||||
else:
|
||||
form: dict[str, str] = {
|
||||
"grant_type": "refresh_token",
|
||||
"client_id": template.oauth.client_id,
|
||||
"refresh_token": str(refresh_token),
|
||||
}
|
||||
if scope_str:
|
||||
form["scope"] = scope_str
|
||||
if template.oauth.client_secret:
|
||||
form["client_secret"] = template.oauth.client_secret
|
||||
headers = {
|
||||
"Content-Type": "application/x-www-form-urlencoded",
|
||||
"Accept": "application/json",
|
||||
}
|
||||
data = form
|
||||
json_body = None
|
||||
|
||||
proxy_config = _get_proxy_config(key, endpoint)
|
||||
|
||||
resp = await post_oauth_token(
|
||||
provider_type=provider_type,
|
||||
token_url=token_url,
|
||||
headers=headers,
|
||||
data=data,
|
||||
json_body=json_body,
|
||||
proxy_config=proxy_config,
|
||||
timeout_seconds=30.0,
|
||||
)
|
||||
|
||||
if 200 <= resp.status_code < 300:
|
||||
token = resp.json()
|
||||
access_token = str(token.get("access_token") or "")
|
||||
new_refresh_token = str(token.get("refresh_token") or "")
|
||||
expires_in = token.get("expires_in")
|
||||
new_expires_at: int | None = None
|
||||
try:
|
||||
if expires_in is not None:
|
||||
new_expires_at = int(time.time()) + int(expires_in)
|
||||
except Exception:
|
||||
new_expires_at = None
|
||||
|
||||
if access_token:
|
||||
token_meta["token_type"] = token.get("token_type")
|
||||
if new_refresh_token:
|
||||
token_meta["refresh_token"] = new_refresh_token
|
||||
token_meta["expires_at"] = new_expires_at
|
||||
token_meta["scope"] = token.get("scope")
|
||||
token_meta["updated_at"] = int(time.time())
|
||||
|
||||
token_meta = await enrich_auth_config(
|
||||
provider_type=provider_type,
|
||||
auth_config=token_meta,
|
||||
token_response=token,
|
||||
access_token=access_token,
|
||||
proxy_config=proxy_config,
|
||||
)
|
||||
|
||||
_persist_refreshed_token(key, access_token, token_meta)
|
||||
else:
|
||||
error_body = ""
|
||||
try:
|
||||
error_body = resp.text or ""
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
logger.warning(
|
||||
"OAuth token refresh failed: provider={}, key_id={}, status={}, body={}",
|
||||
provider_type,
|
||||
getattr(key, "id", "?"),
|
||||
resp.status_code,
|
||||
error_body[:500],
|
||||
)
|
||||
|
||||
# 标记 refresh token 失效(不停用 key,access token 过期前仍可调度)。
|
||||
# 注意:如果上游错误包含账号封禁关键词(如 "deactivated"),
|
||||
# oauth_invalid_reason 会被 account_state 关键词匹配判定为 blocked,这是预期行为。
|
||||
_mark_refresh_token_invalid(key, resp.status_code, error_body)
|
||||
|
||||
return token_meta
|
||||
|
||||
|
||||
# ==============================================================================
|
||||
# Service Account 认证支持
|
||||
# ==============================================================================
|
||||
|
||||
|
||||
async def get_provider_auth(
|
||||
endpoint: "ProviderEndpoint",
|
||||
key: "ProviderAPIKey",
|
||||
*,
|
||||
force_refresh: bool = False,
|
||||
refresh_skew: int | None = None,
|
||||
) -> ProviderAuthInfo | None:
|
||||
"""
|
||||
获取 Provider 的认证信息
|
||||
|
||||
对于标准 API Key,返回 None(由 build_headers 自动处理)。
|
||||
对于 Service Account,异步获取 Access Token 并返回认证信息。
|
||||
|
||||
Args:
|
||||
endpoint: 端点配置
|
||||
key: Provider API Key
|
||||
|
||||
Returns:
|
||||
Service Account 场景: ProviderAuthInfo 对象(包含认证信息和解密后的配置)
|
||||
API Key 场景: None(由 build_headers 处理)
|
||||
|
||||
Raises:
|
||||
InvalidRequestException: 认证配置无效或认证失败
|
||||
"""
|
||||
from src.core.exceptions import InvalidRequestException
|
||||
|
||||
auth_type = getattr(key, "auth_type", "api_key")
|
||||
|
||||
if auth_type == "oauth":
|
||||
# OAuth token 保存在 key.api_key(加密),refresh_token/expires_at 等在 auth_config(加密 JSON)中。
|
||||
# 在请求前做一次懒刷新:接近过期时刷新 access_token,并用 Redis lock 避免并发风暴。
|
||||
|
||||
encrypted_auth_config = getattr(key, "auth_config", None)
|
||||
|
||||
# 先解密 auth_config -- 下游 build_provider_url 等依赖 decrypted_auth_config
|
||||
# 中的 provider_type / project_id / region 等元数据,即使 access_token 命中缓存
|
||||
# 也不能跳过。
|
||||
if encrypted_auth_config:
|
||||
try:
|
||||
decrypted_config = crypto_service.decrypt(encrypted_auth_config)
|
||||
token_meta = json.loads(decrypted_config)
|
||||
except Exception:
|
||||
token_meta = {}
|
||||
else:
|
||||
token_meta = {}
|
||||
|
||||
decrypted_auth_config: dict[str, Any] | None = (
|
||||
token_meta if isinstance(token_meta, dict) and token_meta else None
|
||||
)
|
||||
|
||||
# 快路径:查 Redis token 缓存,命中则跳过 refresh 和 api_key 解密。
|
||||
# 注意:token_meta/decrypted_auth_config 已在上方解密,此处只是跳过后续刷新逻辑。
|
||||
if not force_refresh and encrypted_auth_config:
|
||||
try:
|
||||
from src.services.provider.pool.oauth_cache import get_cached_token
|
||||
|
||||
_cached = await get_cached_token(str(key.id))
|
||||
if _cached:
|
||||
return ProviderAuthInfo(
|
||||
auth_header="Authorization",
|
||||
auth_value=f"Bearer {_cached}",
|
||||
decrypted_auth_config=decrypted_auth_config,
|
||||
)
|
||||
except Exception:
|
||||
logger.debug("OAuth token cache lookup failed for key {}", str(key.id)[:8])
|
||||
|
||||
expires_at = token_meta.get("expires_at")
|
||||
refresh_token = token_meta.get("refresh_token")
|
||||
provider_type = str(token_meta.get("provider_type") or "")
|
||||
cached_access_token = str(token_meta.get("access_token") or "").strip()
|
||||
|
||||
# Refresh skew: providers with pool config use configurable
|
||||
# proactive_refresh_seconds (default 180 s), others use 120 s.
|
||||
# Prefer the caller-supplied value to avoid ORM lazy-load on key.provider.
|
||||
_refresh_skew = refresh_skew if refresh_skew is not None else 120
|
||||
if refresh_skew is None:
|
||||
try:
|
||||
from src.services.provider.pool.config import parse_pool_config
|
||||
|
||||
provider_obj = getattr(key, "provider", None)
|
||||
pcfg = getattr(provider_obj, "config", None) if provider_obj else None
|
||||
pool_cfg = parse_pool_config(pcfg) if pcfg else None
|
||||
if pool_cfg is not None:
|
||||
_refresh_skew = pool_cfg.proactive_refresh_seconds
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
should_refresh = False
|
||||
try:
|
||||
if expires_at is not None:
|
||||
should_refresh = int(time.time()) >= int(expires_at) - _refresh_skew
|
||||
except Exception:
|
||||
should_refresh = False
|
||||
|
||||
if force_refresh:
|
||||
should_refresh = True
|
||||
|
||||
# Kiro 特殊处理:如果没有缓存的 access_token 或 key.api_key 是占位符,强制刷新
|
||||
if provider_type == "kiro" and not should_refresh:
|
||||
if not cached_access_token:
|
||||
should_refresh = True
|
||||
elif crypto_service.decrypt(key.api_key) == "__placeholder__":
|
||||
should_refresh = True
|
||||
|
||||
_refreshed = False
|
||||
_lost_lock = False # 其他实例持有刷新锁,不应标记过期
|
||||
if should_refresh and refresh_token and provider_type:
|
||||
try:
|
||||
from src.core.provider_templates.fixed_providers import FIXED_PROVIDERS
|
||||
from src.core.provider_templates.types import ProviderType
|
||||
|
||||
try:
|
||||
template = FIXED_PROVIDERS.get(ProviderType(provider_type))
|
||||
except Exception:
|
||||
template = None
|
||||
|
||||
redis, got_lock = await _acquire_refresh_lock(key.id)
|
||||
if got_lock or redis is None:
|
||||
try:
|
||||
if provider_type == ProviderType.KIRO.value:
|
||||
token_meta = await _refresh_kiro_token(key, endpoint, token_meta)
|
||||
elif template:
|
||||
token_meta = await _refresh_generic_oauth_token(
|
||||
key, endpoint, template, provider_type, refresh_token, token_meta
|
||||
)
|
||||
_refreshed = True
|
||||
finally:
|
||||
if got_lock:
|
||||
await _release_refresh_lock(redis, key.id)
|
||||
else:
|
||||
_lost_lock = True
|
||||
except Exception:
|
||||
# 刷新失败不阻断请求;后续由上游返回 401 再触发管理端处理
|
||||
pass
|
||||
|
||||
# Refresh 失败(非锁竞争)且 access token 已过期 → 升级标记为 [OAUTH_EXPIRED]
|
||||
# 注意:未获取到锁说明其他实例正在刷新,不应在此标记为过期
|
||||
if should_refresh and not _refreshed and not _lost_lock and expires_at is not None:
|
||||
try:
|
||||
token_truly_expired = int(time.time()) >= int(expires_at)
|
||||
except Exception:
|
||||
token_truly_expired = False
|
||||
if token_truly_expired:
|
||||
_mark_oauth_token_expired(key, expires_at)
|
||||
|
||||
# 获取最终使用的 access_token
|
||||
# Kiro 优先使用 token_meta 中缓存的 access_token(刷新后会更新到 token_meta)
|
||||
if provider_type == "kiro":
|
||||
refreshed_token = str(token_meta.get("access_token") or "").strip()
|
||||
effective_token = refreshed_token or crypto_service.decrypt(key.api_key)
|
||||
else:
|
||||
effective_token = crypto_service.decrypt(key.api_key)
|
||||
|
||||
# 刷新成功后写入 Redis token 缓存(所有 OAuth key 均可受益)
|
||||
if _refreshed and effective_token:
|
||||
try:
|
||||
from src.services.provider.pool.oauth_cache import cache_token
|
||||
|
||||
new_expires_at = token_meta.get("expires_at")
|
||||
if new_expires_at is not None:
|
||||
remaining = int(new_expires_at) - int(time.time())
|
||||
if remaining > 0:
|
||||
await cache_token(str(key.id), effective_token, remaining)
|
||||
except Exception:
|
||||
logger.debug("OAuth token cache write failed for key {}", str(key.id)[:8])
|
||||
|
||||
# 刷新可能更新了 token_meta,同步 decrypted_auth_config
|
||||
if isinstance(token_meta, dict) and token_meta:
|
||||
decrypted_auth_config = token_meta
|
||||
|
||||
return ProviderAuthInfo(
|
||||
auth_header="Authorization",
|
||||
auth_value=f"Bearer {effective_token}",
|
||||
decrypted_auth_config=decrypted_auth_config,
|
||||
)
|
||||
if auth_type in ("service_account", "vertex_ai"):
|
||||
# service_account: GCP Service Account JSON → JWT → Access Token
|
||||
# "vertex_ai" 保留为向后兼容(迁移期间旧数据可能仍使用该值)
|
||||
from src.services.provider.adapters.vertex_ai.auth import _auth_service_account
|
||||
|
||||
return await _auth_service_account(key, endpoint)
|
||||
|
||||
# 标准 API Key:返回 None,由 build_headers 处理
|
||||
return None
|
||||
60
_deprecated_py_src/services/provider/behavior.py
Normal file
60
_deprecated_py_src/services/provider/behavior.py
Normal file
@@ -0,0 +1,60 @@
|
||||
"""Provider behavior 薄封装。
|
||||
|
||||
对外保持既有调用接口,内部统一委托给 core.api_format.capabilities 中的 provider registry。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from src.core.api_format.capabilities import (
|
||||
get_provider_behavior_variants,
|
||||
register_provider_behavior_variant,
|
||||
)
|
||||
from src.core.provider_types import normalize_provider_type
|
||||
from src.services.provider.envelope import ProviderEnvelope, get_provider_envelope
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ProviderBehavior:
|
||||
provider_type: str
|
||||
envelope: ProviderEnvelope | None
|
||||
same_format_variant: str | None
|
||||
cross_format_variant: str | None
|
||||
|
||||
|
||||
def register_behavior_variant(
|
||||
provider_type: str,
|
||||
*,
|
||||
same_format: bool = False,
|
||||
cross_format: bool = False,
|
||||
) -> None:
|
||||
"""兼容入口:注册 provider 的格式变体标志,真实存储位于 core registry。"""
|
||||
register_provider_behavior_variant(
|
||||
provider_type,
|
||||
same_format=same_format,
|
||||
cross_format=cross_format,
|
||||
)
|
||||
|
||||
|
||||
def get_provider_behavior(
|
||||
*,
|
||||
provider_type: str | None,
|
||||
endpoint_sig: str | None,
|
||||
) -> ProviderBehavior:
|
||||
pt = normalize_provider_type(provider_type)
|
||||
envelope = get_provider_envelope(provider_type=pt, endpoint_sig=endpoint_sig)
|
||||
same_format_variant, cross_format_variant = get_provider_behavior_variants(
|
||||
provider_type=pt,
|
||||
endpoint_sig=endpoint_sig or "",
|
||||
)
|
||||
|
||||
return ProviderBehavior(
|
||||
provider_type=pt,
|
||||
envelope=envelope,
|
||||
same_format_variant=same_format_variant,
|
||||
cross_format_variant=cross_format_variant,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["ProviderBehavior", "get_provider_behavior", "register_behavior_variant"]
|
||||
293
_deprecated_py_src/services/provider/delete_cleanup.py
Normal file
293
_deprecated_py_src/services/provider/delete_cleanup.py
Normal file
@@ -0,0 +1,293 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterable, Sequence
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.models.database import (
|
||||
ApiKey,
|
||||
Model,
|
||||
Provider,
|
||||
ProviderAPIKey,
|
||||
ProviderEndpoint,
|
||||
RequestCandidate,
|
||||
Usage,
|
||||
User,
|
||||
UserPreference,
|
||||
VideoTask,
|
||||
)
|
||||
from src.models.database_extensions import ApiKeyProviderMapping, ProviderUsageTracking
|
||||
from src.services.provider_keys.key_side_effects import cleanup_key_references
|
||||
|
||||
_BATCH_SIZE = 2000
|
||||
|
||||
|
||||
def _empty_cleanup_stats() -> dict[str, int]:
|
||||
return {
|
||||
"users": 0,
|
||||
"api_keys": 0,
|
||||
"user_preferences": 0,
|
||||
"usage_provider": 0,
|
||||
"usage_endpoint": 0,
|
||||
"video_tasks_provider": 0,
|
||||
"video_tasks_endpoint": 0,
|
||||
"request_candidates_provider": 0,
|
||||
"request_candidates_endpoint": 0,
|
||||
}
|
||||
|
||||
|
||||
def _empty_delete_stats() -> dict[str, int]:
|
||||
return {
|
||||
"api_key_mappings": 0,
|
||||
"usage_tracking": 0,
|
||||
"models": 0,
|
||||
"api_keys": 0,
|
||||
"endpoints": 0,
|
||||
"providers": 0,
|
||||
}
|
||||
|
||||
|
||||
def _iter_batches(items: Sequence[str], batch_size: int = _BATCH_SIZE) -> list[list[str]]:
|
||||
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)]
|
||||
|
||||
|
||||
def _collect_provider_child_ids(db: Session, provider_id: str) -> tuple[list[str], list[str]]:
|
||||
endpoint_ids = [
|
||||
endpoint_id
|
||||
for endpoint_id, in db.query(ProviderEndpoint.id)
|
||||
.filter(ProviderEndpoint.provider_id == provider_id)
|
||||
.all()
|
||||
]
|
||||
key_ids = [
|
||||
key_id
|
||||
for key_id, in db.query(ProviderAPIKey.id)
|
||||
.filter(ProviderAPIKey.provider_id == provider_id)
|
||||
.all()
|
||||
]
|
||||
return endpoint_ids, key_ids
|
||||
|
||||
|
||||
def prune_allowed_provider_list(
|
||||
allowed_providers: Any, provider_id: str
|
||||
) -> tuple[list[str] | None | Any, bool]:
|
||||
"""从访问限制列表中移除指定 Provider ID。"""
|
||||
if not isinstance(allowed_providers, list):
|
||||
return allowed_providers, False
|
||||
if provider_id not in allowed_providers:
|
||||
return allowed_providers, False
|
||||
|
||||
next_allowed = [value for value in allowed_providers if value != provider_id]
|
||||
return next_allowed, True
|
||||
|
||||
|
||||
def prune_allowed_provider_refs(records: Iterable[Any], provider_id: str) -> int:
|
||||
"""批量移除记录中的 allowed_providers 引用。"""
|
||||
updated = 0
|
||||
for record in records:
|
||||
next_allowed, changed = prune_allowed_provider_list(
|
||||
getattr(record, "allowed_providers", None),
|
||||
provider_id,
|
||||
)
|
||||
if not changed:
|
||||
continue
|
||||
record.allowed_providers = next_allowed
|
||||
updated += 1
|
||||
return updated
|
||||
|
||||
|
||||
def cleanup_deleted_provider_references(
|
||||
db: Session,
|
||||
provider_id: str,
|
||||
*,
|
||||
endpoint_ids: Sequence[str] | None = None,
|
||||
key_ids: Sequence[str] | None = None,
|
||||
) -> dict[str, int]:
|
||||
"""清理 Provider 删除时的大扇出引用,避免依赖数据库级联导致慢删。"""
|
||||
if not provider_id:
|
||||
return _empty_cleanup_stats()
|
||||
|
||||
if endpoint_ids is None or key_ids is None:
|
||||
resolved_endpoint_ids, resolved_key_ids = _collect_provider_child_ids(db, provider_id)
|
||||
endpoint_ids = resolved_endpoint_ids if endpoint_ids is None else list(endpoint_ids)
|
||||
key_ids = resolved_key_ids if key_ids is None else list(key_ids)
|
||||
else:
|
||||
endpoint_ids = list(endpoint_ids)
|
||||
key_ids = list(key_ids)
|
||||
|
||||
updated_users = prune_allowed_provider_refs(
|
||||
db.query(User).filter(User.allowed_providers.isnot(None)).all(),
|
||||
provider_id,
|
||||
)
|
||||
updated_api_keys = prune_allowed_provider_refs(
|
||||
db.query(ApiKey).filter(ApiKey.allowed_providers.isnot(None)).all(),
|
||||
provider_id,
|
||||
)
|
||||
|
||||
cleared_preferences = int(
|
||||
db.query(UserPreference)
|
||||
.filter(UserPreference.default_provider_id == provider_id)
|
||||
.update({UserPreference.default_provider_id: None}, synchronize_session=False)
|
||||
or 0
|
||||
)
|
||||
cleared_usage_providers = int(
|
||||
db.query(Usage)
|
||||
.filter(Usage.provider_id == provider_id)
|
||||
.update({Usage.provider_id: None}, synchronize_session=False)
|
||||
or 0
|
||||
)
|
||||
cleared_video_task_providers = int(
|
||||
db.query(VideoTask)
|
||||
.filter(VideoTask.provider_id == provider_id)
|
||||
.update({VideoTask.provider_id: None}, synchronize_session=False)
|
||||
or 0
|
||||
)
|
||||
|
||||
if key_ids:
|
||||
cleanup_key_references(db, list(key_ids))
|
||||
|
||||
cleared_usage_endpoints = 0
|
||||
cleared_video_task_endpoints = 0
|
||||
deleted_request_candidates_endpoints = 0
|
||||
for batch in _iter_batches(endpoint_ids):
|
||||
cleared_usage_endpoints += int(
|
||||
db.query(Usage)
|
||||
.filter(Usage.provider_endpoint_id.in_(batch))
|
||||
.update({Usage.provider_endpoint_id: None}, synchronize_session=False)
|
||||
or 0
|
||||
)
|
||||
cleared_video_task_endpoints += int(
|
||||
db.query(VideoTask)
|
||||
.filter(VideoTask.endpoint_id.in_(batch))
|
||||
.update({VideoTask.endpoint_id: None}, synchronize_session=False)
|
||||
or 0
|
||||
)
|
||||
deleted_request_candidates_endpoints += int(
|
||||
db.query(RequestCandidate)
|
||||
.filter(RequestCandidate.endpoint_id.in_(batch))
|
||||
.delete(synchronize_session=False)
|
||||
or 0
|
||||
)
|
||||
|
||||
deleted_request_candidates_provider = int(
|
||||
db.query(RequestCandidate)
|
||||
.filter(RequestCandidate.provider_id == provider_id)
|
||||
.delete(synchronize_session=False)
|
||||
or 0
|
||||
)
|
||||
|
||||
stats = {
|
||||
"users": updated_users,
|
||||
"api_keys": updated_api_keys,
|
||||
"user_preferences": cleared_preferences,
|
||||
"usage_provider": cleared_usage_providers,
|
||||
"usage_endpoint": cleared_usage_endpoints,
|
||||
"video_tasks_provider": cleared_video_task_providers,
|
||||
"video_tasks_endpoint": cleared_video_task_endpoints,
|
||||
"request_candidates_provider": deleted_request_candidates_provider,
|
||||
"request_candidates_endpoint": deleted_request_candidates_endpoints,
|
||||
}
|
||||
|
||||
if any(stats.values()) or key_ids:
|
||||
logger.info(
|
||||
"Provider 删除引用清理: provider_id={}, key_refs={}, users={}, api_keys={}, "
|
||||
"user_preferences={}, usage_provider={}, usage_endpoint={}, "
|
||||
"video_tasks_provider={}, video_tasks_endpoint={}, "
|
||||
"request_candidates_provider={}, request_candidates_endpoint={}",
|
||||
provider_id,
|
||||
len(key_ids),
|
||||
stats["users"],
|
||||
stats["api_keys"],
|
||||
stats["user_preferences"],
|
||||
stats["usage_provider"],
|
||||
stats["usage_endpoint"],
|
||||
stats["video_tasks_provider"],
|
||||
stats["video_tasks_endpoint"],
|
||||
stats["request_candidates_provider"],
|
||||
stats["request_candidates_endpoint"],
|
||||
)
|
||||
|
||||
return stats
|
||||
|
||||
|
||||
def delete_provider_tree(db: Session, provider_id: str) -> dict[str, Any]:
|
||||
"""分阶段删除 Provider 及其子资源,降低 ORM/FK 级联导致的超时风险。"""
|
||||
if not provider_id:
|
||||
return {
|
||||
"cleanup": _empty_cleanup_stats(),
|
||||
"deleted": _empty_delete_stats(),
|
||||
"key_count": 0,
|
||||
"endpoint_count": 0,
|
||||
}
|
||||
|
||||
endpoint_ids, key_ids = _collect_provider_child_ids(db, provider_id)
|
||||
cleanup_stats = cleanup_deleted_provider_references(
|
||||
db,
|
||||
provider_id,
|
||||
endpoint_ids=endpoint_ids,
|
||||
key_ids=key_ids,
|
||||
)
|
||||
|
||||
deleted_stats = {
|
||||
"api_key_mappings": int(
|
||||
db.query(ApiKeyProviderMapping)
|
||||
.filter(ApiKeyProviderMapping.provider_id == provider_id)
|
||||
.delete(synchronize_session=False)
|
||||
or 0
|
||||
),
|
||||
"usage_tracking": int(
|
||||
db.query(ProviderUsageTracking)
|
||||
.filter(ProviderUsageTracking.provider_id == provider_id)
|
||||
.delete(synchronize_session=False)
|
||||
or 0
|
||||
),
|
||||
"models": int(
|
||||
db.query(Model)
|
||||
.filter(Model.provider_id == provider_id)
|
||||
.delete(synchronize_session=False)
|
||||
or 0
|
||||
),
|
||||
"api_keys": int(
|
||||
db.query(ProviderAPIKey)
|
||||
.filter(ProviderAPIKey.provider_id == provider_id)
|
||||
.delete(synchronize_session=False)
|
||||
or 0
|
||||
),
|
||||
"endpoints": int(
|
||||
db.query(ProviderEndpoint)
|
||||
.filter(ProviderEndpoint.provider_id == provider_id)
|
||||
.delete(synchronize_session=False)
|
||||
or 0
|
||||
),
|
||||
"providers": int(
|
||||
db.query(Provider).filter(Provider.id == provider_id).delete(synchronize_session=False)
|
||||
or 0
|
||||
),
|
||||
}
|
||||
|
||||
logger.info(
|
||||
"Provider 分阶段删除: provider_id={}, key_count={}, endpoint_count={}, "
|
||||
"deleted_mappings={}, deleted_usage_tracking={}, deleted_models={}, "
|
||||
"deleted_api_keys={}, deleted_endpoints={}, deleted_providers={}",
|
||||
provider_id,
|
||||
len(key_ids),
|
||||
len(endpoint_ids),
|
||||
deleted_stats["api_key_mappings"],
|
||||
deleted_stats["usage_tracking"],
|
||||
deleted_stats["models"],
|
||||
deleted_stats["api_keys"],
|
||||
deleted_stats["endpoints"],
|
||||
deleted_stats["providers"],
|
||||
)
|
||||
|
||||
return {
|
||||
"cleanup": cleanup_stats,
|
||||
"deleted": deleted_stats,
|
||||
"key_count": len(key_ids),
|
||||
"endpoint_count": len(endpoint_ids),
|
||||
}
|
||||
540
_deprecated_py_src/services/provider/delete_task.py
Normal file
540
_deprecated_py_src/services/provider/delete_task.py
Normal file
@@ -0,0 +1,540 @@
|
||||
"""Provider 异步删除任务。
|
||||
|
||||
接口提交后立即返回 task_id,后台分阶段删除 provider 及其子资源。
|
||||
任务状态存储在 Redis 中,支持多 worker 进程共享。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Callable, Sequence
|
||||
from concurrent.futures import Future
|
||||
from typing import Any
|
||||
|
||||
import redis.asyncio as aioredis
|
||||
from sqlalchemy import text
|
||||
|
||||
from src.clients.redis_client import get_redis_client
|
||||
from src.core.logger import logger
|
||||
from src.database import create_session
|
||||
from src.models.database import (
|
||||
ApiKey,
|
||||
Model,
|
||||
Provider,
|
||||
ProviderAPIKey,
|
||||
ProviderEndpoint,
|
||||
RequestCandidate,
|
||||
Usage,
|
||||
User,
|
||||
UserPreference,
|
||||
VideoTask,
|
||||
)
|
||||
from src.models.database_extensions import ApiKeyProviderMapping, ProviderUsageTracking
|
||||
from src.services.cache.model_cache import ModelCacheService
|
||||
from src.services.cache.model_list_cache import invalidate_models_list_cache
|
||||
from src.services.cache.provider_cache import ProviderCacheService
|
||||
from src.services.provider.delete_cleanup import prune_allowed_provider_refs
|
||||
from src.services.provider_keys.key_side_effects import cleanup_key_references
|
||||
|
||||
STATUS_PENDING = "pending"
|
||||
STATUS_RUNNING = "running"
|
||||
STATUS_COMPLETED = "completed"
|
||||
STATUS_FAILED = "failed"
|
||||
|
||||
_TASK_RETAIN_SECONDS = 600
|
||||
_KEY_BATCH_SIZE = 50
|
||||
_ENDPOINT_BATCH_SIZE = 200
|
||||
_BATCH_STATEMENT_TIMEOUT_S = 30
|
||||
_BATCH_LOCK_TIMEOUT_S = 5
|
||||
_TASK_TIMEOUT_S = 1800
|
||||
_REDIS_KEY_PREFIX = "provider_delete_task"
|
||||
|
||||
_running_tasks: set[asyncio.Task[None]] = set()
|
||||
|
||||
|
||||
def _task_key(task_id: str) -> str:
|
||||
return f"{_REDIS_KEY_PREFIX}:{task_id}"
|
||||
|
||||
|
||||
def _provider_lock_key(provider_id: str) -> str:
|
||||
return f"{_REDIS_KEY_PREFIX}:provider:{provider_id}"
|
||||
|
||||
|
||||
class ProviderDeleteTaskInfo:
|
||||
__slots__ = (
|
||||
"task_id",
|
||||
"provider_id",
|
||||
"status",
|
||||
"stage",
|
||||
"total_keys",
|
||||
"deleted_keys",
|
||||
"total_endpoints",
|
||||
"deleted_endpoints",
|
||||
"message",
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
task_id: str,
|
||||
provider_id: str,
|
||||
status: str = STATUS_PENDING,
|
||||
stage: str = "queued",
|
||||
total_keys: int = 0,
|
||||
deleted_keys: int = 0,
|
||||
total_endpoints: int = 0,
|
||||
deleted_endpoints: int = 0,
|
||||
message: str = "",
|
||||
) -> None:
|
||||
self.task_id = task_id
|
||||
self.provider_id = provider_id
|
||||
self.status = status
|
||||
self.stage = stage
|
||||
self.total_keys = total_keys
|
||||
self.deleted_keys = deleted_keys
|
||||
self.total_endpoints = total_endpoints
|
||||
self.deleted_endpoints = deleted_endpoints
|
||||
self.message = message
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"task_id": self.task_id,
|
||||
"provider_id": self.provider_id,
|
||||
"status": self.status,
|
||||
"stage": self.stage,
|
||||
"total_keys": self.total_keys,
|
||||
"deleted_keys": self.deleted_keys,
|
||||
"total_endpoints": self.total_endpoints,
|
||||
"deleted_endpoints": self.deleted_endpoints,
|
||||
"message": self.message,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> "ProviderDeleteTaskInfo":
|
||||
return cls(
|
||||
task_id=str(data["task_id"]),
|
||||
provider_id=str(data["provider_id"]),
|
||||
status=str(data.get("status", STATUS_PENDING)),
|
||||
stage=str(data.get("stage", "queued")),
|
||||
total_keys=int(data.get("total_keys", 0)),
|
||||
deleted_keys=int(data.get("deleted_keys", 0)),
|
||||
total_endpoints=int(data.get("total_endpoints", 0)),
|
||||
deleted_endpoints=int(data.get("deleted_endpoints", 0)),
|
||||
message=str(data.get("message", "")),
|
||||
)
|
||||
|
||||
|
||||
async def _save_task(
|
||||
task: ProviderDeleteTaskInfo,
|
||||
ttl: int = _TASK_RETAIN_SECONDS,
|
||||
r: aioredis.Redis | None = None,
|
||||
) -> None:
|
||||
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()))
|
||||
await r.setex(_provider_lock_key(task.provider_id), ttl, task.task_id)
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to save provider delete task: {}", exc)
|
||||
|
||||
|
||||
async def _load_task(
|
||||
task_id: str,
|
||||
r: aioredis.Redis | None = None,
|
||||
) -> ProviderDeleteTaskInfo | None:
|
||||
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 ProviderDeleteTaskInfo.from_dict(json.loads(data))
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to load provider delete task: {}", exc)
|
||||
return None
|
||||
|
||||
|
||||
async def _update_task_field(
|
||||
task_id: str,
|
||||
r: aioredis.Redis | None = None,
|
||||
**fields: object,
|
||||
) -> None:
|
||||
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 key, value in fields.items():
|
||||
setattr(task, key, value)
|
||||
ttl = (
|
||||
_TASK_RETAIN_SECONDS
|
||||
if task.status in (STATUS_COMPLETED, STATUS_FAILED)
|
||||
else _TASK_RETAIN_SECONDS * 2
|
||||
)
|
||||
await _save_task(task, ttl=ttl, r=r)
|
||||
|
||||
|
||||
async def submit_provider_delete(provider_id: str) -> str:
|
||||
r = await get_redis_client(require_redis=False)
|
||||
if not r:
|
||||
raise RuntimeError("Redis is required for provider delete tasks but is not available")
|
||||
|
||||
existing_task_id = await r.get(_provider_lock_key(provider_id))
|
||||
if isinstance(existing_task_id, bytes):
|
||||
existing_task_id = existing_task_id.decode()
|
||||
if existing_task_id:
|
||||
existing_task = await _load_task(str(existing_task_id), r=r)
|
||||
if existing_task and existing_task.status in (STATUS_PENDING, STATUS_RUNNING):
|
||||
return existing_task.task_id
|
||||
|
||||
task_id = uuid.uuid4().hex[:16]
|
||||
task = ProviderDeleteTaskInfo(
|
||||
task_id=task_id,
|
||||
provider_id=provider_id,
|
||||
message="delete task submitted",
|
||||
)
|
||||
await _save_task(task, ttl=_TASK_RETAIN_SECONDS * 2, r=r)
|
||||
|
||||
def _on_task_done(task: asyncio.Task[None]) -> None:
|
||||
_running_tasks.discard(task)
|
||||
if not task.cancelled() and task.exception():
|
||||
logger.error("[PROVIDER_DELETE_TASK] unhandled error: {}", task.exception())
|
||||
|
||||
bg = asyncio.create_task(
|
||||
_run_provider_delete(task_id, provider_id),
|
||||
name=f"provider-delete-{task_id}",
|
||||
)
|
||||
_running_tasks.add(bg)
|
||||
bg.add_done_callback(_on_task_done)
|
||||
return task_id
|
||||
|
||||
|
||||
async def get_provider_delete_task(task_id: str) -> ProviderDeleteTaskInfo | None:
|
||||
return await _load_task(task_id)
|
||||
|
||||
|
||||
def _iter_batches(items: Sequence[str], batch_size: int) -> list[list[str]]:
|
||||
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)]
|
||||
|
||||
|
||||
def _apply_statement_timeouts(db: Any) -> 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 _collect_ids(db: Any, provider_id: str) -> tuple[list[str], list[str]]:
|
||||
endpoint_ids = [
|
||||
endpoint_id
|
||||
for endpoint_id, in db.query(ProviderEndpoint.id)
|
||||
.filter(ProviderEndpoint.provider_id == provider_id)
|
||||
.all()
|
||||
]
|
||||
key_ids = [
|
||||
key_id
|
||||
for key_id, in db.query(ProviderAPIKey.id)
|
||||
.filter(ProviderAPIKey.provider_id == provider_id)
|
||||
.all()
|
||||
]
|
||||
return endpoint_ids, key_ids
|
||||
|
||||
|
||||
def _sync_delete_provider(
|
||||
provider_id: str,
|
||||
progress_callback: Callable[[dict[str, object]], None] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
db = create_session()
|
||||
try:
|
||||
task_start = time.monotonic()
|
||||
provider = db.query(Provider).filter(Provider.id == provider_id).first()
|
||||
if not provider:
|
||||
raise RuntimeError("provider not found")
|
||||
|
||||
_apply_statement_timeouts(db)
|
||||
endpoint_ids, key_ids = _collect_ids(db, provider_id)
|
||||
if progress_callback is not None:
|
||||
progress_callback(
|
||||
{
|
||||
"stage": "preparing",
|
||||
"total_keys": len(key_ids),
|
||||
"total_endpoints": len(endpoint_ids),
|
||||
"message": f"preparing delete for {len(key_ids)} keys and {len(endpoint_ids)} endpoints",
|
||||
}
|
||||
)
|
||||
|
||||
if getattr(provider, "is_active", False):
|
||||
provider.is_active = False
|
||||
db.commit()
|
||||
if progress_callback is not None:
|
||||
progress_callback(
|
||||
{
|
||||
"stage": "disabling",
|
||||
"message": "provider disabled; starting cleanup",
|
||||
}
|
||||
)
|
||||
|
||||
_apply_statement_timeouts(db)
|
||||
updated_users = prune_allowed_provider_refs(
|
||||
db.query(User).filter(User.allowed_providers.isnot(None)).all(),
|
||||
provider_id,
|
||||
)
|
||||
updated_api_keys = prune_allowed_provider_refs(
|
||||
db.query(ApiKey).filter(ApiKey.allowed_providers.isnot(None)).all(),
|
||||
provider_id,
|
||||
)
|
||||
db.commit()
|
||||
if progress_callback is not None:
|
||||
progress_callback(
|
||||
{
|
||||
"stage": "cleaning_restrictions",
|
||||
"message": f"cleaned access restrictions (users={updated_users}, api_keys={updated_api_keys})",
|
||||
}
|
||||
)
|
||||
|
||||
_apply_statement_timeouts(db)
|
||||
db.query(UserPreference).filter(UserPreference.default_provider_id == provider_id).update(
|
||||
{UserPreference.default_provider_id: None},
|
||||
synchronize_session=False,
|
||||
)
|
||||
db.query(Usage).filter(Usage.provider_id == provider_id).update(
|
||||
{Usage.provider_id: None},
|
||||
synchronize_session=False,
|
||||
)
|
||||
db.query(VideoTask).filter(VideoTask.provider_id == provider_id).update(
|
||||
{VideoTask.provider_id: None},
|
||||
synchronize_session=False,
|
||||
)
|
||||
db.query(RequestCandidate).filter(RequestCandidate.provider_id == provider_id).delete(
|
||||
synchronize_session=False,
|
||||
)
|
||||
db.commit()
|
||||
if progress_callback is not None:
|
||||
progress_callback(
|
||||
{
|
||||
"stage": "cleaning_provider_refs",
|
||||
"message": "cleaned provider-wide history references",
|
||||
}
|
||||
)
|
||||
|
||||
deleted_keys = 0
|
||||
key_batches = _iter_batches(key_ids, _KEY_BATCH_SIZE)
|
||||
for index, batch in enumerate(key_batches, start=1):
|
||||
if time.monotonic() - task_start > _TASK_TIMEOUT_S:
|
||||
raise RuntimeError(f"task timeout after {_TASK_TIMEOUT_S}s")
|
||||
_apply_statement_timeouts(db)
|
||||
cleanup_key_references(db, batch)
|
||||
deleted_batch = int(
|
||||
db.query(ProviderAPIKey)
|
||||
.filter(ProviderAPIKey.provider_id == provider_id, ProviderAPIKey.id.in_(batch))
|
||||
.delete(synchronize_session=False)
|
||||
or 0
|
||||
)
|
||||
db.commit()
|
||||
deleted_keys += deleted_batch
|
||||
if progress_callback is not None:
|
||||
progress_callback(
|
||||
{
|
||||
"stage": "deleting_keys",
|
||||
"deleted_keys": deleted_keys,
|
||||
"message": f"deleted key batch {index}/{max(len(key_batches), 1)}",
|
||||
}
|
||||
)
|
||||
|
||||
deleted_endpoints = 0
|
||||
endpoint_batches = _iter_batches(endpoint_ids, _ENDPOINT_BATCH_SIZE)
|
||||
for index, batch in enumerate(endpoint_batches, start=1):
|
||||
if time.monotonic() - task_start > _TASK_TIMEOUT_S:
|
||||
raise RuntimeError(f"task timeout after {_TASK_TIMEOUT_S}s")
|
||||
_apply_statement_timeouts(db)
|
||||
db.query(Usage).filter(Usage.provider_endpoint_id.in_(batch)).update(
|
||||
{Usage.provider_endpoint_id: None},
|
||||
synchronize_session=False,
|
||||
)
|
||||
db.query(VideoTask).filter(VideoTask.endpoint_id.in_(batch)).update(
|
||||
{VideoTask.endpoint_id: None},
|
||||
synchronize_session=False,
|
||||
)
|
||||
db.query(RequestCandidate).filter(RequestCandidate.endpoint_id.in_(batch)).delete(
|
||||
synchronize_session=False,
|
||||
)
|
||||
deleted_batch = int(
|
||||
db.query(ProviderEndpoint)
|
||||
.filter(ProviderEndpoint.provider_id == provider_id, ProviderEndpoint.id.in_(batch))
|
||||
.delete(synchronize_session=False)
|
||||
or 0
|
||||
)
|
||||
db.commit()
|
||||
deleted_endpoints += deleted_batch
|
||||
if progress_callback is not None:
|
||||
progress_callback(
|
||||
{
|
||||
"stage": "deleting_endpoints",
|
||||
"deleted_endpoints": deleted_endpoints,
|
||||
"message": f"deleted endpoint batch {index}/{max(len(endpoint_batches), 1)}",
|
||||
}
|
||||
)
|
||||
|
||||
_apply_statement_timeouts(db)
|
||||
deleted_models = int(
|
||||
db.query(Model)
|
||||
.filter(Model.provider_id == provider_id)
|
||||
.delete(synchronize_session=False)
|
||||
or 0
|
||||
)
|
||||
deleted_mappings = int(
|
||||
db.query(ApiKeyProviderMapping)
|
||||
.filter(ApiKeyProviderMapping.provider_id == provider_id)
|
||||
.delete(synchronize_session=False)
|
||||
or 0
|
||||
)
|
||||
deleted_usage_tracking = int(
|
||||
db.query(ProviderUsageTracking)
|
||||
.filter(ProviderUsageTracking.provider_id == provider_id)
|
||||
.delete(synchronize_session=False)
|
||||
or 0
|
||||
)
|
||||
deleted_provider = int(
|
||||
db.query(Provider).filter(Provider.id == provider_id).delete(synchronize_session=False)
|
||||
or 0
|
||||
)
|
||||
db.commit()
|
||||
|
||||
return {
|
||||
"provider_id": provider_id,
|
||||
"total_keys": len(key_ids),
|
||||
"deleted_keys": deleted_keys,
|
||||
"total_endpoints": len(endpoint_ids),
|
||||
"deleted_endpoints": deleted_endpoints,
|
||||
"deleted_models": deleted_models,
|
||||
"deleted_mappings": deleted_mappings,
|
||||
"deleted_usage_tracking": deleted_usage_tracking,
|
||||
"deleted_provider": deleted_provider,
|
||||
"elapsed_seconds": time.monotonic() - task_start,
|
||||
}
|
||||
except Exception:
|
||||
try:
|
||||
db.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
raise
|
||||
finally:
|
||||
try:
|
||||
db.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
async def _run_provider_delete(task_id: str, provider_id: str) -> None:
|
||||
r = await get_redis_client(require_redis=False)
|
||||
await _update_task_field(
|
||||
task_id,
|
||||
r=r,
|
||||
status=STATUS_RUNNING,
|
||||
stage="queued",
|
||||
message="delete task started",
|
||||
)
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
progress_futures: list[Future[object]] = []
|
||||
|
||||
def on_progress(fields: dict[str, object]) -> None:
|
||||
try:
|
||||
future = asyncio.run_coroutine_threadsafe(
|
||||
_update_task_field(task_id, r=r, **fields),
|
||||
loop,
|
||||
)
|
||||
progress_futures.append(future)
|
||||
except RuntimeError:
|
||||
pass
|
||||
|
||||
async def _drain_progress() -> None:
|
||||
if progress_futures:
|
||||
await asyncio.gather(
|
||||
*(asyncio.wrap_future(f) for f in progress_futures),
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
try:
|
||||
summary = await asyncio.wait_for(
|
||||
asyncio.to_thread(_sync_delete_provider, provider_id, on_progress),
|
||||
timeout=_TASK_TIMEOUT_S + 60,
|
||||
)
|
||||
|
||||
await _drain_progress()
|
||||
|
||||
try:
|
||||
await invalidate_models_list_cache()
|
||||
await ModelCacheService.invalidate_all_resolve_cache()
|
||||
await ProviderCacheService.invalidate_provider_cache(provider_id)
|
||||
except Exception as exc:
|
||||
logger.error("provider delete cache invalidation failed: {}", exc)
|
||||
|
||||
await _update_task_field(
|
||||
task_id,
|
||||
r=r,
|
||||
status=STATUS_COMPLETED,
|
||||
stage="completed",
|
||||
total_keys=int(summary.get("total_keys", 0)),
|
||||
deleted_keys=int(summary.get("deleted_keys", 0)),
|
||||
total_endpoints=int(summary.get("total_endpoints", 0)),
|
||||
deleted_endpoints=int(summary.get("deleted_endpoints", 0)),
|
||||
message=(
|
||||
f"provider deleted: keys={summary.get('deleted_keys', 0)}, "
|
||||
f"endpoints={summary.get('deleted_endpoints', 0)}"
|
||||
),
|
||||
)
|
||||
logger.info(
|
||||
"[PROVIDER_DELETE_TASK] completed task={} provider={} keys={}/{} endpoints={}/{} elapsed={:.1f}s",
|
||||
task_id,
|
||||
provider_id[:8],
|
||||
summary.get("deleted_keys", 0),
|
||||
summary.get("total_keys", 0),
|
||||
summary.get("deleted_endpoints", 0),
|
||||
summary.get("total_endpoints", 0),
|
||||
float(summary.get("elapsed_seconds", 0.0) or 0.0),
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
await _drain_progress()
|
||||
msg = f"task timeout after {_TASK_TIMEOUT_S + 60}s"
|
||||
await _update_task_field(task_id, r=r, status=STATUS_FAILED, stage="failed", message=msg)
|
||||
logger.error("[PROVIDER_DELETE_TASK] {} task={} provider={}", msg, task_id, provider_id[:8])
|
||||
except asyncio.CancelledError:
|
||||
await _drain_progress()
|
||||
await _update_task_field(
|
||||
task_id,
|
||||
r=r,
|
||||
status=STATUS_FAILED,
|
||||
stage="failed",
|
||||
message="task cancelled (shutdown)",
|
||||
)
|
||||
logger.warning(
|
||||
"[PROVIDER_DELETE_TASK] cancelled task={} provider={}",
|
||||
task_id,
|
||||
provider_id[:8],
|
||||
)
|
||||
except Exception as exc:
|
||||
await _drain_progress()
|
||||
await _update_task_field(
|
||||
task_id,
|
||||
r=r,
|
||||
status=STATUS_FAILED,
|
||||
stage="failed",
|
||||
message=str(exc),
|
||||
)
|
||||
logger.error(
|
||||
"[PROVIDER_DELETE_TASK] failed task={} provider={} error={}",
|
||||
task_id,
|
||||
provider_id[:8],
|
||||
exc,
|
||||
)
|
||||
267
_deprecated_py_src/services/provider/envelope.py
Normal file
267
_deprecated_py_src/services/provider/envelope.py
Normal file
@@ -0,0 +1,267 @@
|
||||
"""Provider request/response envelope hooks.
|
||||
|
||||
Some upstreams expose an API that is *almost* compatible with an existing
|
||||
endpoint signature (family:kind), but wrap the wire format in an extra envelope
|
||||
or require small transport-level behaviors.
|
||||
|
||||
This module provides a small hook mechanism so handlers can stay generic while
|
||||
provider-specific envelopes live in their own service modules.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import threading
|
||||
from collections.abc import Iterable
|
||||
from typing import Any, Protocol
|
||||
|
||||
|
||||
class ProviderEnvelope(Protocol):
|
||||
"""Provider-specific envelope transformation and side-effects."""
|
||||
|
||||
name: str
|
||||
|
||||
def extra_headers(self) -> dict[str, str] | None:
|
||||
"""Extra upstream request headers to merge into the RequestBuilder."""
|
||||
|
||||
def wrap_request(
|
||||
self,
|
||||
request_body: dict[str, Any],
|
||||
*,
|
||||
model: str,
|
||||
url_model: str | None,
|
||||
decrypted_auth_config: dict[str, Any] | None,
|
||||
) -> tuple[dict[str, Any], str | None]:
|
||||
"""Wrap request payload and optionally override url_model (e.g. move model into body)."""
|
||||
|
||||
def unwrap_response(self, data: Any) -> Any:
|
||||
"""Unwrap upstream response payload (streaming chunk or full JSON)."""
|
||||
|
||||
def postprocess_unwrapped_response(self, *, model: str, data: Any) -> None:
|
||||
"""Best-effort post processing after unwrap (e.g. cache signatures)."""
|
||||
|
||||
def capture_selected_base_url(self) -> str | None:
|
||||
"""Capture the base_url selected by transport layer (if any)."""
|
||||
|
||||
def on_http_status(self, *, base_url: str | None, status_code: int) -> None:
|
||||
"""Called after receiving upstream HTTP status code."""
|
||||
|
||||
def on_connection_error(self, *, base_url: str | None, exc: Exception) -> None:
|
||||
"""Called when a connection-type exception happens."""
|
||||
|
||||
def force_stream_rewrite(self) -> bool:
|
||||
"""Whether streaming should always go through the rewrite/conversion path."""
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Optional lifecycle hooks (checked via hasattr before calling)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def prepare_context(
|
||||
self,
|
||||
*,
|
||||
provider_config: Any,
|
||||
key_id: str,
|
||||
user_api_key_id: str | None = None,
|
||||
is_stream: bool,
|
||||
provider_id: str | None = None,
|
||||
key: Any = None,
|
||||
) -> str | None:
|
||||
"""Pre-wrap hook: build provider-specific request context.
|
||||
|
||||
Called before wrap_request(). Returns tls_profile (or None).
|
||||
Implementations typically set contextvars that wrap_request()
|
||||
and extra_headers() will read.
|
||||
"""
|
||||
|
||||
async def post_wrap_request(self, request_body: dict[str, Any]) -> None:
|
||||
"""Post-wrap hook: async processing after wrap_request().
|
||||
|
||||
Called after wrap_request() completes. Use for async operations
|
||||
like distributed session control that cannot run in sync wrap_request().
|
||||
"""
|
||||
|
||||
def excluded_beta_tokens(self) -> frozenset[str]:
|
||||
"""Beta tokens to strip from the merged anthropic-beta header.
|
||||
|
||||
Called by the request builder after merging envelope extra_headers
|
||||
with client original headers. Return an empty frozenset to keep all.
|
||||
"""
|
||||
|
||||
async def extract_error_text(
|
||||
self,
|
||||
source: Any,
|
||||
*,
|
||||
limit: int = 4000,
|
||||
) -> str:
|
||||
"""Extract error text from upstream HTTP error response.
|
||||
|
||||
``source`` is either an ``httpx.Response`` or ``httpx.HTTPStatusError``.
|
||||
Default behavior (when not overridden) is handled by the caller.
|
||||
Implementations may parse provider-specific error formats.
|
||||
"""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Envelope Registry
|
||||
# ---------------------------------------------------------------------------
|
||||
# key: (provider_type, endpoint_sig) — endpoint_sig="" 表示通配
|
||||
_envelope_registry: dict[tuple[str, str], ProviderEnvelope] = {}
|
||||
|
||||
|
||||
def register_envelope(
|
||||
provider_type: str,
|
||||
endpoint_sig: str,
|
||||
envelope: ProviderEnvelope,
|
||||
) -> None:
|
||||
"""注册 provider 特有的 envelope。
|
||||
|
||||
Args:
|
||||
provider_type: 如 "antigravity"
|
||||
endpoint_sig: 如 "gemini:cli",传 "" 表示该 provider 的所有 endpoint
|
||||
envelope: 实现了 ProviderEnvelope 协议的实例
|
||||
"""
|
||||
from src.core.provider_types import normalize_provider_type
|
||||
|
||||
pt = normalize_provider_type(provider_type)
|
||||
sig = str(endpoint_sig or "").strip().lower()
|
||||
_envelope_registry[(pt, sig)] = envelope
|
||||
|
||||
|
||||
def get_provider_envelope(
|
||||
*,
|
||||
provider_type: str | None,
|
||||
endpoint_sig: str | None,
|
||||
) -> ProviderEnvelope | None:
|
||||
"""Return envelope hooks for the given provider_type + endpoint signature."""
|
||||
ensure_providers_bootstrapped(provider_types=[provider_type] if provider_type else None)
|
||||
|
||||
from src.core.provider_types import normalize_provider_type
|
||||
|
||||
pt = normalize_provider_type(provider_type)
|
||||
sig = str(endpoint_sig or "").strip().lower()
|
||||
|
||||
if not pt:
|
||||
return None
|
||||
|
||||
# 精确匹配优先,再尝试通配
|
||||
return _envelope_registry.get((pt, sig)) or _envelope_registry.get((pt, ""))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Provider Bootstrap(惰性 + 幂等)
|
||||
# ---------------------------------------------------------------------------
|
||||
# 所有 registry 共享同一个 bootstrap,首次访问任何 registry 时自动触发。
|
||||
# 不再依赖模块 import 顺序。
|
||||
_bootstrap_lock = threading.Lock()
|
||||
_bootstrap_condition = threading.Condition(_bootstrap_lock)
|
||||
_bootstrap_in_progress = False
|
||||
_bootstrapped_provider_types: set[str] = set()
|
||||
_auto_detected_provider_types: frozenset[str] | None = None
|
||||
|
||||
_PROVIDER_PLUGIN_MODULES: dict[str, str] = {
|
||||
"antigravity": "src.services.provider.adapters.antigravity.plugin",
|
||||
"claude_code": "src.services.provider.adapters.claude_code.plugin",
|
||||
"codex": "src.services.provider.adapters.codex.plugin",
|
||||
"gemini_cli": "src.services.provider.adapters.gemini_cli.plugin",
|
||||
"kiro": "src.services.provider.adapters.kiro.plugin",
|
||||
"vertex_ai": "src.services.provider.adapters.vertex_ai.plugin",
|
||||
}
|
||||
|
||||
|
||||
def _normalize_bootstrap_targets(provider_types: Iterable[str] | None) -> set[str]:
|
||||
from src.core.provider_types import normalize_provider_type
|
||||
|
||||
if provider_types is None:
|
||||
return set()
|
||||
if isinstance(provider_types, str):
|
||||
provider_types = [provider_types]
|
||||
|
||||
targets: set[str] = set()
|
||||
for raw in provider_types:
|
||||
pt = normalize_provider_type(raw)
|
||||
if pt in _PROVIDER_PLUGIN_MODULES:
|
||||
targets.add(pt)
|
||||
return targets
|
||||
|
||||
|
||||
def _discover_active_provider_types() -> set[str]:
|
||||
"""从数据库读取活跃 provider_type,用于按需 bootstrap。"""
|
||||
from src.core.provider_types import normalize_provider_type
|
||||
from src.database.database import create_session
|
||||
from src.models.database import Provider
|
||||
|
||||
db = create_session()
|
||||
try:
|
||||
rows = (
|
||||
db.query(Provider.provider_type).filter(Provider.is_active.is_(True)).distinct().all()
|
||||
)
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
discovered: set[str] = set()
|
||||
for (raw_provider_type,) in rows:
|
||||
pt = normalize_provider_type(raw_provider_type)
|
||||
if pt in _PROVIDER_PLUGIN_MODULES:
|
||||
discovered.add(pt)
|
||||
return discovered
|
||||
|
||||
|
||||
def _bootstrap_provider_type(provider_type: str) -> None:
|
||||
module_path = _PROVIDER_PLUGIN_MODULES[provider_type]
|
||||
module = importlib.import_module(module_path)
|
||||
register_all = getattr(module, "register_all", None)
|
||||
if callable(register_all):
|
||||
register_all()
|
||||
|
||||
|
||||
def ensure_providers_bootstrapped(provider_types: Iterable[str] | None = None) -> None:
|
||||
"""确保 provider plugins 已注册(幂等,支持按 provider_type 精准注册)。"""
|
||||
global _auto_detected_provider_types, _bootstrap_in_progress # noqa: PLW0603
|
||||
|
||||
targets = _normalize_bootstrap_targets(provider_types)
|
||||
|
||||
# DB 查询在锁外执行,避免慢查询阻塞其他线程的 bootstrap 操作。
|
||||
need_discover = not targets and _auto_detected_provider_types is None
|
||||
if need_discover:
|
||||
try:
|
||||
detected = _discover_active_provider_types()
|
||||
except Exception:
|
||||
detected = set()
|
||||
else:
|
||||
detected = set()
|
||||
|
||||
with _bootstrap_condition:
|
||||
if not targets:
|
||||
if _auto_detected_provider_types is None:
|
||||
# 回退策略:DB 不可用/无记录时,保持原有全量 bootstrap 语义。
|
||||
_auto_detected_provider_types = frozenset(
|
||||
detected if detected else _PROVIDER_PLUGIN_MODULES.keys()
|
||||
)
|
||||
targets = set(_auto_detected_provider_types)
|
||||
|
||||
while _bootstrap_in_progress:
|
||||
_bootstrap_condition.wait()
|
||||
|
||||
missing = targets - _bootstrapped_provider_types
|
||||
if not missing:
|
||||
return
|
||||
_bootstrap_in_progress = True
|
||||
|
||||
bootstrapped_now: set[str] = set()
|
||||
try:
|
||||
for pt in sorted(missing):
|
||||
_bootstrap_provider_type(pt)
|
||||
bootstrapped_now.add(pt)
|
||||
finally:
|
||||
with _bootstrap_condition:
|
||||
_bootstrapped_provider_types.update(bootstrapped_now)
|
||||
_bootstrap_in_progress = False
|
||||
_bootstrap_condition.notify_all()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ProviderEnvelope",
|
||||
"ensure_providers_bootstrapped",
|
||||
"get_provider_envelope",
|
||||
"register_envelope",
|
||||
]
|
||||
52
_deprecated_py_src/services/provider/export.py
Normal file
52
_deprecated_py_src/services/provider/export.py
Normal file
@@ -0,0 +1,52 @@
|
||||
"""OAuth Key 导出:provider-specific export builders.
|
||||
|
||||
每个 Provider adapter 在 ``register_all()`` 时注册自己的 ``build_export_data``,
|
||||
导出端点通过 ``build_export_data()`` 分发。
|
||||
|
||||
未注册的 provider_type 使用默认实现(strip null + 临时字段)。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Callable
|
||||
|
||||
# 所有 provider 共用的临时/无用字段
|
||||
_DEFAULT_SKIP_KEYS = frozenset(
|
||||
{
|
||||
"access_token",
|
||||
"expires_at",
|
||||
"updated_at",
|
||||
"token_type",
|
||||
"scope",
|
||||
}
|
||||
)
|
||||
|
||||
ExportBuilder = Callable[[dict[str, Any], dict[str, Any] | None], dict[str, Any]]
|
||||
|
||||
_BUILDERS: dict[str, ExportBuilder] = {}
|
||||
|
||||
|
||||
def register_export_builder(provider_type: str, builder: ExportBuilder) -> None:
|
||||
_BUILDERS[provider_type] = builder
|
||||
|
||||
|
||||
def _default_builder(
|
||||
auth_config: dict[str, Any],
|
||||
upstream_metadata: dict[str, Any] | None,
|
||||
) -> dict[str, Any]:
|
||||
"""默认导出:去掉 null、空字符串、临时字段。"""
|
||||
return {
|
||||
k: v
|
||||
for k, v in auth_config.items()
|
||||
if k not in _DEFAULT_SKIP_KEYS and v is not None and v != ""
|
||||
}
|
||||
|
||||
|
||||
def build_export_data(
|
||||
provider_type: str,
|
||||
auth_config: dict[str, Any],
|
||||
upstream_metadata: dict[str, Any] | None,
|
||||
) -> dict[str, Any]:
|
||||
"""调用 provider-specific builder 构建导出数据。"""
|
||||
builder = _BUILDERS.get(provider_type, _default_builder)
|
||||
return builder(auth_config, upstream_metadata)
|
||||
352
_deprecated_py_src/services/provider/fingerprint.py
Normal file
352
_deprecated_py_src/services/provider/fingerprint.py
Normal file
@@ -0,0 +1,352 @@
|
||||
"""Per-key request fingerprint generation and lazy persistence helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import random
|
||||
import secrets
|
||||
import threading
|
||||
import uuid
|
||||
from dataclasses import asdict, dataclass
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from src.core.logger import logger
|
||||
|
||||
# curl_cffi impersonate profile pool (Chrome family)
|
||||
CHROME_IMPERSONATE_PROFILES: tuple[str, ...] = (
|
||||
"chrome110",
|
||||
"chrome116",
|
||||
"chrome119",
|
||||
"chrome120",
|
||||
"chrome123",
|
||||
"chrome124",
|
||||
"chrome131",
|
||||
"chrome133",
|
||||
)
|
||||
|
||||
KNOWN_IMPERSONATE_PROFILES: frozenset[str] = frozenset(CHROME_IMPERSONATE_PROFILES)
|
||||
|
||||
_CHROME_VERSION_BY_PROFILE: dict[str, str] = {
|
||||
"chrome110": "110.0.5481.177",
|
||||
"chrome116": "116.0.5845.188",
|
||||
"chrome119": "119.0.6045.214",
|
||||
"chrome120": "120.0.6099.216",
|
||||
"chrome123": "123.0.6312.122",
|
||||
"chrome124": "124.0.6367.243",
|
||||
"chrome131": "131.0.6778.265",
|
||||
"chrome133": "133.0.6943.142",
|
||||
}
|
||||
|
||||
_PLATFORM_VARIANTS: tuple[tuple[str, str, str, str], ...] = (
|
||||
("Linux", "x64", "X11; Linux x86_64", "Linux x86_64"),
|
||||
("Linux", "arm64", "X11; Linux arm64", "Linux arm64"),
|
||||
("Windows", "x64", "Windows NT 10.0; Win64; x64", "Windows x64"),
|
||||
("MacOS", "x64", "Macintosh; Intel Mac OS X 10_15_7", "Darwin x64"),
|
||||
("MacOS", "arm64", "Macintosh; ARM Mac OS X 14_0_0", "Darwin arm64"),
|
||||
)
|
||||
|
||||
_STAINLESS_PACKAGE_VERSIONS: tuple[str, ...] = ("0.68.0", "0.69.0", "0.70.0", "0.71.0")
|
||||
_NODE_VERSIONS: tuple[str, ...] = ("v20.18.1", "v22.12.0", "v22.14.0", "v24.13.0")
|
||||
_ELECTRON_VERSIONS: tuple[str, ...] = ("35.5.1", "36.7.1", "37.3.0", "38.7.0", "39.2.3")
|
||||
_STAINLESS_TIMEOUTS: tuple[str, ...] = ("600", "900")
|
||||
|
||||
_PENDING_LAZY_PERSIST: set[str] = set()
|
||||
_PENDING_LAZY_PERSIST_LOCK = threading.Lock()
|
||||
_PENDING_LAZY_PERSIST_MAX = 2000
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class FingerprintProfile:
|
||||
# TLS layer
|
||||
impersonate: str
|
||||
# Claude Code / Claude Chat feature dimensions
|
||||
stainless_package_version: str
|
||||
stainless_os: str
|
||||
stainless_arch: str
|
||||
stainless_runtime_version: str
|
||||
stainless_timeout: str
|
||||
# Generic dimensions
|
||||
user_agent: str
|
||||
node_version: str
|
||||
chrome_version: str
|
||||
electron_version: str
|
||||
vscode_session_id: str
|
||||
platform_info: str
|
||||
|
||||
|
||||
def _build_rng(seed: str | None) -> random.Random:
|
||||
if seed is None:
|
||||
return random.Random(secrets.randbits(64))
|
||||
digest = hashlib.sha256(seed.encode("utf-8")).digest()
|
||||
return random.Random(int.from_bytes(digest[:8], "big", signed=False))
|
||||
|
||||
|
||||
def _normalize_impersonate(raw: Any, fallback: str) -> str:
|
||||
value = str(raw or "").strip().lower()
|
||||
return value if value in KNOWN_IMPERSONATE_PROFILES else fallback
|
||||
|
||||
|
||||
def _build_user_agent(
|
||||
platform_token: str,
|
||||
chrome_version: str,
|
||||
electron_version: str,
|
||||
) -> str:
|
||||
return (
|
||||
f"Mozilla/5.0 ({platform_token}) AppleWebKit/537.36 (KHTML, like Gecko) "
|
||||
f"Chrome/{chrome_version} Electron/{electron_version} Safari/537.36"
|
||||
)
|
||||
|
||||
|
||||
def resolve_platform_token(
|
||||
fp: dict[str, str] | None = None,
|
||||
*,
|
||||
platform_info: str | None = None,
|
||||
stainless_os: str | None = None,
|
||||
stainless_arch: str | None = None,
|
||||
) -> str:
|
||||
"""Resolve a platform token string from fingerprint dict or explicit kwargs."""
|
||||
if fp is not None:
|
||||
os_name = str(fp.get("stainless_os") or "").lower()
|
||||
arch = str(fp.get("stainless_arch") or "").lower()
|
||||
info = str(fp.get("platform_info") or "").lower()
|
||||
else:
|
||||
os_name = str(stainless_os or "").lower()
|
||||
arch = str(stainless_arch or "").lower()
|
||||
info = str(platform_info or "").lower()
|
||||
|
||||
if os_name.startswith("win") or "windows" in info:
|
||||
return "Windows NT 10.0; Win64; x64"
|
||||
if os_name in {"darwin", "mac", "macos"} or "darwin" in info:
|
||||
return (
|
||||
"Macintosh; ARM Mac OS X 14_0_0"
|
||||
if arch in {"arm64", "aarch64"}
|
||||
else "Macintosh; Intel Mac OS X 10_15_7"
|
||||
)
|
||||
return "X11; Linux arm64" if arch in {"arm64", "aarch64"} else "X11; Linux x86_64"
|
||||
|
||||
|
||||
def _normalize_text(value: Any, fallback: str) -> str:
|
||||
text = str(value or "").strip()
|
||||
return text or fallback
|
||||
|
||||
|
||||
def _sanitize_fingerprint_dict(raw: dict[str, Any], key_id: str) -> dict[str, str]:
|
||||
generated = generate_fingerprint(seed=key_id or None)
|
||||
fp: dict[str, str] = {k: str(v) for k, v in generated.items()}
|
||||
for key, value in raw.items():
|
||||
if isinstance(value, str) and value.strip():
|
||||
fp[key] = value.strip()
|
||||
|
||||
fp["impersonate"] = _normalize_impersonate(fp.get("impersonate"), generated["impersonate"])
|
||||
fp["chrome_version"] = _normalize_text(
|
||||
fp.get("chrome_version"),
|
||||
_CHROME_VERSION_BY_PROFILE.get(fp["impersonate"], generated["chrome_version"]),
|
||||
)
|
||||
fp["node_version"] = _normalize_text(fp.get("node_version"), generated["node_version"])
|
||||
fp["electron_version"] = _normalize_text(
|
||||
fp.get("electron_version"),
|
||||
generated["electron_version"],
|
||||
)
|
||||
fp["stainless_package_version"] = _normalize_text(
|
||||
fp.get("stainless_package_version"),
|
||||
generated["stainless_package_version"],
|
||||
)
|
||||
fp["stainless_os"] = _normalize_text(fp.get("stainless_os"), generated["stainless_os"])
|
||||
fp["stainless_arch"] = _normalize_text(fp.get("stainless_arch"), generated["stainless_arch"])
|
||||
fp["stainless_runtime_version"] = _normalize_text(
|
||||
fp.get("stainless_runtime_version"),
|
||||
generated["stainless_runtime_version"],
|
||||
)
|
||||
fp["stainless_timeout"] = _normalize_text(
|
||||
fp.get("stainless_timeout"),
|
||||
generated["stainless_timeout"],
|
||||
)
|
||||
fp["platform_info"] = _normalize_text(fp.get("platform_info"), generated["platform_info"])
|
||||
fp["vscode_session_id"] = _normalize_text(
|
||||
fp.get("vscode_session_id"),
|
||||
generated["vscode_session_id"],
|
||||
)
|
||||
|
||||
if not str(fp.get("user_agent") or "").strip():
|
||||
fp["user_agent"] = _build_user_agent(
|
||||
resolve_platform_token(fp),
|
||||
fp["chrome_version"],
|
||||
fp["electron_version"],
|
||||
)
|
||||
|
||||
return fp
|
||||
|
||||
|
||||
def generate_fingerprint(seed: str | None = None) -> dict[str, str]:
|
||||
"""Generate a serializable fingerprint profile dict."""
|
||||
rng = _build_rng(seed)
|
||||
|
||||
impersonate = rng.choice(CHROME_IMPERSONATE_PROFILES)
|
||||
chrome_version = _CHROME_VERSION_BY_PROFILE.get(impersonate, "120.0.6099.216")
|
||||
node_version = rng.choice(_NODE_VERSIONS)
|
||||
electron_version = rng.choice(_ELECTRON_VERSIONS)
|
||||
stainless_os, stainless_arch, platform_token, platform_info = rng.choice(_PLATFORM_VARIANTS)
|
||||
|
||||
if seed is None:
|
||||
vscode_session_id = uuid.uuid4().hex
|
||||
else:
|
||||
vscode_session_id = uuid.uuid5(uuid.NAMESPACE_URL, f"aether:fingerprint:{seed}").hex
|
||||
|
||||
return {
|
||||
"impersonate": impersonate,
|
||||
"stainless_package_version": rng.choice(_STAINLESS_PACKAGE_VERSIONS),
|
||||
"stainless_os": stainless_os,
|
||||
"stainless_arch": stainless_arch,
|
||||
"stainless_runtime_version": node_version,
|
||||
"stainless_timeout": rng.choice(_STAINLESS_TIMEOUTS),
|
||||
"node_version": node_version,
|
||||
"chrome_version": chrome_version,
|
||||
"electron_version": electron_version,
|
||||
"vscode_session_id": vscode_session_id,
|
||||
"platform_info": platform_info,
|
||||
"user_agent": _build_user_agent(platform_token, chrome_version, electron_version),
|
||||
}
|
||||
|
||||
|
||||
def _dict_to_profile(d: dict[str, str]) -> FingerprintProfile:
|
||||
return FingerprintProfile(
|
||||
impersonate=d["impersonate"],
|
||||
stainless_package_version=d["stainless_package_version"],
|
||||
stainless_os=d["stainless_os"],
|
||||
stainless_arch=d["stainless_arch"],
|
||||
stainless_runtime_version=d["stainless_runtime_version"],
|
||||
stainless_timeout=d["stainless_timeout"],
|
||||
user_agent=d["user_agent"],
|
||||
node_version=d["node_version"],
|
||||
chrome_version=d["chrome_version"],
|
||||
electron_version=d["electron_version"],
|
||||
vscode_session_id=d["vscode_session_id"],
|
||||
platform_info=d["platform_info"],
|
||||
)
|
||||
|
||||
|
||||
def load_fingerprint(raw: dict[str, Any] | None, key_id: str) -> FingerprintProfile:
|
||||
"""Load fingerprint from DB JSON with deterministic fallback by key_id."""
|
||||
if raw and isinstance(raw, dict):
|
||||
normalized = _sanitize_fingerprint_dict(raw, key_id)
|
||||
else:
|
||||
normalized = generate_fingerprint(seed=key_id or None)
|
||||
return _dict_to_profile(normalized)
|
||||
|
||||
|
||||
def serialize_fingerprint(fp: FingerprintProfile) -> dict[str, str]:
|
||||
return {k: str(v) for k, v in asdict(fp).items()}
|
||||
|
||||
|
||||
def normalize_fingerprint(raw: dict[str, Any], key_id: str) -> dict[str, str]:
|
||||
return serialize_fingerprint(load_fingerprint(raw, key_id))
|
||||
|
||||
|
||||
def _persist_fingerprint_if_missing_sync(key_id: str, fp: dict[str, str]) -> None:
|
||||
from src.database import create_session
|
||||
from src.models.database import ProviderAPIKey
|
||||
|
||||
db = create_session()
|
||||
try:
|
||||
updated = (
|
||||
db.query(ProviderAPIKey)
|
||||
.filter(
|
||||
ProviderAPIKey.id == key_id,
|
||||
ProviderAPIKey.fingerprint.is_(None),
|
||||
)
|
||||
.update(
|
||||
{
|
||||
ProviderAPIKey.fingerprint: fp,
|
||||
ProviderAPIKey.updated_at: datetime.now(timezone.utc),
|
||||
},
|
||||
synchronize_session=False,
|
||||
)
|
||||
)
|
||||
if updated > 0:
|
||||
db.commit()
|
||||
except Exception as exc:
|
||||
db.rollback()
|
||||
logger.debug("lazy fingerprint persist failed for key {}: {}", key_id[:8], str(exc))
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
def _mark_pending_persist(key_id: str) -> bool:
|
||||
with _PENDING_LAZY_PERSIST_LOCK:
|
||||
if key_id in _PENDING_LAZY_PERSIST:
|
||||
return False
|
||||
if len(_PENDING_LAZY_PERSIST) >= _PENDING_LAZY_PERSIST_MAX:
|
||||
return False
|
||||
_PENDING_LAZY_PERSIST.add(key_id)
|
||||
return True
|
||||
|
||||
|
||||
def _clear_pending_persist(key_id: str) -> None:
|
||||
with _PENDING_LAZY_PERSIST_LOCK:
|
||||
_PENDING_LAZY_PERSIST.discard(key_id)
|
||||
|
||||
|
||||
def schedule_lazy_fingerprint_persist(key_id: str, fp: dict[str, str]) -> None:
|
||||
"""Persist generated fingerprint in background when key.fingerprint is missing."""
|
||||
key_id = str(key_id or "").strip()
|
||||
if not key_id:
|
||||
return
|
||||
if not _mark_pending_persist(key_id):
|
||||
return
|
||||
|
||||
payload = dict(fp)
|
||||
|
||||
async def _persist_async() -> None:
|
||||
try:
|
||||
await asyncio.to_thread(_persist_fingerprint_if_missing_sync, key_id, payload)
|
||||
finally:
|
||||
_clear_pending_persist(key_id)
|
||||
|
||||
try:
|
||||
asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
try:
|
||||
_persist_fingerprint_if_missing_sync(key_id, payload)
|
||||
finally:
|
||||
_clear_pending_persist(key_id)
|
||||
return
|
||||
|
||||
from src.utils.async_utils import safe_create_task
|
||||
|
||||
safe_create_task(_persist_async())
|
||||
|
||||
|
||||
def ensure_key_fingerprint(
|
||||
key: Any,
|
||||
*,
|
||||
persist_if_missing: bool = False,
|
||||
) -> FingerprintProfile:
|
||||
"""Get a key fingerprint, generating deterministic fallback if missing."""
|
||||
key_id = str(getattr(key, "id", "") or "").strip()
|
||||
raw = getattr(key, "fingerprint", None)
|
||||
|
||||
if isinstance(raw, dict) and raw:
|
||||
return load_fingerprint(raw, key_id)
|
||||
|
||||
generated = generate_fingerprint(seed=key_id or None)
|
||||
|
||||
if persist_if_missing and key_id:
|
||||
schedule_lazy_fingerprint_persist(key_id, generated)
|
||||
|
||||
return _dict_to_profile(generated)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"CHROME_IMPERSONATE_PROFILES",
|
||||
"FingerprintProfile",
|
||||
"KNOWN_IMPERSONATE_PROFILES",
|
||||
"ensure_key_fingerprint",
|
||||
"generate_fingerprint",
|
||||
"load_fingerprint",
|
||||
"normalize_fingerprint",
|
||||
"resolve_platform_token",
|
||||
"schedule_lazy_fingerprint_persist",
|
||||
"serialize_fingerprint",
|
||||
]
|
||||
45
_deprecated_py_src/services/provider/format.py
Normal file
45
_deprecated_py_src/services/provider/format.py
Normal file
@@ -0,0 +1,45 @@
|
||||
"""
|
||||
Endpoint signature 辅助函数。
|
||||
|
||||
调度/编排链路使用 endpoint signature key:`family:kind`(如 "claude:chat", "openai:cli")。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from src.core.api_format.enums import ApiFamily, EndpointKind
|
||||
from src.core.api_format.signature import (
|
||||
EndpointSignature,
|
||||
make_signature_key,
|
||||
normalize_signature_key,
|
||||
)
|
||||
|
||||
DEFAULT_ENDPOINT_SIGNATURE: str = make_signature_key(ApiFamily.CLAUDE, EndpointKind.CHAT)
|
||||
|
||||
|
||||
def normalize_endpoint_signature(
|
||||
value: str | EndpointSignature | tuple[ApiFamily, EndpointKind] | None,
|
||||
*,
|
||||
default: str = DEFAULT_ENDPOINT_SIGNATURE,
|
||||
) -> str:
|
||||
"""
|
||||
将任意输入归一化为 canonical signature key(`family:kind`,小写)。
|
||||
|
||||
不支持旧格式(如 "CLAUDE_CLI"),仅接受 `family:kind` 格式。
|
||||
解析失败时返回默认值。
|
||||
"""
|
||||
if value is None:
|
||||
return default
|
||||
if isinstance(value, EndpointSignature):
|
||||
return value.key
|
||||
if isinstance(value, tuple) and len(value) == 2:
|
||||
fam, kind = value
|
||||
if isinstance(fam, ApiFamily) and isinstance(kind, EndpointKind):
|
||||
return make_signature_key(fam, kind)
|
||||
return default
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
return normalize_signature_key(value)
|
||||
except ValueError:
|
||||
# 如果解析失败,返回默认值
|
||||
return default
|
||||
return default
|
||||
213
_deprecated_py_src/services/provider/oauth_token.py
Normal file
213
_deprecated_py_src/services/provider/oauth_token.py
Normal file
@@ -0,0 +1,213 @@
|
||||
"""Provider OAuth token helpers.
|
||||
|
||||
These helpers are for *upstream Provider* OAuth keys (ProviderAPIKey.auth_type == "oauth"),
|
||||
not for user-login OAuth.
|
||||
|
||||
Why:
|
||||
- Request path uses `get_provider_auth()` which may refresh the access_token lazily.
|
||||
- Some background/admin paths (model fetch/query, etc.) need the same behavior but must
|
||||
avoid sharing a SQLAlchemy Session across concurrent async tasks.
|
||||
|
||||
Strategy:
|
||||
- Run `get_provider_auth()` on a detached key-like object (no DB session held during HTTP).
|
||||
- If refresh updated encrypted fields, persist them back to DB in a short transaction.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.database import create_session
|
||||
from src.models.database import ProviderAPIKey
|
||||
from src.services.provider.pool.account_state import (
|
||||
OAUTH_EXPIRED_PREFIX,
|
||||
OAUTH_REFRESH_FAILED_PREFIX,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Account-level block 结构化标记
|
||||
# ---------------------------------------------------------------------------
|
||||
# oauth_invalid_reason 以此前缀开头的,属于"账号级别"异常(如 Google 要求验证账号);
|
||||
# 刷新 token 无法修复,必须由用户手动解决后再由管理员手动清除。
|
||||
# 其余 reason 属于 token 级别异常,成功刷新 token 后自动清除。
|
||||
OAUTH_ACCOUNT_BLOCK_PREFIX = "[ACCOUNT_BLOCK] "
|
||||
|
||||
# 上游返回 "token 已失效" 语义的关键词(小写匹配)。
|
||||
# 被 codex_refresher 前向分类和 oauth_token 回溯清理共用。
|
||||
TOKEN_INVALIDATED_KEYWORDS: tuple[str, ...] = (
|
||||
"authentication token has been invalidated",
|
||||
"token has been invalidated",
|
||||
)
|
||||
|
||||
# 回溯清理专用:历史写入的中文 reason 也需匹配
|
||||
_LEGACY_TOKEN_INVALID_KEYWORDS: tuple[str, ...] = (
|
||||
*TOKEN_INVALIDATED_KEYWORDS,
|
||||
"codex token 无效或已过期",
|
||||
)
|
||||
|
||||
|
||||
def looks_like_token_invalidated(message: str | None) -> bool:
|
||||
"""判断上游错误消息是否表示 access token 已失效/被轮换。"""
|
||||
lowered = str(message or "").strip().lower()
|
||||
return any(keyword in lowered for keyword in TOKEN_INVALIDATED_KEYWORDS)
|
||||
|
||||
|
||||
def _is_refresh_recoverable_account_block(reason: str | None) -> bool:
|
||||
"""历史兼容:部分 token 级异常曾被错误写成 [ACCOUNT_BLOCK]。
|
||||
|
||||
这类原因在手动刷新成功后应自动清除,否则前端会继续展示
|
||||
"Token 失效/账号异常",并阻止 Key 恢复调度。
|
||||
"""
|
||||
if not reason:
|
||||
return False
|
||||
text = str(reason)
|
||||
if not text.startswith(OAUTH_ACCOUNT_BLOCK_PREFIX):
|
||||
return False
|
||||
lowered = text[len(OAUTH_ACCOUNT_BLOCK_PREFIX) :].strip().lower()
|
||||
return any(keyword in lowered for keyword in _LEGACY_TOKEN_INVALID_KEYWORDS)
|
||||
|
||||
|
||||
def is_account_level_block(reason: str | None) -> bool:
|
||||
"""判断 oauth_invalid_reason 是否属于账号级别的 block(刷新 token 无法修复)。"""
|
||||
if not reason:
|
||||
return False
|
||||
text = str(reason)
|
||||
return text.startswith(
|
||||
OAUTH_ACCOUNT_BLOCK_PREFIX
|
||||
) and not _is_refresh_recoverable_account_block(text)
|
||||
|
||||
|
||||
async def verify_oauth_before_account_block(
|
||||
*,
|
||||
endpoint: Any,
|
||||
key: Any,
|
||||
candidate_reason: str,
|
||||
request_id: str | None = None,
|
||||
key_display: str | None = None,
|
||||
) -> bool:
|
||||
"""Before applying an account-level block, distinguish it from OAuth expiry."""
|
||||
display = key_display or str(getattr(key, "id", "?") or "?")
|
||||
try:
|
||||
from src.services.provider.auth import get_provider_auth
|
||||
|
||||
await get_provider_auth(endpoint, key, force_refresh=True, refresh_skew=0)
|
||||
except Exception as exc:
|
||||
logger.debug(
|
||||
"[OAUTH_VERIFY] [{}] {} account-block precheck failed: {}",
|
||||
request_id,
|
||||
display,
|
||||
exc,
|
||||
)
|
||||
|
||||
latest_reason = str(getattr(key, "oauth_invalid_reason", None) or "").strip()
|
||||
if latest_reason.startswith(OAUTH_EXPIRED_PREFIX) or latest_reason.startswith(
|
||||
OAUTH_REFRESH_FAILED_PREFIX
|
||||
):
|
||||
logger.info(
|
||||
"[OAUTH_VERIFY] [{}] {} candidate account block ({}) skipped due to {}",
|
||||
request_id,
|
||||
display,
|
||||
candidate_reason,
|
||||
latest_reason[:120],
|
||||
)
|
||||
return False
|
||||
|
||||
logger.debug(
|
||||
"[OAUTH_VERIFY] [{}] {} proceeding with account block ({}), post-refresh reason: {}",
|
||||
request_id,
|
||||
display,
|
||||
candidate_reason,
|
||||
latest_reason[:120] if latest_reason else "<none>",
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class OAuthAccessTokenResult:
|
||||
access_token: str
|
||||
decrypted_auth_config: dict[str, Any] | None
|
||||
refreshed: bool
|
||||
|
||||
|
||||
async def resolve_oauth_access_token(
|
||||
*,
|
||||
key_id: str,
|
||||
encrypted_api_key: str,
|
||||
encrypted_auth_config: str | None,
|
||||
provider_proxy_config: dict[str, Any] | None = None,
|
||||
endpoint_api_format: str | None = None,
|
||||
) -> OAuthAccessTokenResult:
|
||||
"""Resolve (and lazily refresh) OAuth access_token for a ProviderAPIKey.
|
||||
|
||||
This helper is safe to call from concurrent async tasks because it does not
|
||||
rely on the caller's SQLAlchemy Session:
|
||||
- It runs refresh logic without an ORM session.
|
||||
- If refresh succeeded (encrypted fields changed), it persists the new encrypted
|
||||
values to DB using a short, independent session.
|
||||
"""
|
||||
|
||||
# Local import to avoid circular imports during app startup.
|
||||
from src.services.provider.auth import get_provider_auth
|
||||
|
||||
# Build detached key-like objects for get_provider_auth().
|
||||
provider_obj = (
|
||||
SimpleNamespace(proxy=provider_proxy_config) if provider_proxy_config is not None else None
|
||||
)
|
||||
endpoint_obj = SimpleNamespace(api_format=str(endpoint_api_format or ""))
|
||||
key_obj = SimpleNamespace(
|
||||
id=str(key_id),
|
||||
auth_type="oauth",
|
||||
api_key=encrypted_api_key,
|
||||
auth_config=encrypted_auth_config,
|
||||
provider=provider_obj,
|
||||
)
|
||||
|
||||
orig_api_key = key_obj.api_key
|
||||
orig_auth_config = key_obj.auth_config
|
||||
|
||||
auth_info = await get_provider_auth(endpoint_obj, key_obj) # type: ignore[arg-type]
|
||||
if auth_info is None:
|
||||
# Should not happen for auth_type="oauth", but keep defensive.
|
||||
return OAuthAccessTokenResult(access_token="", decrypted_auth_config=None, refreshed=False)
|
||||
|
||||
access_token = str(auth_info.auth_value or "").removeprefix("Bearer ").strip()
|
||||
refreshed = (key_obj.api_key != orig_api_key) or (key_obj.auth_config != orig_auth_config)
|
||||
|
||||
if refreshed:
|
||||
# Persist refreshed token/config back to DB.
|
||||
try:
|
||||
with create_session() as db:
|
||||
row = db.query(ProviderAPIKey).filter(ProviderAPIKey.id == str(key_id)).first()
|
||||
if row is not None:
|
||||
row.api_key = key_obj.api_key
|
||||
row.auth_config = key_obj.auth_config
|
||||
# Refresh succeeded => only clear recoverable token errors.
|
||||
# True account-level blocks must be cleared explicitly.
|
||||
current_reason = str(getattr(row, "oauth_invalid_reason", None) or "")
|
||||
if row.oauth_invalid_at is not None and not is_account_level_block(
|
||||
current_reason
|
||||
):
|
||||
row.oauth_invalid_at = None
|
||||
row.oauth_invalid_reason = None
|
||||
db.commit()
|
||||
except Exception as e:
|
||||
# Don't fail caller path; token is still usable for this request.
|
||||
logger.debug("[OAUTH_REFRESH] persist refreshed token failed for key {}: {}", key_id, e)
|
||||
|
||||
return OAuthAccessTokenResult(
|
||||
access_token=access_token,
|
||||
decrypted_auth_config=auth_info.decrypted_auth_config,
|
||||
refreshed=refreshed,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"OAuthAccessTokenResult",
|
||||
"TOKEN_INVALIDATED_KEYWORDS",
|
||||
"verify_oauth_before_account_block",
|
||||
"looks_like_token_invalidated",
|
||||
"resolve_oauth_access_token",
|
||||
]
|
||||
20
_deprecated_py_src/services/provider/pool/__init__.py
Normal file
20
_deprecated_py_src/services/provider/pool/__init__.py
Normal file
@@ -0,0 +1,20 @@
|
||||
"""Generic Account Pool management for any Provider type.
|
||||
|
||||
Re-exports the main public API for convenience.
|
||||
"""
|
||||
|
||||
from src.services.provider.pool.config import (
|
||||
PoolConfig,
|
||||
ScoringWeights,
|
||||
UnschedulableRule,
|
||||
parse_pool_config,
|
||||
)
|
||||
from src.services.provider.pool.manager import PoolManager
|
||||
|
||||
__all__ = [
|
||||
"PoolConfig",
|
||||
"PoolManager",
|
||||
"ScoringWeights",
|
||||
"UnschedulableRule",
|
||||
"parse_pool_config",
|
||||
]
|
||||
594
_deprecated_py_src/services/provider/pool/account_state.py
Normal file
594
_deprecated_py_src/services/provider/pool/account_state.py
Normal file
@@ -0,0 +1,594 @@
|
||||
"""Pool account state helpers.
|
||||
|
||||
Provides a shared way to classify account-level hard-block states
|
||||
from upstream metadata and OAuth invalid reasons.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from src.services.provider_keys.quota_reader import get_quota_reader
|
||||
|
||||
OAUTH_ACCOUNT_BLOCK_PREFIX = "[ACCOUNT_BLOCK] "
|
||||
OAUTH_REFRESH_FAILED_PREFIX = "[REFRESH_FAILED] "
|
||||
OAUTH_EXPIRED_PREFIX = "[OAUTH_EXPIRED] "
|
||||
OAUTH_REQUEST_FAILED_PREFIX = "[REQUEST_FAILED] "
|
||||
|
||||
# -- 按原因细分的关键词组 --
|
||||
# 封禁类 (suspended / banned)
|
||||
_KEYWORDS_SUSPENDED: tuple[str, ...] = (
|
||||
"suspended",
|
||||
"account_block",
|
||||
"account blocked",
|
||||
"封禁",
|
||||
"封号",
|
||||
"被封",
|
||||
"账户已封禁",
|
||||
"账号异常",
|
||||
)
|
||||
|
||||
# 停用类 (disabled / deactivated)
|
||||
_KEYWORDS_DISABLED: tuple[str, ...] = (
|
||||
"account has been disabled",
|
||||
"account disabled",
|
||||
"account has been deactivated",
|
||||
"account_deactivated",
|
||||
"account deactivated",
|
||||
"organization has been disabled",
|
||||
"organization_disabled",
|
||||
"deactivated_workspace",
|
||||
"deactivated",
|
||||
"访问被禁止",
|
||||
"账户访问被禁止",
|
||||
)
|
||||
|
||||
_TOKEN_INVALID_KEYWORDS: tuple[str, ...] = (
|
||||
"authentication token has been invalidated",
|
||||
"token has been invalidated",
|
||||
"codex token 无效或已过期",
|
||||
)
|
||||
|
||||
# 需要验证类
|
||||
_KEYWORDS_VERIFICATION: tuple[str, ...] = (
|
||||
"validation_required",
|
||||
"verify your account",
|
||||
"需要验证",
|
||||
"验证账号",
|
||||
"验证身份",
|
||||
)
|
||||
|
||||
# 合并的完整列表(用于 is_account_level_block_reason 快速判断)
|
||||
ACCOUNT_BLOCK_REASON_KEYWORDS: tuple[str, ...] = (
|
||||
*_KEYWORDS_SUSPENDED,
|
||||
*_KEYWORDS_DISABLED,
|
||||
*_TOKEN_INVALID_KEYWORDS,
|
||||
*_KEYWORDS_VERIFICATION,
|
||||
)
|
||||
|
||||
AUTO_REMOVABLE_ACCOUNT_STATE_CODES: frozenset[str] = frozenset(
|
||||
{
|
||||
"account_banned",
|
||||
"account_suspended",
|
||||
"account_disabled",
|
||||
"workspace_deactivated",
|
||||
"account_forbidden",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _classify_block_reason(text: str) -> tuple[str, str]:
|
||||
"""Return (code, label) based on the oauth_invalid_reason text."""
|
||||
lowered = text.lower()
|
||||
if any(kw in lowered for kw in _TOKEN_INVALID_KEYWORDS):
|
||||
return "oauth_expired", "Token 失效"
|
||||
if any(kw in lowered for kw in _KEYWORDS_VERIFICATION):
|
||||
return "account_verification", "需要验证"
|
||||
if "deactivated_workspace" in lowered:
|
||||
return "workspace_deactivated", "工作区停用"
|
||||
if any(kw in lowered for kw in _KEYWORDS_DISABLED):
|
||||
return "account_disabled", "账号停用"
|
||||
if any(kw in lowered for kw in _KEYWORDS_SUSPENDED):
|
||||
return "account_suspended", "账号封禁"
|
||||
return "account_blocked", "账号异常"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PoolAccountState:
|
||||
"""Resolved account-level state for one key."""
|
||||
|
||||
blocked: bool
|
||||
code: str | None = None # account_banned / account_forbidden / account_blocked
|
||||
label: str | None = None
|
||||
reason: str | None = None
|
||||
source: str | None = None # metadata / oauth_invalid / oauth_refresh / oauth_request
|
||||
recoverable: bool = False
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class OAuthStatusSnapshot:
|
||||
code: str = "none" # none / valid / expiring / expired / invalid / check_failed
|
||||
label: str | None = None
|
||||
reason: str | None = None
|
||||
expires_at: int | None = None
|
||||
invalid_at: int | None = None
|
||||
source: str | None = None
|
||||
requires_reauth: bool = False
|
||||
expiring_soon: bool = False
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AccountStatusSnapshot:
|
||||
code: str = "ok"
|
||||
label: str | None = None
|
||||
reason: str | None = None
|
||||
blocked: bool = False
|
||||
source: str | None = None
|
||||
recoverable: bool = False
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class QuotaStatusSnapshot:
|
||||
code: str = "unknown" # unknown / ok / exhausted
|
||||
label: str | None = None
|
||||
reason: str | None = None
|
||||
exhausted: bool = False
|
||||
usage_ratio: float | None = None
|
||||
updated_at: int | None = None
|
||||
reset_seconds: float | None = None
|
||||
plan_type: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ProviderKeyStatusSnapshot:
|
||||
oauth: OAuthStatusSnapshot
|
||||
account: AccountStatusSnapshot
|
||||
quota: QuotaStatusSnapshot
|
||||
|
||||
|
||||
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 _clean_text(value: Any) -> str | None:
|
||||
if not isinstance(value, str):
|
||||
return None
|
||||
text = value.strip()
|
||||
return text or None
|
||||
|
||||
|
||||
def _extract_reason(source: dict[str, Any] | None, *fields: str) -> str | None:
|
||||
if not isinstance(source, dict):
|
||||
return None
|
||||
for field in fields:
|
||||
text = _clean_text(source.get(field))
|
||||
if text:
|
||||
return text
|
||||
return None
|
||||
|
||||
|
||||
def _is_workspace_deactivated_reason(reason: str | None) -> bool:
|
||||
text = _clean_text(reason)
|
||||
return bool(text and "deactivated_workspace" in text.lower())
|
||||
|
||||
|
||||
_TAGGED_REASON_PATTERN = re.compile(
|
||||
r"(?:^|\n)\[(?P<tag>[A-Z_]+)\]\s*(?P<detail>.*?)(?=\n\[[A-Z_]+\]|\Z)",
|
||||
re.S,
|
||||
)
|
||||
|
||||
|
||||
def _extract_tagged_reason_sections(reason: str | None) -> dict[str, str]:
|
||||
text = _clean_text(reason)
|
||||
if not text:
|
||||
return {}
|
||||
sections: dict[str, str] = {}
|
||||
for match in _TAGGED_REASON_PATTERN.finditer(text):
|
||||
tag = str(match.group("tag") or "").strip().upper()
|
||||
if not tag or tag in sections:
|
||||
continue
|
||||
detail = str(match.group("detail") or "").strip()
|
||||
sections[tag] = detail
|
||||
return sections
|
||||
|
||||
|
||||
def _resolve_from_metadata(
|
||||
provider_type: str | None,
|
||||
upstream_metadata: Any,
|
||||
) -> PoolAccountState | None:
|
||||
if not isinstance(upstream_metadata, dict):
|
||||
return None
|
||||
|
||||
normalized_provider = str(provider_type or "").strip().lower()
|
||||
provider_bucket: dict[str, Any] | None = None
|
||||
if normalized_provider:
|
||||
maybe_bucket = upstream_metadata.get(normalized_provider)
|
||||
if isinstance(maybe_bucket, dict):
|
||||
provider_bucket = maybe_bucket
|
||||
|
||||
quota_block = get_quota_reader(normalized_provider, upstream_metadata).account_block()
|
||||
if quota_block.blocked:
|
||||
return PoolAccountState(
|
||||
blocked=True,
|
||||
code=quota_block.code,
|
||||
label=quota_block.label,
|
||||
reason=quota_block.reason,
|
||||
source="metadata",
|
||||
)
|
||||
|
||||
for source in (provider_bucket, upstream_metadata):
|
||||
if not isinstance(source, dict):
|
||||
continue
|
||||
if _is_truthy_flag(source.get("is_banned")):
|
||||
reason = _extract_reason(source, "ban_reason", "forbidden_reason", "reason", "message")
|
||||
return PoolAccountState(
|
||||
blocked=True,
|
||||
code="account_banned",
|
||||
label="账号封禁",
|
||||
reason=reason or "账号已封禁",
|
||||
source="metadata",
|
||||
)
|
||||
if _is_truthy_flag(source.get("is_forbidden")) or _is_truthy_flag(
|
||||
source.get("account_disabled")
|
||||
):
|
||||
reason = _extract_reason(source, "forbidden_reason", "ban_reason", "reason", "message")
|
||||
if _is_workspace_deactivated_reason(reason):
|
||||
return PoolAccountState(
|
||||
blocked=True,
|
||||
code="workspace_deactivated",
|
||||
label="工作区停用",
|
||||
reason=reason or "工作区已停用",
|
||||
source="metadata",
|
||||
)
|
||||
return PoolAccountState(
|
||||
blocked=True,
|
||||
code="account_forbidden",
|
||||
label="访问受限",
|
||||
reason=reason or "账号访问受限",
|
||||
source="metadata",
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _resolve_from_oauth_invalid_reason(reason: str | None) -> PoolAccountState | None:
|
||||
text = _clean_text(reason)
|
||||
if not text:
|
||||
return None
|
||||
|
||||
if text.startswith(OAUTH_ACCOUNT_BLOCK_PREFIX):
|
||||
cleaned = text[len(OAUTH_ACCOUNT_BLOCK_PREFIX) :].strip()
|
||||
code, label = (
|
||||
_classify_block_reason(cleaned) if cleaned else ("account_blocked", "账号异常")
|
||||
)
|
||||
return PoolAccountState(
|
||||
blocked=True,
|
||||
code=code,
|
||||
label=label,
|
||||
reason=cleaned or "账号异常",
|
||||
source="oauth_invalid",
|
||||
)
|
||||
|
||||
if text.startswith(OAUTH_EXPIRED_PREFIX):
|
||||
cleaned = text[len(OAUTH_EXPIRED_PREFIX) :].strip()
|
||||
return PoolAccountState(
|
||||
blocked=True,
|
||||
code="oauth_expired",
|
||||
label="Token 失效",
|
||||
reason=cleaned or "OAuth Token 已过期且无法续期",
|
||||
source="oauth_invalid",
|
||||
recoverable=True,
|
||||
)
|
||||
|
||||
if text.startswith(OAUTH_REFRESH_FAILED_PREFIX):
|
||||
cleaned = text[len(OAUTH_REFRESH_FAILED_PREFIX) :].strip()
|
||||
return PoolAccountState(
|
||||
blocked=False,
|
||||
code="oauth_refresh_failed",
|
||||
label="续期失败",
|
||||
reason=cleaned or "OAuth Token 续期失败",
|
||||
source="oauth_refresh",
|
||||
recoverable=True,
|
||||
)
|
||||
|
||||
if text.startswith(OAUTH_REQUEST_FAILED_PREFIX):
|
||||
cleaned = text[len(OAUTH_REQUEST_FAILED_PREFIX) :].strip()
|
||||
return PoolAccountState(
|
||||
blocked=False,
|
||||
code="oauth_request_failed",
|
||||
label="请求失败",
|
||||
reason=cleaned or "账号状态检查失败",
|
||||
source="oauth_request",
|
||||
recoverable=True,
|
||||
)
|
||||
|
||||
if text.startswith("["):
|
||||
return None
|
||||
|
||||
lowered = text.lower()
|
||||
if any(keyword in lowered for keyword in ACCOUNT_BLOCK_REASON_KEYWORDS):
|
||||
code, label = _classify_block_reason(text)
|
||||
return PoolAccountState(
|
||||
blocked=True,
|
||||
code=code,
|
||||
label=label,
|
||||
reason=text,
|
||||
source="oauth_invalid",
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def resolve_account_status_snapshot(
|
||||
*,
|
||||
provider_type: str | None,
|
||||
upstream_metadata: Any,
|
||||
oauth_invalid_reason: str | None,
|
||||
) -> AccountStatusSnapshot:
|
||||
from_metadata = _resolve_from_metadata(provider_type, upstream_metadata)
|
||||
if from_metadata is not None:
|
||||
return AccountStatusSnapshot(
|
||||
code=from_metadata.code or "ok",
|
||||
label=from_metadata.label,
|
||||
reason=from_metadata.reason,
|
||||
blocked=from_metadata.blocked,
|
||||
source=from_metadata.source,
|
||||
recoverable=from_metadata.recoverable,
|
||||
)
|
||||
|
||||
text = _clean_text(oauth_invalid_reason)
|
||||
if not text:
|
||||
return AccountStatusSnapshot()
|
||||
|
||||
tagged_sections = _extract_tagged_reason_sections(text)
|
||||
if "ACCOUNT_BLOCK" in tagged_sections:
|
||||
cleaned = tagged_sections["ACCOUNT_BLOCK"]
|
||||
code, label = (
|
||||
_classify_block_reason(cleaned) if cleaned else ("account_blocked", "账号异常")
|
||||
)
|
||||
return AccountStatusSnapshot(
|
||||
code=code,
|
||||
label=label,
|
||||
reason=cleaned or "账号异常",
|
||||
blocked=True,
|
||||
source="oauth_invalid",
|
||||
)
|
||||
|
||||
if text.startswith("["):
|
||||
return AccountStatusSnapshot()
|
||||
|
||||
lowered = text.lower()
|
||||
if any(keyword in lowered for keyword in ACCOUNT_BLOCK_REASON_KEYWORDS):
|
||||
code, label = _classify_block_reason(text)
|
||||
return AccountStatusSnapshot(
|
||||
code=code,
|
||||
label=label,
|
||||
reason=text,
|
||||
blocked=True,
|
||||
source="oauth_invalid",
|
||||
)
|
||||
|
||||
return AccountStatusSnapshot()
|
||||
|
||||
|
||||
def resolve_oauth_status_snapshot(
|
||||
*,
|
||||
auth_type: str | None,
|
||||
oauth_expires_at: int | None,
|
||||
oauth_invalid_at: int | None,
|
||||
oauth_invalid_reason: str | None,
|
||||
now_ts: int | None = None,
|
||||
) -> OAuthStatusSnapshot:
|
||||
if str(auth_type or "").strip().lower() != "oauth":
|
||||
return OAuthStatusSnapshot()
|
||||
|
||||
now = int(now_ts if now_ts is not None else time.time())
|
||||
invalid_at = int(oauth_invalid_at) if isinstance(oauth_invalid_at, int) else None
|
||||
tagged_sections = _extract_tagged_reason_sections(oauth_invalid_reason)
|
||||
raw_reason = _clean_text(oauth_invalid_reason)
|
||||
|
||||
expired_reason = tagged_sections.get("OAUTH_EXPIRED")
|
||||
if expired_reason:
|
||||
return OAuthStatusSnapshot(
|
||||
code="invalid",
|
||||
label="已失效",
|
||||
reason=expired_reason,
|
||||
invalid_at=invalid_at,
|
||||
expires_at=oauth_expires_at,
|
||||
source="oauth_invalid",
|
||||
requires_reauth=True,
|
||||
)
|
||||
|
||||
refresh_failed_reason = tagged_sections.get("REFRESH_FAILED")
|
||||
if refresh_failed_reason:
|
||||
return OAuthStatusSnapshot(
|
||||
code="invalid",
|
||||
label="已失效",
|
||||
reason=refresh_failed_reason,
|
||||
invalid_at=invalid_at,
|
||||
expires_at=oauth_expires_at,
|
||||
source="oauth_refresh",
|
||||
requires_reauth=True,
|
||||
)
|
||||
|
||||
request_failed_reason = tagged_sections.get("REQUEST_FAILED")
|
||||
if request_failed_reason:
|
||||
return OAuthStatusSnapshot(
|
||||
code="check_failed",
|
||||
label="检查失败",
|
||||
reason=request_failed_reason,
|
||||
expires_at=oauth_expires_at,
|
||||
source="oauth_request",
|
||||
)
|
||||
|
||||
account_snapshot = resolve_account_status_snapshot(
|
||||
provider_type=None,
|
||||
upstream_metadata=None,
|
||||
oauth_invalid_reason=raw_reason,
|
||||
)
|
||||
if account_snapshot.blocked:
|
||||
if oauth_expires_at is None:
|
||||
return OAuthStatusSnapshot()
|
||||
elif raw_reason or invalid_at is not None:
|
||||
return OAuthStatusSnapshot(
|
||||
code="invalid",
|
||||
label="已失效",
|
||||
reason=raw_reason,
|
||||
invalid_at=invalid_at,
|
||||
expires_at=oauth_expires_at,
|
||||
source="oauth_invalid",
|
||||
requires_reauth=True,
|
||||
)
|
||||
|
||||
expires_at = int(oauth_expires_at) if isinstance(oauth_expires_at, int) else None
|
||||
if expires_at is None:
|
||||
return OAuthStatusSnapshot()
|
||||
if expires_at <= now:
|
||||
return OAuthStatusSnapshot(
|
||||
code="expired",
|
||||
label="已过期",
|
||||
reason="Token 已过期,请重新授权",
|
||||
expires_at=expires_at,
|
||||
source="expires_at",
|
||||
requires_reauth=True,
|
||||
)
|
||||
expiring_soon = (expires_at - now) < 24 * 3600
|
||||
return OAuthStatusSnapshot(
|
||||
code="expiring" if expiring_soon else "valid",
|
||||
label="即将过期" if expiring_soon else "有效",
|
||||
expires_at=expires_at,
|
||||
source="expires_at",
|
||||
expiring_soon=expiring_soon,
|
||||
)
|
||||
|
||||
|
||||
def resolve_quota_status_snapshot(
|
||||
*,
|
||||
provider_type: str | None,
|
||||
upstream_metadata: Any,
|
||||
) -> QuotaStatusSnapshot:
|
||||
normalized_provider = str(provider_type or "").strip().lower()
|
||||
reader = get_quota_reader(normalized_provider, upstream_metadata)
|
||||
quota_state = reader.is_exhausted()
|
||||
usage_ratio = reader.usage_ratio()
|
||||
updated_at = reader.updated_at()
|
||||
reset_seconds = reader.reset_seconds()
|
||||
plan_type = reader.plan_type()
|
||||
|
||||
if quota_state.exhausted:
|
||||
return QuotaStatusSnapshot(
|
||||
code="exhausted",
|
||||
label="额度耗尽",
|
||||
reason=quota_state.reason,
|
||||
exhausted=True,
|
||||
usage_ratio=usage_ratio,
|
||||
updated_at=updated_at,
|
||||
reset_seconds=reset_seconds,
|
||||
plan_type=plan_type,
|
||||
)
|
||||
|
||||
if any(value is not None for value in (usage_ratio, updated_at, reset_seconds, plan_type)):
|
||||
return QuotaStatusSnapshot(
|
||||
code="ok",
|
||||
exhausted=False,
|
||||
usage_ratio=usage_ratio,
|
||||
updated_at=updated_at,
|
||||
reset_seconds=reset_seconds,
|
||||
plan_type=plan_type,
|
||||
)
|
||||
|
||||
return QuotaStatusSnapshot()
|
||||
|
||||
|
||||
def build_provider_key_status_snapshot(
|
||||
*,
|
||||
auth_type: str | None,
|
||||
oauth_expires_at: int | None,
|
||||
oauth_invalid_at: int | None,
|
||||
oauth_invalid_reason: str | None,
|
||||
provider_type: str | None,
|
||||
upstream_metadata: Any,
|
||||
now_ts: int | None = None,
|
||||
) -> ProviderKeyStatusSnapshot:
|
||||
account = resolve_account_status_snapshot(
|
||||
provider_type=provider_type,
|
||||
upstream_metadata=upstream_metadata,
|
||||
oauth_invalid_reason=oauth_invalid_reason,
|
||||
)
|
||||
oauth = resolve_oauth_status_snapshot(
|
||||
auth_type=auth_type,
|
||||
oauth_expires_at=oauth_expires_at,
|
||||
oauth_invalid_at=oauth_invalid_at,
|
||||
oauth_invalid_reason=oauth_invalid_reason,
|
||||
now_ts=now_ts,
|
||||
)
|
||||
quota = resolve_quota_status_snapshot(
|
||||
provider_type=provider_type,
|
||||
upstream_metadata=upstream_metadata,
|
||||
)
|
||||
return ProviderKeyStatusSnapshot(oauth=oauth, account=account, quota=quota)
|
||||
|
||||
|
||||
def resolve_pool_account_state(
|
||||
*,
|
||||
provider_type: str | None,
|
||||
upstream_metadata: Any,
|
||||
oauth_invalid_reason: str | None,
|
||||
) -> PoolAccountState:
|
||||
"""Resolve account-level hard-block state for pool scheduling."""
|
||||
|
||||
from_metadata = _resolve_from_metadata(provider_type, upstream_metadata)
|
||||
if from_metadata is not None:
|
||||
return from_metadata
|
||||
|
||||
from_oauth = _resolve_from_oauth_invalid_reason(oauth_invalid_reason)
|
||||
if from_oauth is not None:
|
||||
return from_oauth
|
||||
|
||||
return PoolAccountState(blocked=False)
|
||||
|
||||
|
||||
def should_auto_remove_account_state(state: PoolAccountState) -> bool:
|
||||
"""Whether a resolved account state is safe to auto-remove.
|
||||
|
||||
Auto-removal is limited to hard, non-recoverable account abnormalities.
|
||||
Pure token failures (`oauth_expired`, `oauth_refresh_failed`) and
|
||||
softer/manual-recoverable states like `account_verification` are excluded.
|
||||
"""
|
||||
|
||||
return bool(
|
||||
state.blocked
|
||||
and not state.recoverable
|
||||
and str(state.code or "").strip().lower() in AUTO_REMOVABLE_ACCOUNT_STATE_CODES
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ACCOUNT_BLOCK_REASON_KEYWORDS",
|
||||
"AUTO_REMOVABLE_ACCOUNT_STATE_CODES",
|
||||
"AccountStatusSnapshot",
|
||||
"OAUTH_ACCOUNT_BLOCK_PREFIX",
|
||||
"OAUTH_EXPIRED_PREFIX",
|
||||
"OAUTH_REFRESH_FAILED_PREFIX",
|
||||
"OAUTH_REQUEST_FAILED_PREFIX",
|
||||
"OAuthStatusSnapshot",
|
||||
"PoolAccountState",
|
||||
"ProviderKeyStatusSnapshot",
|
||||
"QuotaStatusSnapshot",
|
||||
"build_provider_key_status_snapshot",
|
||||
"resolve_account_status_snapshot",
|
||||
"resolve_oauth_status_snapshot",
|
||||
"resolve_pool_account_state",
|
||||
"resolve_quota_status_snapshot",
|
||||
"should_auto_remove_account_state",
|
||||
]
|
||||
341
_deprecated_py_src/services/provider/pool/config.py
Normal file
341
_deprecated_py_src/services/provider/pool/config.py
Normal file
@@ -0,0 +1,341 @@
|
||||
"""Account Pool configuration (provider-agnostic)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.services.provider.pool.dimensions import get_preset_dimension, get_preset_names
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ScoringWeights:
|
||||
"""Weights used by multi-score scheduling."""
|
||||
|
||||
lru: float = 0.3
|
||||
latency: float = 0.25
|
||||
health: float = 0.2
|
||||
cost_remaining: float = 0.25
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SchedulingPreset:
|
||||
"""Single scheduling preset item with enable/disable and optional sub-config."""
|
||||
|
||||
preset: str
|
||||
enabled: bool = True
|
||||
mode: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UnschedulableRule:
|
||||
"""Keyword-based temporary unschedule rule."""
|
||||
|
||||
keyword: str
|
||||
duration_minutes: int = 5
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PoolConfig:
|
||||
"""Parsed pool configuration for any Provider.
|
||||
|
||||
All transient state lives in Redis; this dataclass only holds
|
||||
the *configuration* that controls pool behaviour.
|
||||
"""
|
||||
|
||||
# -- Sticky Session -------------------------------------------------------
|
||||
sticky_session_ttl_seconds: int = 3600 # 1 hour
|
||||
# Key 优先模式下号池整体优先级(None 时回退 provider_priority)
|
||||
global_priority: int | None = None
|
||||
|
||||
# -- Load-Aware Selection -------------------------------------------------
|
||||
load_threshold_percent: int = 80
|
||||
|
||||
# -- Scheduling (unified preset list) -------------------------------------
|
||||
scheduling_presets: tuple[SchedulingPreset, ...] = (
|
||||
SchedulingPreset(preset="cache_affinity", enabled=True),
|
||||
)
|
||||
# Derived from scheduling_presets at parse time (backward compat for consumers)
|
||||
lru_enabled: bool = True
|
||||
scheduling_mode: str = "lru" # lru | multi_score
|
||||
|
||||
scoring_weights: ScoringWeights = field(default_factory=ScoringWeights)
|
||||
latency_window_seconds: int = 3600
|
||||
latency_sample_limit: int = 50
|
||||
|
||||
# -- Rolling-Window Cost Tracking -----------------------------------------
|
||||
cost_window_seconds: int = 18000 # 5 hours
|
||||
cost_limit_per_key_tokens: int | None = None # None = unlimited
|
||||
cost_soft_threshold_percent: int = 80
|
||||
|
||||
# -- Cooldown Defaults ----------------------------------------------------
|
||||
rate_limit_cooldown_seconds: int = 300 # 429
|
||||
overload_cooldown_seconds: int = 30 # 529
|
||||
|
||||
# -- OAuth Proactive Refresh ----------------------------------------------
|
||||
proactive_refresh_seconds: int = 180 # 3 minutes before expiry
|
||||
|
||||
# -- Health Policy --------------------------------------------------------
|
||||
health_policy_enabled: bool = True
|
||||
|
||||
# -- Temporary Unschedulable Rules ----------------------------------------
|
||||
unschedulable_rules: list[UnschedulableRule] = field(default_factory=list)
|
||||
|
||||
# -- Batch Operations -----------------------------------------------------
|
||||
batch_concurrency: int = 8
|
||||
|
||||
# -- Quota Probing --------------------------------------------------------
|
||||
probing_enabled: bool = False
|
||||
probing_interval_minutes: int = 10
|
||||
auto_remove_banned_keys: bool = False
|
||||
|
||||
# -- Stream Timeout Auto-Pause --------------------------------------------
|
||||
stream_timeout_threshold: int = 3 # N timeouts within window trigger cooldown
|
||||
stream_timeout_window_seconds: int = 1800 # 30 min counting window
|
||||
stream_timeout_cooldown_seconds: int = 300 # 5 min cooldown
|
||||
|
||||
# -- Pluggable Strategies -------------------------------------------------
|
||||
strategies: tuple[str, ...] = ()
|
||||
|
||||
|
||||
def parse_pool_config(provider_config: Any) -> PoolConfig | None:
|
||||
"""Parse PoolConfig from ``Provider.config``.
|
||||
|
||||
Only looks for the explicit ``pool_advanced`` key. Returns ``None``
|
||||
when the provider has no pool section configured, meaning the caller
|
||||
should use the normal (non-pool) scheduling path.
|
||||
"""
|
||||
config_dict = provider_config if isinstance(provider_config, dict) else {}
|
||||
|
||||
raw_advanced = config_dict.get("pool_advanced")
|
||||
if raw_advanced is None:
|
||||
return None
|
||||
|
||||
if not isinstance(raw_advanced, dict):
|
||||
# Could be a pre-validated Pydantic model; grab its dict.
|
||||
try:
|
||||
raw_advanced = raw_advanced.model_dump() # type: ignore[union-attr]
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"PoolConfig: advanced config type invalid ({}), falling back to defaults",
|
||||
type(raw_advanced).__name__,
|
||||
)
|
||||
return PoolConfig()
|
||||
|
||||
rules: list[UnschedulableRule] = []
|
||||
raw_rules = raw_advanced.get("unschedulable_rules")
|
||||
if isinstance(raw_rules, list):
|
||||
for r in raw_rules:
|
||||
if isinstance(r, dict) and isinstance(r.get("keyword"), str):
|
||||
rules.append(
|
||||
UnschedulableRule(
|
||||
keyword=r["keyword"],
|
||||
duration_minutes=int(r.get("duration_minutes", 5)),
|
||||
)
|
||||
)
|
||||
|
||||
def _int_or(key: str, default: int) -> int:
|
||||
v = raw_advanced.get(key)
|
||||
if v is None:
|
||||
return default
|
||||
try:
|
||||
return int(v)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
def _bool_or(key: str, default: bool) -> bool:
|
||||
v = raw_advanced.get(key)
|
||||
if v is None:
|
||||
return default
|
||||
return bool(v)
|
||||
|
||||
def _opt_int(key: str) -> int | None:
|
||||
v = raw_advanced.get(key)
|
||||
if v is None:
|
||||
return None
|
||||
try:
|
||||
return int(v)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
scoring_weights = _parse_scoring_weights(raw_advanced.get("scoring_weights"))
|
||||
|
||||
# Parse scheduling presets (new object-list format or legacy string-list)
|
||||
presets = _parse_scheduling_presets_v2(
|
||||
raw_advanced.get("scheduling_presets"),
|
||||
legacy_mode=raw_advanced.get("scheduling_mode"),
|
||||
legacy_lru=raw_advanced.get("lru_enabled"),
|
||||
)
|
||||
|
||||
# Derive scheduling_mode and lru_enabled from the presets list
|
||||
enabled = [p for p in presets if p.enabled]
|
||||
lru_enabled = any(p.preset == "lru" for p in enabled)
|
||||
non_lru_enabled = [p for p in enabled if p.preset != "lru"]
|
||||
scheduling_mode = "multi_score" if non_lru_enabled else "lru"
|
||||
|
||||
strategies = list(_parse_strategies(raw_advanced.get("strategies")))
|
||||
if scheduling_mode == "multi_score" and "multi_score" not in strategies:
|
||||
strategies.append("multi_score")
|
||||
|
||||
return PoolConfig(
|
||||
sticky_session_ttl_seconds=_int_or("sticky_session_ttl_seconds", 3600),
|
||||
global_priority=_opt_int("global_priority"),
|
||||
load_threshold_percent=_int_or("load_threshold_percent", 80),
|
||||
scheduling_presets=presets,
|
||||
lru_enabled=lru_enabled,
|
||||
scheduling_mode=scheduling_mode,
|
||||
scoring_weights=scoring_weights,
|
||||
latency_window_seconds=_int_or("latency_window_seconds", 3600),
|
||||
latency_sample_limit=_int_or("latency_sample_limit", 50),
|
||||
cost_window_seconds=_int_or("cost_window_seconds", 18000),
|
||||
cost_limit_per_key_tokens=_opt_int("cost_limit_per_key_tokens"),
|
||||
cost_soft_threshold_percent=_int_or("cost_soft_threshold_percent", 80),
|
||||
rate_limit_cooldown_seconds=_int_or("rate_limit_cooldown_seconds", 300),
|
||||
overload_cooldown_seconds=_int_or("overload_cooldown_seconds", 30),
|
||||
proactive_refresh_seconds=_int_or("proactive_refresh_seconds", 180),
|
||||
health_policy_enabled=_bool_or("health_policy_enabled", True),
|
||||
unschedulable_rules=rules,
|
||||
batch_concurrency=max(1, min(_int_or("batch_concurrency", 8), 32)),
|
||||
probing_enabled=_bool_or("probing_enabled", False),
|
||||
probing_interval_minutes=max(1, min(_int_or("probing_interval_minutes", 10), 1440)),
|
||||
auto_remove_banned_keys=_bool_or("auto_remove_banned_keys", False),
|
||||
stream_timeout_threshold=_int_or("stream_timeout_threshold", 3),
|
||||
stream_timeout_window_seconds=_int_or("stream_timeout_window_seconds", 1800),
|
||||
stream_timeout_cooldown_seconds=_int_or("stream_timeout_cooldown_seconds", 300),
|
||||
strategies=tuple(strategies),
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Internal parsers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _allowed_preset_names() -> set[str]:
|
||||
return get_preset_names() | {"lru"}
|
||||
|
||||
|
||||
def _get_preset_mode_meta(preset_name: str) -> tuple[tuple[str, ...], str | None]:
|
||||
dim = get_preset_dimension(preset_name)
|
||||
if dim is None or not dim.modes:
|
||||
return (), None
|
||||
|
||||
modes = tuple(str(mode).strip().lower() for mode in dim.modes if str(mode).strip())
|
||||
if not modes:
|
||||
return (), None
|
||||
|
||||
raw_default = str(dim.default_mode or "").strip().lower()
|
||||
default_mode = raw_default if raw_default in modes else modes[0]
|
||||
return modes, default_mode
|
||||
|
||||
|
||||
def _parse_strategies(raw: Any) -> tuple[str, ...]:
|
||||
"""Parse strategy names from config (list[str] -> tuple[str, ...])."""
|
||||
if not isinstance(raw, list):
|
||||
return ()
|
||||
return tuple(str(s) for s in raw if isinstance(s, str) and s)
|
||||
|
||||
|
||||
def _parse_scoring_weights(raw: Any) -> ScoringWeights:
|
||||
"""Parse scoring weights with graceful fallback."""
|
||||
if not isinstance(raw, dict):
|
||||
return ScoringWeights()
|
||||
|
||||
def _float_or(value: Any, default: float) -> float:
|
||||
try:
|
||||
parsed = float(value)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
return max(0.0, min(parsed, 1.0))
|
||||
|
||||
return ScoringWeights(
|
||||
lru=_float_or(raw.get("lru"), 0.3),
|
||||
latency=_float_or(raw.get("latency"), 0.25),
|
||||
health=_float_or(raw.get("health"), 0.2),
|
||||
cost_remaining=_float_or(raw.get("cost_remaining"), 0.25),
|
||||
)
|
||||
|
||||
|
||||
def _parse_scheduling_presets_v2(
|
||||
raw: Any,
|
||||
*,
|
||||
legacy_mode: Any = None,
|
||||
legacy_lru: Any = None,
|
||||
) -> tuple[SchedulingPreset, ...]:
|
||||
"""Parse scheduling presets, supporting both new and legacy formats.
|
||||
|
||||
New format::
|
||||
|
||||
[{"preset": "lru", "enabled": true},
|
||||
{"preset": "free_team_first", "enabled": true, "mode": "free_only"},
|
||||
...]
|
||||
|
||||
Legacy format::
|
||||
|
||||
["free_team_first", "recent_refresh"] (with separate scheduling_mode / lru_enabled)
|
||||
"""
|
||||
if isinstance(raw, list) and raw:
|
||||
first = raw[0]
|
||||
if isinstance(first, dict):
|
||||
return _parse_preset_object_list(raw)
|
||||
if isinstance(first, str):
|
||||
return _convert_legacy_string_list(raw, legacy_mode, legacy_lru)
|
||||
|
||||
# No presets at all: derive from legacy fields
|
||||
return _build_from_legacy_fields(legacy_mode, legacy_lru)
|
||||
|
||||
|
||||
def _parse_preset_object_list(raw: list[Any]) -> tuple[SchedulingPreset, ...]:
|
||||
"""Parse new-format object list into SchedulingPreset tuple."""
|
||||
allowed = _allowed_preset_names()
|
||||
ordered: list[SchedulingPreset] = []
|
||||
seen: set[str] = set()
|
||||
for item in raw:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
name = str(item.get("preset", "")).strip().lower()
|
||||
if name not in allowed or name in seen:
|
||||
continue
|
||||
seen.add(name)
|
||||
enabled = bool(item.get("enabled", True))
|
||||
mode: str | None = None
|
||||
modes, default_mode = _get_preset_mode_meta(name)
|
||||
if modes:
|
||||
raw_mode = str(item.get("mode", default_mode) or "").strip().lower()
|
||||
mode = raw_mode if raw_mode in modes else default_mode
|
||||
ordered.append(SchedulingPreset(preset=name, enabled=enabled, mode=mode))
|
||||
return tuple(ordered) if ordered else (SchedulingPreset(preset="lru", enabled=True),)
|
||||
|
||||
|
||||
def _convert_legacy_string_list(
|
||||
raw: list[Any],
|
||||
legacy_mode: Any,
|
||||
legacy_lru: Any,
|
||||
) -> tuple[SchedulingPreset, ...]:
|
||||
"""Convert legacy string list + mode/lru fields to new format."""
|
||||
lru_enabled = legacy_lru if isinstance(legacy_lru, bool) else True
|
||||
|
||||
allowed_non_lru = _allowed_preset_names() - {"lru"}
|
||||
items: list[SchedulingPreset] = [SchedulingPreset(preset="lru", enabled=lru_enabled)]
|
||||
seen: set[str] = {"lru"}
|
||||
for p in raw:
|
||||
if not isinstance(p, str):
|
||||
continue
|
||||
name = p.strip().lower()
|
||||
if name not in allowed_non_lru or name in seen:
|
||||
continue
|
||||
seen.add(name)
|
||||
items.append(SchedulingPreset(preset=name, enabled=True))
|
||||
return tuple(items)
|
||||
|
||||
|
||||
def _build_from_legacy_fields(legacy_mode: Any, legacy_lru: Any) -> tuple[SchedulingPreset, ...]:
|
||||
"""Build presets from legacy scheduling_mode / lru_enabled only."""
|
||||
# Explicit lru_enabled=True -> LRU; explicit lru_enabled=False -> cache_affinity.
|
||||
# No legacy fields at all -> default to cache_affinity.
|
||||
if isinstance(legacy_lru, bool):
|
||||
if legacy_lru:
|
||||
return (SchedulingPreset(preset="lru", enabled=True),)
|
||||
return (SchedulingPreset(preset="cache_affinity", enabled=True),)
|
||||
return (SchedulingPreset(preset="cache_affinity", enabled=True),)
|
||||
66
_deprecated_py_src/services/provider/pool/cost_tracker.py
Normal file
66
_deprecated_py_src/services/provider/pool/cost_tracker.py
Normal file
@@ -0,0 +1,66 @@
|
||||
"""Rolling-window cost tracking for the Account Pool.
|
||||
|
||||
Each key has a configurable token budget per rolling window (e.g. 5 hours).
|
||||
When the budget is exhausted the key is marked as unschedulable by the pool
|
||||
manager. A "soft threshold" (default 80 %) causes the pool to *prefer*
|
||||
other keys but still allows traffic if no alternatives exist.
|
||||
|
||||
All state is stored in Redis sorted sets via :mod:`redis_ops`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from src.services.provider.pool import redis_ops
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from src.services.provider.pool.config import PoolConfig
|
||||
|
||||
|
||||
async def record_usage(
|
||||
provider_id: str,
|
||||
key_id: str,
|
||||
tokens: int,
|
||||
config: PoolConfig,
|
||||
) -> None:
|
||||
"""Record *tokens* used by *key_id* in the rolling cost window."""
|
||||
if tokens <= 0:
|
||||
return
|
||||
if config.cost_limit_per_key_tokens is None:
|
||||
return # cost tracking disabled
|
||||
await redis_ops.add_cost_entry(provider_id, key_id, tokens, config.cost_window_seconds)
|
||||
|
||||
|
||||
async def get_window_usage(
|
||||
provider_id: str,
|
||||
key_id: str,
|
||||
config: PoolConfig,
|
||||
) -> int:
|
||||
"""Return total tokens used by *key_id* within the current window."""
|
||||
return await redis_ops.get_cost_window_total(provider_id, key_id, config.cost_window_seconds)
|
||||
|
||||
|
||||
async def is_at_limit(
|
||||
provider_id: str,
|
||||
key_id: str,
|
||||
config: PoolConfig,
|
||||
) -> bool:
|
||||
"""Return ``True`` if the key has exhausted its budget."""
|
||||
if config.cost_limit_per_key_tokens is None:
|
||||
return False
|
||||
total = await get_window_usage(provider_id, key_id, config)
|
||||
return total >= config.cost_limit_per_key_tokens
|
||||
|
||||
|
||||
async def is_approaching_limit(
|
||||
provider_id: str,
|
||||
key_id: str,
|
||||
config: PoolConfig,
|
||||
) -> bool:
|
||||
"""Return ``True`` if the key is above the soft threshold."""
|
||||
if config.cost_limit_per_key_tokens is None:
|
||||
return False
|
||||
total = await get_window_usage(provider_id, key_id, config)
|
||||
threshold = config.cost_limit_per_key_tokens * config.cost_soft_threshold_percent / 100
|
||||
return total >= threshold
|
||||
@@ -0,0 +1,39 @@
|
||||
"""Pool scheduling preset dimensions.
|
||||
|
||||
Importing this package registers all built-in preset dimensions.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from . import cache_affinity # noqa: F401
|
||||
from . import cost_first # noqa: F401
|
||||
from . import free_first # noqa: F401
|
||||
from . import free_team_first # noqa: F401
|
||||
from . import health_first # noqa: F401
|
||||
from . import latency_first # noqa: F401
|
||||
from . import load_balance # noqa: F401
|
||||
from . import plus_first # noqa: F401
|
||||
from . import priority_first # noqa: F401
|
||||
from . import quota_balanced # noqa: F401
|
||||
from . import recent_refresh # noqa: F401
|
||||
from . import single_account # noqa: F401
|
||||
from . import team_first # noqa: F401
|
||||
from .registry import (
|
||||
PresetDimensionBase,
|
||||
PresetDimensionMeta,
|
||||
get_all_preset_dimensions,
|
||||
get_preset_dimension,
|
||||
get_preset_dimension_metas,
|
||||
get_preset_names,
|
||||
register_preset_dimension,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"PresetDimensionBase",
|
||||
"PresetDimensionMeta",
|
||||
"get_all_preset_dimensions",
|
||||
"get_preset_dimension",
|
||||
"get_preset_dimension_metas",
|
||||
"get_preset_names",
|
||||
"register_preset_dimension",
|
||||
]
|
||||
284
_deprecated_py_src/services/provider/pool/dimensions/_helpers.py
Normal file
284
_deprecated_py_src/services/provider/pool/dimensions/_helpers.py
Normal file
@@ -0,0 +1,284 @@
|
||||
"""Shared helpers for pool preset dimensions."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from src.core.provider_types import ProviderType, normalize_provider_type
|
||||
from src.services.provider_keys.quota_reader import get_quota_reader
|
||||
|
||||
|
||||
def safe_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 safe_metadata(key_obj: Any) -> dict[str, Any]:
|
||||
raw = getattr(key_obj, "upstream_metadata", None)
|
||||
return raw if isinstance(raw, dict) else {}
|
||||
|
||||
|
||||
def normalize_plan(value: Any) -> str | None:
|
||||
if not isinstance(value, str):
|
||||
return None
|
||||
normalized = value.strip().lower()
|
||||
return normalized or None
|
||||
|
||||
|
||||
def rank_ascending(key_id: str, scores: dict[str, float], all_ids: list[str]) -> float:
|
||||
"""Rank score within all IDs; lower value means better rank."""
|
||||
|
||||
if not all_ids:
|
||||
return 0.0
|
||||
|
||||
valid_count = sum(1 for kid in all_ids if safe_float(scores.get(kid)) is not None)
|
||||
if valid_count <= 0:
|
||||
return 0.5
|
||||
|
||||
decorated: list[tuple[int, float, int, str]] = []
|
||||
for idx, kid in enumerate(all_ids):
|
||||
score_raw = safe_float(scores.get(kid))
|
||||
if score_raw is None:
|
||||
decorated.append((1, float("inf"), idx, kid))
|
||||
else:
|
||||
decorated.append((0, score_raw, idx, kid))
|
||||
|
||||
decorated.sort(key=lambda item: (item[0], item[1], item[2]))
|
||||
rank_idx = 0
|
||||
for idx, (_missing, _value, _order, kid) in enumerate(decorated):
|
||||
if kid == key_id:
|
||||
rank_idx = idx
|
||||
break
|
||||
|
||||
n = len(all_ids)
|
||||
if n <= 1:
|
||||
return 0.0
|
||||
return rank_idx / float(n - 1)
|
||||
|
||||
|
||||
def rank_descending(key_id: str, scores: dict[str, float], all_ids: list[str]) -> float:
|
||||
"""Rank score within all IDs; higher value means better rank."""
|
||||
|
||||
if not all_ids:
|
||||
return 0.0
|
||||
|
||||
valid_count = sum(1 for kid in all_ids if safe_float(scores.get(kid)) is not None)
|
||||
if valid_count <= 0:
|
||||
return 0.5
|
||||
|
||||
decorated: list[tuple[int, float, int, str]] = []
|
||||
for idx, kid in enumerate(all_ids):
|
||||
score_raw = safe_float(scores.get(kid))
|
||||
if score_raw is None:
|
||||
decorated.append((1, float("inf"), idx, kid))
|
||||
else:
|
||||
# 排序时取负值,使分值越大排名越靠前(rank 越小)
|
||||
decorated.append((0, -score_raw, idx, kid))
|
||||
|
||||
decorated.sort(key=lambda item: (item[0], item[1], item[2]))
|
||||
rank_idx = 0
|
||||
for idx, (_missing, _value, _order, kid) in enumerate(decorated):
|
||||
if kid == key_id:
|
||||
rank_idx = idx
|
||||
break
|
||||
|
||||
n = len(all_ids)
|
||||
if n <= 1:
|
||||
return 0.0
|
||||
return rank_idx / float(n - 1)
|
||||
|
||||
|
||||
def extract_plan_type(key_obj: Any) -> str | None:
|
||||
direct = normalize_plan(getattr(key_obj, "oauth_plan_type", None))
|
||||
if direct:
|
||||
return direct
|
||||
|
||||
metadata = safe_metadata(key_obj)
|
||||
for provider_type in (ProviderType.CODEX, ProviderType.KIRO, ProviderType.ANTIGRAVITY):
|
||||
plan_type = get_quota_reader(provider_type, metadata).plan_type()
|
||||
if plan_type:
|
||||
return plan_type
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _resolve_key_provider_type(key_obj: Any, provider_type: str | None = None) -> str | None:
|
||||
explicit = normalize_provider_type(provider_type)
|
||||
if explicit:
|
||||
return explicit
|
||||
|
||||
direct = normalize_provider_type(getattr(key_obj, "provider_type", None))
|
||||
if direct:
|
||||
return direct
|
||||
|
||||
provider = getattr(key_obj, "provider", None)
|
||||
related = normalize_provider_type(getattr(provider, "provider_type", None))
|
||||
if related:
|
||||
return related
|
||||
|
||||
metadata = safe_metadata(key_obj)
|
||||
candidates = [
|
||||
provider.value
|
||||
for provider in (ProviderType.CODEX, ProviderType.KIRO, ProviderType.ANTIGRAVITY)
|
||||
if isinstance(metadata.get(provider.value), dict)
|
||||
]
|
||||
if len(candidates) == 1:
|
||||
return candidates[0]
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _extract_codex_weekly_reset_seconds(metadata: dict[str, Any]) -> float | None:
|
||||
codex = metadata.get(ProviderType.CODEX.value)
|
||||
if not isinstance(codex, dict):
|
||||
return None
|
||||
|
||||
weekly_used_percent = safe_float(codex.get("primary_used_percent"))
|
||||
if weekly_used_percent is not None:
|
||||
clamped_used = max(0.0, min(weekly_used_percent, 100.0))
|
||||
if clamped_used <= 1e-6:
|
||||
# 周额度仍为满额时,不启用周窗口重置倒计时。
|
||||
return None
|
||||
|
||||
now = time.time()
|
||||
|
||||
# 优先绝对时间戳,避免 reset_seconds 快照随时间漂移。
|
||||
reset_at = safe_float(codex.get("primary_reset_at"))
|
||||
if reset_at is not None and reset_at > 0:
|
||||
remaining = reset_at - now
|
||||
return remaining if remaining > 0 else 0.0
|
||||
|
||||
reset_seconds = safe_float(codex.get("primary_reset_seconds"))
|
||||
if reset_seconds is None or reset_seconds < 0:
|
||||
return None
|
||||
|
||||
updated_at = safe_float(codex.get("updated_at"))
|
||||
if updated_at is not None and updated_at > 0:
|
||||
# 时钟偏移下 updated_at 可能晚于当前时间,elapsed 需要下限钳制到 0。
|
||||
elapsed = max(now - updated_at, 0.0)
|
||||
return max(reset_seconds - elapsed, 0.0)
|
||||
|
||||
return reset_seconds
|
||||
|
||||
|
||||
def extract_reset_seconds(key_obj: Any, provider_type: str | None = None) -> float | None:
|
||||
metadata = safe_metadata(key_obj)
|
||||
resolved_provider_type = _resolve_key_provider_type(key_obj, provider_type)
|
||||
|
||||
if resolved_provider_type == ProviderType.CODEX:
|
||||
# Codex metadata 已统一约定:primary_* 表示周限额,secondary_* 表示 5H 限额。
|
||||
return _extract_codex_weekly_reset_seconds(metadata)
|
||||
|
||||
if resolved_provider_type in (ProviderType.KIRO, ProviderType.ANTIGRAVITY):
|
||||
return get_quota_reader(resolved_provider_type, metadata).reset_seconds()
|
||||
|
||||
candidates: list[float] = []
|
||||
|
||||
for provider_type in (ProviderType.CODEX, ProviderType.KIRO, ProviderType.ANTIGRAVITY):
|
||||
reset_seconds = get_quota_reader(provider_type, metadata).reset_seconds()
|
||||
if reset_seconds is None:
|
||||
continue
|
||||
candidates.append(reset_seconds)
|
||||
|
||||
if not candidates:
|
||||
return None
|
||||
return min(candidates)
|
||||
|
||||
|
||||
def extract_usage_ratio(key_obj: Any) -> float | None:
|
||||
metadata = safe_metadata(key_obj)
|
||||
|
||||
for provider_type in (ProviderType.CODEX, ProviderType.KIRO, ProviderType.ANTIGRAVITY):
|
||||
usage_ratio = get_quota_reader(provider_type, metadata).usage_ratio()
|
||||
if usage_ratio is not None:
|
||||
return usage_ratio
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def extract_internal_priority(key_obj: Any) -> int:
|
||||
raw = getattr(key_obj, "internal_priority", None)
|
||||
parsed = safe_float(raw)
|
||||
if parsed is None:
|
||||
return 999999
|
||||
return max(0, int(parsed))
|
||||
|
||||
|
||||
def extract_health_score(key_obj: Any) -> float | None:
|
||||
direct = safe_float(getattr(key_obj, "health_score", None))
|
||||
if direct is not None:
|
||||
return max(0.0, min(direct, 1.0))
|
||||
|
||||
health_by_format = getattr(key_obj, "health_by_format", None)
|
||||
if not isinstance(health_by_format, dict) or not health_by_format:
|
||||
return None
|
||||
|
||||
scores: list[float] = []
|
||||
for payload in health_by_format.values():
|
||||
if not isinstance(payload, dict):
|
||||
continue
|
||||
score = safe_float(payload.get("health_score"))
|
||||
if score is None:
|
||||
continue
|
||||
scores.append(max(0.0, min(score, 1.0)))
|
||||
|
||||
if not scores:
|
||||
return None
|
||||
return min(scores)
|
||||
|
||||
|
||||
def plan_priority_score(plan_type: str | None, mode: str | None = None) -> float:
|
||||
"""Score a key based on plan type and scheduling mode.
|
||||
|
||||
Lower score = higher priority.
|
||||
"""
|
||||
|
||||
effective_mode = (mode or "both").strip().lower()
|
||||
if effective_mode == "free_only":
|
||||
if plan_type == "free":
|
||||
return 0.0
|
||||
if plan_type == "team":
|
||||
return 0.5
|
||||
elif effective_mode == "team_only":
|
||||
if plan_type == "team":
|
||||
return 0.0
|
||||
if plan_type == "free":
|
||||
return 0.5
|
||||
elif effective_mode == "plus_only":
|
||||
if plan_type in {"plus", "pro"}:
|
||||
return 0.0
|
||||
if plan_type in {"enterprise", "business"}:
|
||||
return 0.3
|
||||
else:
|
||||
# "both" or unrecognized -> original behavior
|
||||
if plan_type in {"free", "team"}:
|
||||
return 0.0
|
||||
if plan_type in {"enterprise", "business"}:
|
||||
return 0.2
|
||||
if plan_type in {"plus", "pro"}:
|
||||
return 0.6
|
||||
if plan_type:
|
||||
return 0.7
|
||||
return 0.8
|
||||
|
||||
|
||||
__all__ = [
|
||||
"extract_health_score",
|
||||
"extract_internal_priority",
|
||||
"extract_plan_type",
|
||||
"extract_reset_seconds",
|
||||
"extract_usage_ratio",
|
||||
"normalize_plan",
|
||||
"plan_priority_score",
|
||||
"rank_ascending",
|
||||
"rank_descending",
|
||||
"safe_float",
|
||||
"safe_metadata",
|
||||
]
|
||||
@@ -0,0 +1,45 @@
|
||||
"""cache_affinity preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import rank_descending
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class CacheAffinityDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "cache_affinity"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "缓存亲和"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "优先复用最近使用过的 Key,利用 Prompt Caching"
|
||||
|
||||
@property
|
||||
def mutex_group(self) -> str | None:
|
||||
return "distribution_mode"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "依据 LRU 时间戳(最近使用优先,与 LRU 轮转相反)"
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
return rank_descending(key_id, lru_scores, all_key_ids)
|
||||
|
||||
|
||||
register_preset_dimension(CacheAffinityDimension())
|
||||
@@ -0,0 +1,62 @@
|
||||
"""cost_first preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import extract_usage_ratio, rank_ascending, safe_float
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class CostFirstDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "cost_first"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "成本优先"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "优先选择窗口消耗更低的账号"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "依据窗口成本/Token 用量,缺失时回退配额使用率"
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
cost_totals = context.get("cost_totals")
|
||||
if not isinstance(cost_totals, dict):
|
||||
cost_totals = {}
|
||||
cost_limit = safe_float(context.get("cost_limit_per_key_tokens"))
|
||||
|
||||
cost_scores: dict[str, float] = {}
|
||||
for kid in all_key_ids:
|
||||
used = safe_float(cost_totals.get(kid))
|
||||
if used is not None and used >= 0:
|
||||
if cost_limit is not None and cost_limit > 0:
|
||||
cost_scores[kid] = max(0.0, min(used / cost_limit, 1.0))
|
||||
else:
|
||||
cost_scores[kid] = min(1.0, used / (used + 10000.0))
|
||||
continue
|
||||
|
||||
usage_ratio = extract_usage_ratio(keys_by_id.get(kid))
|
||||
if usage_ratio is not None:
|
||||
cost_scores[kid] = usage_ratio
|
||||
|
||||
if not cost_scores:
|
||||
return rank_ascending(key_id, lru_scores, all_key_ids)
|
||||
return rank_ascending(key_id, cost_scores, all_key_ids)
|
||||
|
||||
|
||||
register_preset_dimension(CostFirstDimension())
|
||||
@@ -0,0 +1,51 @@
|
||||
"""free_first preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import extract_plan_type, plan_priority_score, rank_ascending
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class FreeFirstDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "free_first"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "Free 优先"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "优先消耗 Free 账号(依赖 plan_type)"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "依据 plan_type(Free 账号优先调度)"
|
||||
|
||||
@property
|
||||
def providers(self) -> tuple[str, ...]:
|
||||
return ("codex", "kiro")
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
plan_scores = {
|
||||
kid: plan_priority_score(extract_plan_type(keys_by_id.get(kid)), "free_only")
|
||||
for kid in all_key_ids
|
||||
}
|
||||
if len(set(plan_scores.values())) <= 1:
|
||||
return rank_ascending(key_id, lru_scores, all_key_ids)
|
||||
return rank_ascending(key_id, plan_scores, all_key_ids)
|
||||
|
||||
|
||||
register_preset_dimension(FreeFirstDimension())
|
||||
@@ -0,0 +1,63 @@
|
||||
"""free_team_first preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import extract_plan_type, plan_priority_score, rank_ascending
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class FreeTeamFirstDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "free_team_first"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "Free/Team 优先"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "优先消耗低档账号(依赖 plan_type)"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "依据 plan_type(oauth_plan_type 或 upstream_metadata)"
|
||||
|
||||
@property
|
||||
def providers(self) -> tuple[str, ...]:
|
||||
return ("codex", "kiro")
|
||||
|
||||
@property
|
||||
def modes(self) -> tuple[str, ...] | None:
|
||||
return ("free_only", "team_only", "both")
|
||||
|
||||
@property
|
||||
def default_mode(self) -> str | None:
|
||||
return "both"
|
||||
|
||||
@property
|
||||
def hidden(self) -> bool:
|
||||
return True
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
plan_scores = {
|
||||
kid: plan_priority_score(extract_plan_type(keys_by_id.get(kid)), mode)
|
||||
for kid in all_key_ids
|
||||
}
|
||||
if len(set(plan_scores.values())) <= 1:
|
||||
return rank_ascending(key_id, lru_scores, all_key_ids)
|
||||
return rank_ascending(key_id, plan_scores, all_key_ids)
|
||||
|
||||
|
||||
register_preset_dimension(FreeTeamFirstDimension())
|
||||
@@ -0,0 +1,57 @@
|
||||
"""health_first preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import extract_health_score, rank_ascending, safe_float
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class HealthFirstDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "health_first"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "健康优先"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "优先选择健康分更高、失败更少的账号"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "依据 health_by_format 聚合分(含熔断/失败衰减)"
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
health_scores_ctx = context.get("health_scores")
|
||||
if not isinstance(health_scores_ctx, dict):
|
||||
health_scores_ctx = {}
|
||||
|
||||
penalty_scores: dict[str, float] = {}
|
||||
for kid in all_key_ids:
|
||||
score = safe_float(health_scores_ctx.get(kid))
|
||||
if score is None:
|
||||
score = extract_health_score(keys_by_id.get(kid))
|
||||
if score is None:
|
||||
continue
|
||||
normalized = max(0.0, min(score, 1.0))
|
||||
penalty_scores[kid] = 1.0 - normalized
|
||||
|
||||
if not penalty_scores:
|
||||
return rank_ascending(key_id, lru_scores, all_key_ids)
|
||||
return rank_ascending(key_id, penalty_scores, all_key_ids)
|
||||
|
||||
|
||||
register_preset_dimension(HealthFirstDimension())
|
||||
@@ -0,0 +1,54 @@
|
||||
"""latency_first preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import rank_ascending, safe_float
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class LatencyFirstDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "latency_first"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "延迟优先"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "优先选择最近延迟更低的账号"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "依据号池延迟窗口均值(latency_window_seconds)"
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
latency_avgs = context.get("latency_avgs")
|
||||
if not isinstance(latency_avgs, dict):
|
||||
latency_avgs = {}
|
||||
|
||||
latency_scores: dict[str, float] = {}
|
||||
for kid in all_key_ids:
|
||||
latency = safe_float(latency_avgs.get(kid))
|
||||
if latency is None or latency < 0:
|
||||
continue
|
||||
latency_scores[kid] = latency
|
||||
|
||||
if not latency_scores:
|
||||
return rank_ascending(key_id, lru_scores, all_key_ids)
|
||||
return rank_ascending(key_id, latency_scores, all_key_ids)
|
||||
|
||||
|
||||
register_preset_dimension(LatencyFirstDimension())
|
||||
@@ -0,0 +1,45 @@
|
||||
"""load_balance preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
from typing import Any
|
||||
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class LoadBalanceDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "load_balance"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "负载均衡"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "随机分散 Key 使用,均匀分摊负载"
|
||||
|
||||
@property
|
||||
def mutex_group(self) -> str | None:
|
||||
return "distribution_mode"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "每次随机分值,实现完全均匀分散"
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
return random.random()
|
||||
|
||||
|
||||
register_preset_dimension(LoadBalanceDimension())
|
||||
@@ -0,0 +1,51 @@
|
||||
"""plus_first preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import extract_plan_type, plan_priority_score, rank_ascending
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class PlusFirstDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "plus_first"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "Plus 优先"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "优先消耗 Plus/Pro 账号(依赖 plan_type)"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "依据 plan_type(Plus/Pro 账号优先调度)"
|
||||
|
||||
@property
|
||||
def providers(self) -> tuple[str, ...]:
|
||||
return ("codex", "kiro")
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
plan_scores = {
|
||||
kid: plan_priority_score(extract_plan_type(keys_by_id.get(kid)), "plus_only")
|
||||
for kid in all_key_ids
|
||||
}
|
||||
if len(set(plan_scores.values())) <= 1:
|
||||
return rank_ascending(key_id, lru_scores, all_key_ids)
|
||||
return rank_ascending(key_id, plan_scores, all_key_ids)
|
||||
|
||||
|
||||
register_preset_dimension(PlusFirstDimension())
|
||||
@@ -0,0 +1,46 @@
|
||||
"""priority_first preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import extract_internal_priority, rank_ascending
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class PriorityFirstDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "priority_first"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "优先级优先"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "按账号优先级顺序调度(数字越小越优先)"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "依据 internal_priority(支持拖拽/手工编辑)"
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
priority_scores = {
|
||||
kid: float(extract_internal_priority(keys_by_id.get(kid))) for kid in all_key_ids
|
||||
}
|
||||
if len(set(priority_scores.values())) <= 1:
|
||||
return rank_ascending(key_id, lru_scores, all_key_ids)
|
||||
return rank_ascending(key_id, priority_scores, all_key_ids)
|
||||
|
||||
|
||||
register_preset_dimension(PriorityFirstDimension())
|
||||
@@ -0,0 +1,61 @@
|
||||
"""quota_balanced preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import extract_usage_ratio, rank_ascending, safe_float
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class QuotaBalancedDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "quota_balanced"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "额度平均"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "优先选额度消耗最少的账号"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "依据账号配额使用率;无配额时回退到窗口成本使用"
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
usage_scores: dict[str, float] = {}
|
||||
cost_totals = context.get("cost_totals")
|
||||
if not isinstance(cost_totals, dict):
|
||||
cost_totals = {}
|
||||
cost_limit = safe_float(context.get("cost_limit_per_key_tokens"))
|
||||
for kid in all_key_ids:
|
||||
key_obj = keys_by_id.get(kid)
|
||||
usage_ratio = extract_usage_ratio(key_obj)
|
||||
if usage_ratio is None:
|
||||
used = safe_float(cost_totals.get(kid))
|
||||
if used is not None and used >= 0:
|
||||
if cost_limit is not None and cost_limit > 0:
|
||||
usage_ratio = max(0.0, min(used / cost_limit, 1.0))
|
||||
else:
|
||||
# 无明确上限时用 log 归一化,确保维度仍有区分能力。
|
||||
usage_ratio = min(1.0, used / (used + 10000.0))
|
||||
if usage_ratio is not None:
|
||||
usage_scores[kid] = usage_ratio
|
||||
if not usage_scores:
|
||||
return rank_ascending(key_id, lru_scores, all_key_ids)
|
||||
return rank_ascending(key_id, usage_scores, all_key_ids)
|
||||
|
||||
|
||||
register_preset_dimension(QuotaBalancedDimension())
|
||||
@@ -0,0 +1,53 @@
|
||||
"""recent_refresh preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import extract_reset_seconds, rank_ascending
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class RecentRefreshDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "recent_refresh"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "额度刷新优先"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "优先选即将刷新额度的账号"
|
||||
|
||||
@property
|
||||
def providers(self) -> tuple[str, ...]:
|
||||
return ("codex", "kiro")
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "依据账号额度重置倒计时(next_reset / reset_seconds)"
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
provider_type = context.get("provider_type")
|
||||
reset_scores: dict[str, float] = {}
|
||||
for kid in all_key_ids:
|
||||
reset_seconds = extract_reset_seconds(keys_by_id.get(kid), provider_type=provider_type)
|
||||
if reset_seconds is not None:
|
||||
reset_scores[kid] = reset_seconds
|
||||
if not reset_scores:
|
||||
return rank_ascending(key_id, lru_scores, all_key_ids)
|
||||
return rank_ascending(key_id, reset_scores, all_key_ids)
|
||||
|
||||
|
||||
register_preset_dimension(RecentRefreshDimension())
|
||||
265
_deprecated_py_src/services/provider/pool/dimensions/registry.py
Normal file
265
_deprecated_py_src/services/provider/pool/dimensions/registry.py
Normal file
@@ -0,0 +1,265 @@
|
||||
"""Preset dimension registry for pool multi-score scheduling."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from threading import RLock
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PresetDimensionMeta:
|
||||
"""Serializable metadata for one preset dimension."""
|
||||
|
||||
name: str
|
||||
label: str
|
||||
description: str
|
||||
providers: tuple[str, ...]
|
||||
modes: tuple[str, ...] | None
|
||||
default_mode: str | None
|
||||
mutex_group: str | None
|
||||
evidence_hint: str | None
|
||||
|
||||
|
||||
class PresetDimensionBase(ABC):
|
||||
"""Base class of one pool scheduling preset dimension."""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def name(self) -> str:
|
||||
"""Stable preset key, e.g. ``free_team_first``."""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def label(self) -> str:
|
||||
"""User-facing label."""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def description(self) -> str:
|
||||
"""User-facing description."""
|
||||
|
||||
@property
|
||||
def providers(self) -> tuple[str, ...]:
|
||||
"""Supported provider types.
|
||||
|
||||
Empty tuple means the dimension is universal and applies to all providers.
|
||||
"""
|
||||
|
||||
return ()
|
||||
|
||||
@property
|
||||
def modes(self) -> tuple[str, ...] | None:
|
||||
"""Optional sub-modes for this dimension."""
|
||||
|
||||
return None
|
||||
|
||||
@property
|
||||
def default_mode(self) -> str | None:
|
||||
"""Default mode when mode is omitted."""
|
||||
|
||||
return None
|
||||
|
||||
@property
|
||||
def mutex_group(self) -> str | None:
|
||||
"""Optional mutual-exclusion group key.
|
||||
|
||||
Presets in the same group are expected to be mutually exclusive in UI.
|
||||
"""
|
||||
|
||||
return None
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
"""Human-readable hint about which data this preset uses."""
|
||||
|
||||
return None
|
||||
|
||||
@property
|
||||
def hidden(self) -> bool:
|
||||
"""If True, this dimension is excluded from API metadata listings.
|
||||
|
||||
The dimension remains functional for backward compatibility but
|
||||
will not appear in the scheduling dialog.
|
||||
"""
|
||||
|
||||
return False
|
||||
|
||||
@abstractmethod
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
"""Compute normalized metric in [0, 1], lower is better."""
|
||||
|
||||
def is_applicable(self, provider_type: str) -> bool:
|
||||
"""Return whether this dimension applies to the given provider type."""
|
||||
|
||||
if not self.providers:
|
||||
return True
|
||||
normalized = _normalize_name(provider_type)
|
||||
return normalized in self.providers
|
||||
|
||||
|
||||
def _normalize_name(value: Any) -> str:
|
||||
if not isinstance(value, str):
|
||||
return ""
|
||||
return value.strip().lower()
|
||||
|
||||
|
||||
def _normalize_names(values: tuple[str, ...] | list[str]) -> tuple[str, ...]:
|
||||
normalized = [_normalize_name(item) for item in values]
|
||||
return tuple(item for item in normalized if item)
|
||||
|
||||
|
||||
_registry_lock = RLock()
|
||||
_registry: dict[str, PresetDimensionBase] = {}
|
||||
|
||||
|
||||
def register_preset_dimension(dim: PresetDimensionBase) -> None:
|
||||
"""Register or replace one preset dimension by name."""
|
||||
|
||||
name = _normalize_name(dim.name)
|
||||
if not name:
|
||||
raise ValueError("preset dimension name must be a non-empty string")
|
||||
|
||||
providers = _normalize_names(dim.providers)
|
||||
modes = _normalize_names(dim.modes or ())
|
||||
default_mode = _normalize_name(dim.default_mode)
|
||||
|
||||
if modes and default_mode and default_mode not in modes:
|
||||
raise ValueError(f"default_mode must be one of modes for preset '{name}'")
|
||||
|
||||
class _NormalizedDimension(PresetDimensionBase):
|
||||
# Lightweight wrapper to keep normalized metadata while preserving compute logic.
|
||||
def __init__(self, wrapped: PresetDimensionBase) -> None:
|
||||
self._wrapped = wrapped
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return name
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return self._wrapped.label
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return self._wrapped.description
|
||||
|
||||
@property
|
||||
def providers(self) -> tuple[str, ...]:
|
||||
return providers
|
||||
|
||||
@property
|
||||
def modes(self) -> tuple[str, ...] | None:
|
||||
return modes or None
|
||||
|
||||
@property
|
||||
def default_mode(self) -> str | None:
|
||||
if not modes:
|
||||
return None
|
||||
if default_mode:
|
||||
return default_mode
|
||||
return modes[0]
|
||||
|
||||
@property
|
||||
def mutex_group(self) -> str | None:
|
||||
raw = _normalize_name(self._wrapped.mutex_group)
|
||||
return raw or None
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
raw = str(self._wrapped.evidence_hint or "").strip()
|
||||
return raw or None
|
||||
|
||||
@property
|
||||
def hidden(self) -> bool:
|
||||
return self._wrapped.hidden
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
return self._wrapped.compute_metric(
|
||||
key_id=key_id,
|
||||
all_key_ids=all_key_ids,
|
||||
keys_by_id=keys_by_id,
|
||||
lru_scores=lru_scores,
|
||||
context=context,
|
||||
mode=mode,
|
||||
)
|
||||
|
||||
normalized = _NormalizedDimension(dim)
|
||||
with _registry_lock:
|
||||
_registry[name] = normalized
|
||||
|
||||
|
||||
def get_preset_dimension(name: str) -> PresetDimensionBase | None:
|
||||
"""Get one registered preset dimension by name."""
|
||||
|
||||
key = _normalize_name(name)
|
||||
if not key:
|
||||
return None
|
||||
with _registry_lock:
|
||||
return _registry.get(key)
|
||||
|
||||
|
||||
def get_all_preset_dimensions() -> list[PresetDimensionBase]:
|
||||
"""Get all registered preset dimensions in registration order."""
|
||||
|
||||
with _registry_lock:
|
||||
return list(_registry.values())
|
||||
|
||||
|
||||
def get_preset_names() -> set[str]:
|
||||
"""Get all registered preset names."""
|
||||
|
||||
with _registry_lock:
|
||||
return set(_registry.keys())
|
||||
|
||||
|
||||
def get_preset_dimension_metas() -> list[PresetDimensionMeta]:
|
||||
"""Get serializable metadata for all preset dimensions."""
|
||||
|
||||
metas: list[PresetDimensionMeta] = []
|
||||
for dim in get_all_preset_dimensions():
|
||||
if dim.hidden:
|
||||
continue
|
||||
metas.append(
|
||||
PresetDimensionMeta(
|
||||
name=dim.name,
|
||||
label=dim.label,
|
||||
description=dim.description,
|
||||
providers=dim.providers,
|
||||
modes=dim.modes,
|
||||
default_mode=dim.default_mode,
|
||||
mutex_group=dim.mutex_group,
|
||||
evidence_hint=dim.evidence_hint,
|
||||
)
|
||||
)
|
||||
return metas
|
||||
|
||||
|
||||
__all__ = [
|
||||
"PresetDimensionBase",
|
||||
"PresetDimensionMeta",
|
||||
"get_all_preset_dimensions",
|
||||
"get_preset_dimension",
|
||||
"get_preset_dimension_metas",
|
||||
"get_preset_names",
|
||||
"register_preset_dimension",
|
||||
]
|
||||
@@ -0,0 +1,51 @@
|
||||
"""single_account preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import extract_internal_priority, rank_ascending, rank_descending
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class SingleAccountDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "single_account"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "单号优先"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "集中使用同一账号(反向 LRU)"
|
||||
|
||||
@property
|
||||
def mutex_group(self) -> str | None:
|
||||
return "distribution_mode"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "先按账号优先级(internal_priority),同级再按反向 LRU 集中"
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
priority_scores = {
|
||||
kid: float(extract_internal_priority(keys_by_id.get(kid))) for kid in all_key_ids
|
||||
}
|
||||
priority_rank = rank_ascending(key_id, priority_scores, all_key_ids)
|
||||
lru_concentrate_rank = rank_descending(key_id, lru_scores, all_key_ids)
|
||||
# 强化“单号优先”的可控性:优先级优先,反向 LRU 作为次级聚合。
|
||||
return max(0.0, min(priority_rank * 0.75 + lru_concentrate_rank * 0.25, 1.0))
|
||||
|
||||
|
||||
register_preset_dimension(SingleAccountDimension())
|
||||
@@ -0,0 +1,51 @@
|
||||
"""team_first preset dimension."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ._helpers import extract_plan_type, plan_priority_score, rank_ascending
|
||||
from .registry import PresetDimensionBase, register_preset_dimension
|
||||
|
||||
|
||||
class TeamFirstDimension(PresetDimensionBase):
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "team_first"
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
return "Team 优先"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "优先消耗 Team 账号(依赖 plan_type)"
|
||||
|
||||
@property
|
||||
def evidence_hint(self) -> str | None:
|
||||
return "依据 plan_type(Team 账号优先调度)"
|
||||
|
||||
@property
|
||||
def providers(self) -> tuple[str, ...]:
|
||||
return ("codex", "kiro")
|
||||
|
||||
def compute_metric(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
keys_by_id: dict[str, Any],
|
||||
lru_scores: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
mode: str | None,
|
||||
) -> float:
|
||||
plan_scores = {
|
||||
kid: plan_priority_score(extract_plan_type(keys_by_id.get(kid)), "team_only")
|
||||
for kid in all_key_ids
|
||||
}
|
||||
if len(set(plan_scores.values())) <= 1:
|
||||
return rank_ascending(key_id, lru_scores, all_key_ids)
|
||||
return rank_ascending(key_id, plan_scores, all_key_ids)
|
||||
|
||||
|
||||
register_preset_dimension(TeamFirstDimension())
|
||||
104
_deprecated_py_src/services/provider/pool/health_cache.py
Normal file
104
_deprecated_py_src/services/provider/pool/health_cache.py
Normal file
@@ -0,0 +1,104 @@
|
||||
"""In-process pool health score cache.
|
||||
|
||||
This cache avoids recomputing per-key health aggregation on every request.
|
||||
It does not replace persistent health storage; source data still comes from
|
||||
``ProviderAPIKey.health_by_format`` carried on key objects.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
_TTL_SECONDS = 30.0
|
||||
_MAX_PROVIDERS = 500 # 最大缓存 provider 数量,防止无界增长
|
||||
_LOCK = threading.Lock()
|
||||
_CACHE: dict[str, tuple[float, dict[str, float]]] = {}
|
||||
|
||||
|
||||
def aggregate_health_score(health_by_format: Any) -> float:
|
||||
"""Aggregate health score from ``health_by_format`` (lower-bound strategy)."""
|
||||
if not isinstance(health_by_format, dict) or not health_by_format:
|
||||
return 1.0
|
||||
scores: list[float] = []
|
||||
for item in health_by_format.values():
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
try:
|
||||
score = float(item.get("health_score") or 1.0)
|
||||
except (TypeError, ValueError):
|
||||
score = 1.0
|
||||
scores.append(max(0.0, min(score, 1.0)))
|
||||
if not scores:
|
||||
return 1.0
|
||||
return min(scores)
|
||||
|
||||
|
||||
def get_health_scores(provider_id: str, keys: list[Any]) -> dict[str, float]:
|
||||
"""Return key health scores with per-provider TTL cache.
|
||||
|
||||
Uses incremental merge: if the cache is still valid but missing some keys,
|
||||
only the missing keys are computed and merged into the existing cache entry.
|
||||
"""
|
||||
now = time.monotonic()
|
||||
keys_by_id: dict[str, Any] = {}
|
||||
for k in keys:
|
||||
kid = str(getattr(k, "id", "") or "")
|
||||
if kid:
|
||||
keys_by_id[kid] = k
|
||||
if not keys_by_id:
|
||||
return {}
|
||||
|
||||
with _LOCK:
|
||||
cached = _CACHE.get(provider_id)
|
||||
if cached is not None:
|
||||
expires_at, payload = cached
|
||||
if now < expires_at:
|
||||
missing_ids = [kid for kid in keys_by_id if kid not in payload]
|
||||
if not missing_ids:
|
||||
return {kid: payload[kid] for kid in keys_by_id}
|
||||
# Compute only for missing keys, merge into existing cache
|
||||
for kid in missing_ids:
|
||||
payload[kid] = aggregate_health_score(
|
||||
getattr(keys_by_id[kid], "health_by_format", None)
|
||||
)
|
||||
# Trim stale key entries that are no longer in the current key set
|
||||
stale_ids = [sid for sid in payload if sid not in keys_by_id]
|
||||
for sid in stale_ids:
|
||||
del payload[sid]
|
||||
return {kid: payload[kid] for kid in keys_by_id}
|
||||
|
||||
fresh: dict[str, float] = {}
|
||||
for kid, key in keys_by_id.items():
|
||||
fresh[kid] = aggregate_health_score(getattr(key, "health_by_format", None))
|
||||
|
||||
with _LOCK:
|
||||
_CACHE[provider_id] = (now + _TTL_SECONDS, fresh)
|
||||
# 超出上限时清理过期条目,仍超限则淘汰最旧条目
|
||||
if len(_CACHE) > _MAX_PROVIDERS:
|
||||
expired = [k for k, (exp, _) in _CACHE.items() if now >= exp]
|
||||
for k in expired:
|
||||
del _CACHE[k]
|
||||
if len(_CACHE) > _MAX_PROVIDERS:
|
||||
oldest_key = min(_CACHE, key=lambda k: _CACHE[k][0])
|
||||
del _CACHE[oldest_key]
|
||||
return dict(fresh)
|
||||
|
||||
|
||||
def invalidate_provider_health_scores(provider_id: str) -> None:
|
||||
"""Invalidate health-score cache for one provider."""
|
||||
with _LOCK:
|
||||
_CACHE.pop(provider_id, None)
|
||||
|
||||
|
||||
def _clear_cache_for_tests() -> None:
|
||||
with _LOCK:
|
||||
_CACHE.clear()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"aggregate_health_score",
|
||||
"get_health_scores",
|
||||
"invalidate_provider_health_scores",
|
||||
]
|
||||
359
_deprecated_py_src/services/provider/pool/health_policy.py
Normal file
359
_deprecated_py_src/services/provider/pool/health_policy.py
Normal file
@@ -0,0 +1,359 @@
|
||||
"""Account Pool health policy: error code classification and key state management.
|
||||
|
||||
Maps upstream HTTP status codes to pool-level actions:
|
||||
|
||||
| Code | Action |
|
||||
|--------------|------------------------------------------------------------|
|
||||
| 401 | Invalidate OAuth token cache; permanent (deactivated) 1h, else no cooldown |
|
||||
| 402 | Long cooldown (payment issue) |
|
||||
| 403 | Graded cooldown: severe (suspended/banned) 1h, else 300s+ |
|
||||
| 400 | Check body for "organization has been disabled" -> cooldown |
|
||||
| 429 | Cooldown (retry-after or rate_limit_cooldown_seconds) |
|
||||
| 529 | Cooldown (overload_cooldown_seconds) |
|
||||
| * | Check unschedulable_rules keyword matching |
|
||||
| 408/5xx/etc | Transient cooldown (overload_cooldown_seconds) |
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.services.provider.pool import redis_ops
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from src.services.provider.pool.config import PoolConfig
|
||||
|
||||
# Patterns in 400 error body that indicate account-level issues.
|
||||
_ACCOUNT_DISABLE_PATTERNS = (
|
||||
"organization has been disabled",
|
||||
"organization_disabled",
|
||||
"account has been disabled",
|
||||
"account_disabled",
|
||||
"account has been deactivated",
|
||||
"account_deactivated",
|
||||
"account deactivated",
|
||||
)
|
||||
|
||||
# 需要更长冷却的账号异常语义(403 body 关键字)。
|
||||
_FORBIDDEN_ACCOUNT_PATTERNS = (
|
||||
"account suspended",
|
||||
"account banned",
|
||||
"account deactivated",
|
||||
"subscription inactive",
|
||||
"suspended",
|
||||
"banned",
|
||||
"deactivated",
|
||||
)
|
||||
|
||||
_TRANSIENT_STATUS_COOLDOWN_REASON: dict[int, str] = {
|
||||
408: "request_timeout_408",
|
||||
409: "conflict_409",
|
||||
423: "locked_423",
|
||||
425: "too_early_425",
|
||||
500: "server_error_500",
|
||||
502: "bad_gateway_502",
|
||||
503: "service_unavailable_503",
|
||||
504: "gateway_timeout_504",
|
||||
}
|
||||
|
||||
|
||||
def _parse_retry_after(headers: dict[str, str] | None) -> int | None:
|
||||
"""Extract retry-after seconds from response headers."""
|
||||
if not headers:
|
||||
return None
|
||||
raw = headers.get("retry-after") or headers.get("Retry-After")
|
||||
if not raw:
|
||||
return None
|
||||
try:
|
||||
val = int(raw)
|
||||
return max(1, min(val, 3600))
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
|
||||
|
||||
def _extract_error_message(error_body: str | None) -> str:
|
||||
"""Best-effort extraction of error message from JSON body."""
|
||||
if not error_body:
|
||||
return ""
|
||||
try:
|
||||
data = json.loads(error_body)
|
||||
if isinstance(data, dict):
|
||||
error_obj = data.get("error")
|
||||
if isinstance(error_obj, dict):
|
||||
return str(error_obj.get("message", ""))
|
||||
if isinstance(error_obj, str):
|
||||
return error_obj
|
||||
return str(data.get("message", ""))
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
return error_body[:500]
|
||||
|
||||
|
||||
def _resolve_transient_cooldown_ttl(
|
||||
*,
|
||||
status_code: int,
|
||||
retry_after_seconds: int | None,
|
||||
config: PoolConfig,
|
||||
) -> int:
|
||||
"""Resolve cooldown ttl for transient upstream status codes."""
|
||||
if status_code in (429, 503):
|
||||
if retry_after_seconds is not None:
|
||||
return retry_after_seconds
|
||||
if status_code == 429:
|
||||
return config.rate_limit_cooldown_seconds
|
||||
# 408/409/423/425/5xx: 统一走短时过载冷却,避免雪崩重试。
|
||||
return config.overload_cooldown_seconds
|
||||
|
||||
|
||||
def _parse_google_quota_cooldown(error_body: str | None) -> int | None:
|
||||
"""Parse Google-specific quota cooldown from error body.
|
||||
|
||||
Safe to call for any provider: returns None unless the error body
|
||||
contains Google-specific fields (quotaResetTimeStamp / quotaResetDelay /
|
||||
"reset after" message pattern).
|
||||
"""
|
||||
if not error_body:
|
||||
return None
|
||||
try:
|
||||
from src.services.provider.adapters.gemini_cli.quota import extract_quota_cooldown_seconds
|
||||
|
||||
return extract_quota_cooldown_seconds(error_body)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
async def apply_health_policy(
|
||||
*,
|
||||
provider_id: str,
|
||||
key_id: str,
|
||||
status_code: int,
|
||||
error_body: str | None,
|
||||
response_headers: dict[str, str] | None,
|
||||
config: PoolConfig,
|
||||
) -> None:
|
||||
"""Apply health policy for an upstream error.
|
||||
|
||||
This is fire-and-forget; exceptions are caught and logged.
|
||||
"""
|
||||
if not config.health_policy_enabled:
|
||||
return
|
||||
|
||||
try:
|
||||
await _apply(
|
||||
provider_id=provider_id,
|
||||
key_id=key_id,
|
||||
status_code=status_code,
|
||||
error_body=error_body,
|
||||
response_headers=response_headers,
|
||||
config=config,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Pool health policy failed for key {}: {}",
|
||||
key_id[:8],
|
||||
str(exc),
|
||||
)
|
||||
|
||||
|
||||
async def _apply(
|
||||
*,
|
||||
provider_id: str,
|
||||
key_id: str,
|
||||
status_code: int,
|
||||
error_body: str | None,
|
||||
response_headers: dict[str, str] | None,
|
||||
config: PoolConfig,
|
||||
) -> None:
|
||||
error_msg = _extract_error_message(error_body)
|
||||
|
||||
# --- 401 Unauthorized ---------------------------------------------------
|
||||
if status_code == 401:
|
||||
await redis_ops.invalidate_oauth_token_cache(key_id)
|
||||
# Check if the 401 body indicates a permanent account-level deactivation
|
||||
# (e.g. OpenAI "account_deactivated"). These deserve a long cooldown.
|
||||
error_lower = error_msg.lower()
|
||||
is_permanent = any(p in error_lower for p in _ACCOUNT_DISABLE_PATTERNS)
|
||||
if is_permanent:
|
||||
await redis_ops.set_cooldown(provider_id, key_id, "account_deactivated_401", ttl=3600)
|
||||
logger.warning(
|
||||
"Pool[{}]: key {} got 401 with account deactivation, cooldown 1h",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
)
|
||||
else:
|
||||
# Transient auth failure (e.g. expired token) — token cache already
|
||||
# invalidated above; the next request will trigger a token refresh.
|
||||
# No cooldown needed: the key should be retried immediately after refresh.
|
||||
logger.info(
|
||||
"Pool[{}]: key {} got 401, token cache invalidated (no cooldown)",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
)
|
||||
return
|
||||
|
||||
# --- 402 Payment Required ------------------------------------------------
|
||||
if status_code == 402:
|
||||
await redis_ops.set_cooldown(provider_id, key_id, "payment_required_402", ttl=3600)
|
||||
logger.warning(
|
||||
"Pool[{}]: key {} got 402 (payment required), cooldown 1h",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
)
|
||||
return
|
||||
|
||||
# --- 403 Forbidden -------------------------------------------------------
|
||||
if status_code == 403:
|
||||
error_lower = error_msg.lower()
|
||||
severe = any(pattern in error_lower for pattern in _FORBIDDEN_ACCOUNT_PATTERNS)
|
||||
ttl = 3600 if severe else max(config.rate_limit_cooldown_seconds, 300)
|
||||
await redis_ops.set_cooldown(provider_id, key_id, "forbidden_403", ttl=ttl)
|
||||
logger.warning(
|
||||
"Pool[{}]: key {} got 403 (forbidden), cooldown {}s",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
ttl,
|
||||
)
|
||||
return
|
||||
|
||||
# --- 400 with account-disable pattern ------------------------------------
|
||||
if status_code == 400:
|
||||
error_lower = error_msg.lower()
|
||||
for pattern in _ACCOUNT_DISABLE_PATTERNS:
|
||||
if pattern in error_lower:
|
||||
await redis_ops.set_cooldown(
|
||||
provider_id, key_id, f"account_disabled_400:{pattern}", ttl=3600
|
||||
)
|
||||
logger.warning(
|
||||
"Pool[{}]: key {} got 400 with '{}', cooldown 1h",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
pattern,
|
||||
)
|
||||
return
|
||||
|
||||
# --- 429 Rate Limited ----------------------------------------------------
|
||||
if status_code == 429:
|
||||
retry_after = _parse_retry_after(response_headers)
|
||||
if retry_after is None:
|
||||
retry_after = _parse_google_quota_cooldown(error_body)
|
||||
ttl = _resolve_transient_cooldown_ttl(
|
||||
status_code=status_code,
|
||||
retry_after_seconds=retry_after,
|
||||
config=config,
|
||||
)
|
||||
await redis_ops.set_cooldown(provider_id, key_id, "rate_limited_429", ttl=ttl)
|
||||
logger.info(
|
||||
"Pool[{}]: key {} got 429, cooldown {}s",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
ttl,
|
||||
)
|
||||
return
|
||||
|
||||
# --- 529 Overloaded ------------------------------------------------------
|
||||
if status_code == 529:
|
||||
ttl = config.overload_cooldown_seconds
|
||||
await redis_ops.set_cooldown(provider_id, key_id, "overloaded_529", ttl=ttl)
|
||||
logger.info(
|
||||
"Pool[{}]: key {} got 529, cooldown {}s",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
ttl,
|
||||
)
|
||||
return
|
||||
|
||||
# --- Keyword-based unschedulable rules -----------------------------------
|
||||
if config.unschedulable_rules and error_msg:
|
||||
error_lower = error_msg.lower()
|
||||
for rule in config.unschedulable_rules:
|
||||
if rule.keyword.lower() in error_lower:
|
||||
ttl = max(60, rule.duration_minutes * 60)
|
||||
await redis_ops.set_cooldown(
|
||||
provider_id,
|
||||
key_id,
|
||||
f"rule:{rule.keyword}",
|
||||
ttl=ttl,
|
||||
)
|
||||
logger.info(
|
||||
"Pool[{}]: key {} matched rule '{}', cooldown {}m",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
rule.keyword,
|
||||
rule.duration_minutes,
|
||||
)
|
||||
return
|
||||
|
||||
# --- Transient status bucket (408/409/423/425/5xx) ----------------------
|
||||
reason = _TRANSIENT_STATUS_COOLDOWN_REASON.get(status_code)
|
||||
if reason:
|
||||
retry_after = _parse_retry_after(response_headers)
|
||||
ttl = _resolve_transient_cooldown_ttl(
|
||||
status_code=status_code,
|
||||
retry_after_seconds=retry_after,
|
||||
config=config,
|
||||
)
|
||||
await redis_ops.set_cooldown(provider_id, key_id, reason, ttl=ttl)
|
||||
logger.info(
|
||||
"Pool[{}]: key {} got {}, cooldown {}s ({})",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
status_code,
|
||||
ttl,
|
||||
reason,
|
||||
)
|
||||
return
|
||||
|
||||
|
||||
async def apply_stream_timeout_policy(
|
||||
*,
|
||||
provider_id: str,
|
||||
key_id: str,
|
||||
config: PoolConfig,
|
||||
) -> None:
|
||||
"""Record a stream timeout event and apply cooldown if threshold is reached.
|
||||
|
||||
Called when an upstream stream response times out (no data within the
|
||||
configured interval). Increments a per-key counter in Redis and sets
|
||||
a cooldown if the count reaches the configured threshold.
|
||||
"""
|
||||
if not config.health_policy_enabled:
|
||||
return
|
||||
|
||||
try:
|
||||
count = await redis_ops.incr_stream_timeout_count(
|
||||
provider_id,
|
||||
key_id,
|
||||
config.stream_timeout_window_seconds,
|
||||
)
|
||||
if count >= config.stream_timeout_threshold:
|
||||
ttl = config.stream_timeout_cooldown_seconds
|
||||
await redis_ops.set_cooldown(
|
||||
provider_id,
|
||||
key_id,
|
||||
f"stream_timeout_x{count}",
|
||||
ttl=ttl,
|
||||
)
|
||||
logger.warning(
|
||||
"Pool[{}]: key {} stream timeout count {} >= threshold {}, cooldown {}s",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
count,
|
||||
config.stream_timeout_threshold,
|
||||
ttl,
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
"Pool[{}]: key {} stream timeout count {}/{}",
|
||||
provider_id[:8],
|
||||
key_id[:8],
|
||||
count,
|
||||
config.stream_timeout_threshold,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Pool stream timeout policy failed for key {}: {}",
|
||||
key_id[:8],
|
||||
str(exc),
|
||||
)
|
||||
91
_deprecated_py_src/services/provider/pool/hooks.py
Normal file
91
_deprecated_py_src/services/provider/pool/hooks.py
Normal file
@@ -0,0 +1,91 @@
|
||||
"""Pool scheduling hooks -- provider-type-specific pool behaviour.
|
||||
|
||||
Some provider types need custom logic during pool scheduling (e.g. extracting
|
||||
a session UUID for sticky binding). This module provides a small Protocol +
|
||||
registry so the pool layer stays generic while provider-specific behaviour
|
||||
lives alongside each adapter.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from typing import Any, Protocol, runtime_checkable
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Protocol
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class PoolSchedulingHook(Protocol):
|
||||
"""Provider-type-specific pool scheduling behaviour.
|
||||
|
||||
Each provider type can optionally register a hook to customize:
|
||||
- Session UUID extraction (for sticky sessions)
|
||||
- Post-success / post-error callbacks
|
||||
|
||||
Optional methods (checked via ``hasattr`` by callers):
|
||||
- ``on_pool_success``
|
||||
- ``on_pool_error``
|
||||
"""
|
||||
|
||||
name: str
|
||||
|
||||
def extract_session_uuid(self, request_body: dict[str, Any]) -> str | None:
|
||||
"""Extract a session UUID for sticky binding from the request body."""
|
||||
...
|
||||
|
||||
# -- Optional lifecycle callbacks -----------------------------------------
|
||||
# These are checked via ``hasattr`` so existing implementations that
|
||||
# don't define them will continue to work.
|
||||
|
||||
def on_pool_success(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
session_uuid: str | None,
|
||||
context: dict[str, Any],
|
||||
) -> None:
|
||||
"""Called after a successful pool request (provider-specific logic)."""
|
||||
...
|
||||
|
||||
def on_pool_error(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
status_code: int,
|
||||
context: dict[str, Any],
|
||||
) -> None:
|
||||
"""Called after a failed pool request (provider-specific logic)."""
|
||||
...
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Registry
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_hook_registry: dict[str, PoolSchedulingHook] = {}
|
||||
_registry_lock = threading.Lock()
|
||||
|
||||
|
||||
def register_pool_hook(provider_type: str, hook: PoolSchedulingHook) -> None:
|
||||
"""Register a pool scheduling hook for a provider type."""
|
||||
from src.core.provider_types import normalize_provider_type
|
||||
|
||||
pt = normalize_provider_type(provider_type)
|
||||
with _registry_lock:
|
||||
_hook_registry[pt] = hook
|
||||
|
||||
|
||||
def get_pool_hook(provider_type: str | None) -> PoolSchedulingHook | None:
|
||||
"""Return the pool scheduling hook for a provider type, or ``None``."""
|
||||
if not provider_type:
|
||||
return None
|
||||
from src.services.provider.envelope import ensure_providers_bootstrapped
|
||||
|
||||
ensure_providers_bootstrapped(provider_types=[provider_type])
|
||||
|
||||
from src.core.provider_types import normalize_provider_type
|
||||
|
||||
pt = normalize_provider_type(provider_type)
|
||||
return _hook_registry.get(pt)
|
||||
621
_deprecated_py_src/services/provider/pool/manager.py
Normal file
621
_deprecated_py_src/services/provider/pool/manager.py
Normal file
@@ -0,0 +1,621 @@
|
||||
"""Account Pool Manager (provider-agnostic).
|
||||
|
||||
Stateless facade that coordinates pool operations for any Provider with
|
||||
pool configuration enabled. All state lives in Redis via :mod:`redis_ops`.
|
||||
|
||||
Usage::
|
||||
|
||||
mgr = PoolManager(provider_id, pool_config)
|
||||
reordered = await mgr.reorder_candidates(session_uuid, candidates)
|
||||
# ... execute request ...
|
||||
await mgr.on_request_success(session_uuid=..., key_id=..., tokens_used=...)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import random
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, Any, TypeVar
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.services.provider.pool import redis_ops
|
||||
from src.services.provider.pool.account_state import resolve_pool_account_state
|
||||
from src.services.provider.pool.config import PoolConfig
|
||||
from src.services.provider.pool.health_cache import get_health_scores
|
||||
from src.services.provider.pool.trace import PoolCandidateTrace, PoolSchedulingTrace
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from src.models.database import ProviderAPIKey
|
||||
from src.services.scheduling.schemas import ProviderCandidate
|
||||
|
||||
|
||||
class PoolManager:
|
||||
"""Coordinate pool-level scheduling for a single Provider."""
|
||||
|
||||
__slots__ = ("provider_id", "config", "provider_type")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider_id: str,
|
||||
config: PoolConfig,
|
||||
provider_type: str | None = None,
|
||||
) -> None:
|
||||
self.provider_id = provider_id
|
||||
self.config = config
|
||||
self.provider_type = str(provider_type or "").strip().lower() or None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Core scheduling: reorder candidate list for pool-aware selection
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def reorder_candidates(
|
||||
self,
|
||||
session_uuid: str | None,
|
||||
candidates: list[ProviderCandidate],
|
||||
) -> list[ProviderCandidate]:
|
||||
"""Reorder *candidates* according to pool rules.
|
||||
|
||||
The returned list keeps the same elements but in a new order:
|
||||
|
||||
1. **Sticky session hit** -- if the session is already bound to a key
|
||||
and that key appears in *candidates* and is not in cooldown, move it
|
||||
to position 0.
|
||||
2. **Filter** out keys in account-blocked / cooldown / cost-exhausted state (mark
|
||||
``is_skipped``).
|
||||
3. **LRU sort** -- among remaining candidates at the same priority
|
||||
level, sort by least-recently-used.
|
||||
4. **Random tiebreak** -- among candidates with identical LRU score.
|
||||
|
||||
Also builds a :class:`PoolSchedulingTrace` and attaches per-candidate
|
||||
trace data via ``_pool_extra_data`` / ``_pool_scheduling_trace``
|
||||
attributes on candidate objects.
|
||||
"""
|
||||
if not candidates:
|
||||
return candidates
|
||||
|
||||
pid = self.provider_id
|
||||
|
||||
# Build trace
|
||||
trace = PoolSchedulingTrace(
|
||||
provider_id=pid,
|
||||
total_keys=len(candidates),
|
||||
session_uuid=session_uuid[:8] if session_uuid else None,
|
||||
)
|
||||
|
||||
# --- Strategy: before_select ----------------------------------
|
||||
strategies = _get_active_strategies(self.config)
|
||||
key_ids = [str(c.key.id) for c in candidates]
|
||||
strategy_context: dict[str, Any] = {"session_uuid": session_uuid}
|
||||
for strategy in strategies:
|
||||
if hasattr(strategy, "on_before_select"):
|
||||
try:
|
||||
filtered = strategy.on_before_select(
|
||||
provider_id=pid,
|
||||
key_ids=key_ids,
|
||||
config=self.config,
|
||||
context=strategy_context,
|
||||
)
|
||||
if filtered is not None:
|
||||
key_ids = filtered
|
||||
except Exception:
|
||||
logger.opt(exception=True).debug(
|
||||
"Pool[{}]: strategy before_select failed", pid[:8]
|
||||
)
|
||||
|
||||
# --- 1. Sticky session ----------------------------------------
|
||||
sticky_key_id: str | None = None
|
||||
if session_uuid and self.config.sticky_session_ttl_seconds > 0:
|
||||
sticky_key_id = await redis_ops.get_sticky_binding(
|
||||
pid, session_uuid, self.config.sticky_session_ttl_seconds
|
||||
)
|
||||
|
||||
provider_type = self.provider_type
|
||||
if provider_type is None and candidates:
|
||||
first_provider = getattr(candidates[0], "provider", None)
|
||||
provider_type = str(getattr(first_provider, "provider_type", "") or "").strip().lower()
|
||||
if not provider_type:
|
||||
provider_type = None
|
||||
|
||||
# --- 2. Batch fetch pool state (parallel) ---------------------
|
||||
all_key_ids = [str(c.key.id) for c in candidates]
|
||||
|
||||
# Fire independent Redis queries concurrently.
|
||||
# Only fetch reason (no TTL) on the scheduling hot path -- TTL is only
|
||||
# used for trace display and costs an extra pipeline command per key.
|
||||
_cooldown_coro = redis_ops.batch_get_cooldowns(pid, all_key_ids, include_ttl=False)
|
||||
_cost_coro = (
|
||||
redis_ops.batch_get_cost_totals(pid, all_key_ids, self.config.cost_window_seconds)
|
||||
if (
|
||||
self.config.cost_limit_per_key_tokens is not None
|
||||
or self.config.scheduling_mode == "multi_score"
|
||||
)
|
||||
else None
|
||||
)
|
||||
# LRU scores are needed both for plain LRU sorting and for multi_score
|
||||
# dimensions (e.g. cache_affinity / single_account) that rely on
|
||||
# lru_scores data.
|
||||
_need_lru = self.config.lru_enabled or self.config.scheduling_mode == "multi_score"
|
||||
_lru_coro = redis_ops.get_lru_scores(pid, all_key_ids) if _need_lru else None
|
||||
_latency_coro = (
|
||||
redis_ops.batch_get_latency_avgs(pid, all_key_ids, self.config.latency_window_seconds)
|
||||
if self.config.scheduling_mode == "multi_score"
|
||||
else None
|
||||
)
|
||||
|
||||
# Gather all non-None coroutines in parallel.
|
||||
coros: list[Any] = [_cooldown_coro]
|
||||
_cost_idx = -1
|
||||
_lru_idx = -1
|
||||
_latency_idx = -1
|
||||
if _cost_coro is not None:
|
||||
_cost_idx = len(coros)
|
||||
coros.append(_cost_coro)
|
||||
if _lru_coro is not None:
|
||||
_lru_idx = len(coros)
|
||||
coros.append(_lru_coro)
|
||||
if _latency_coro is not None:
|
||||
_latency_idx = len(coros)
|
||||
coros.append(_latency_coro)
|
||||
|
||||
gathered = await asyncio.gather(*coros)
|
||||
|
||||
cooldowns: dict[str, str | None] = gathered[0]
|
||||
|
||||
# Cost check
|
||||
cost_exhausted: set[str] = set()
|
||||
cost_soft: set[str] = set()
|
||||
cost_totals: dict[str, int] = {}
|
||||
if _cost_idx >= 0:
|
||||
cost_totals = gathered[_cost_idx]
|
||||
limit = self.config.cost_limit_per_key_tokens
|
||||
if limit is not None:
|
||||
for kid, total in cost_totals.items():
|
||||
if total >= limit:
|
||||
cost_exhausted.add(kid)
|
||||
elif total >= limit * self.config.cost_soft_threshold_percent / 100:
|
||||
cost_soft.add(kid)
|
||||
|
||||
# LRU scores
|
||||
lru_scores: dict[str, float] = {}
|
||||
if _lru_idx >= 0:
|
||||
lru_scores = gathered[_lru_idx]
|
||||
|
||||
# Latency averages
|
||||
latency_avgs: dict[str, float] = {}
|
||||
if _latency_idx >= 0:
|
||||
latency_avgs = gathered[_latency_idx]
|
||||
|
||||
# Health scores (TTL cached, no Redis round-trip) -- only needed for multi_score
|
||||
health_scores: dict[str, float] = {}
|
||||
if self.config.scheduling_mode == "multi_score":
|
||||
health_scores = get_health_scores(pid, [c.key for c in candidates])
|
||||
|
||||
strategy_context.update(
|
||||
{
|
||||
"provider_type": provider_type,
|
||||
"all_key_ids": all_key_ids,
|
||||
"lru_scores": lru_scores,
|
||||
"cost_totals": cost_totals,
|
||||
"cost_limit_per_key_tokens": self.config.cost_limit_per_key_tokens,
|
||||
"latency_avgs": latency_avgs,
|
||||
"health_scores": health_scores,
|
||||
"keys_by_id": {str(c.key.id): c.key for c in candidates},
|
||||
}
|
||||
)
|
||||
|
||||
# --- Strategy: compute_score ----------------------------------
|
||||
custom_scores: dict[str, float] = {}
|
||||
for strategy in strategies:
|
||||
if hasattr(strategy, "compute_score"):
|
||||
for kid in all_key_ids:
|
||||
try:
|
||||
custom = strategy.compute_score(
|
||||
key_id=kid,
|
||||
config=self.config,
|
||||
context=strategy_context,
|
||||
)
|
||||
if custom is not None:
|
||||
custom_scores[kid] = float(custom)
|
||||
lru_scores[kid] = float(custom)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# --- 3. Classify candidates -----------------------------------
|
||||
# Use precomputed account states from CandidateBuilder when available
|
||||
# (upstream_metadata is deferred on pool keys to save memory).
|
||||
# Fall back to on-the-fly resolution for non-pool candidates.
|
||||
account_states: dict[str, Any] = {}
|
||||
for c in candidates:
|
||||
kid = str(c.key.id)
|
||||
if kid not in account_states:
|
||||
precomputed = getattr(c.key, "_pool_account_state", None)
|
||||
if precomputed is not None:
|
||||
account_states[kid] = precomputed
|
||||
else:
|
||||
account_states[kid] = resolve_pool_account_state(
|
||||
provider_type=provider_type,
|
||||
upstream_metadata=getattr(c.key, "upstream_metadata", None),
|
||||
oauth_invalid_reason=getattr(c.key, "oauth_invalid_reason", None),
|
||||
)
|
||||
|
||||
sticky_candidate: ProviderCandidate | None = None
|
||||
available: list[ProviderCandidate] = []
|
||||
skipped: list[ProviderCandidate] = []
|
||||
|
||||
for c in candidates:
|
||||
kid = str(c.key.id)
|
||||
ct = PoolCandidateTrace(key_id=kid)
|
||||
ct.scoring_mode = self.config.scheduling_mode
|
||||
ct.latency_avg_ms = float(latency_avgs.get(kid, 0.0) or 0.0)
|
||||
ct.health_score = float(health_scores.get(kid, 1.0) or 1.0)
|
||||
if kid in custom_scores:
|
||||
ct.composite_score = float(custom_scores[kid])
|
||||
|
||||
# Already skipped upstream?
|
||||
if c.is_skipped:
|
||||
skipped.append(c)
|
||||
ct.skipped = True
|
||||
ct.skip_type = "upstream"
|
||||
trace.candidate_traces[kid] = ct
|
||||
continue
|
||||
|
||||
# Account blocked?
|
||||
account_state = account_states[kid]
|
||||
if account_state.blocked:
|
||||
c.is_skipped = True
|
||||
skip_reason = account_state.reason or account_state.label or "account blocked"
|
||||
c.skip_reason = f"pool account blocked: {skip_reason}"
|
||||
skipped.append(c)
|
||||
ct.skipped = True
|
||||
ct.skip_type = "account_blocked"
|
||||
ct.account_block_code = account_state.code
|
||||
ct.account_block_label = account_state.label
|
||||
ct.account_block_reason = account_state.reason
|
||||
_attach_pool_extra(c, ct)
|
||||
trace.candidate_traces[kid] = ct
|
||||
continue
|
||||
|
||||
cd_reason = cooldowns.get(kid)
|
||||
if cd_reason is not None:
|
||||
c.is_skipped = True
|
||||
c.skip_reason = f"pool cooldown: {cd_reason}"
|
||||
skipped.append(c)
|
||||
ct.skipped = True
|
||||
ct.skip_type = "cooldown"
|
||||
ct.cooldown_reason = cd_reason
|
||||
ct.cooldown_ttl = None # TTL skipped on hot path for perf
|
||||
_attach_pool_extra(c, ct)
|
||||
trace.candidate_traces[kid] = ct
|
||||
continue
|
||||
|
||||
# Cost exhausted?
|
||||
if kid in cost_exhausted:
|
||||
c.is_skipped = True
|
||||
c.skip_reason = "pool cost limit reached"
|
||||
skipped.append(c)
|
||||
ct.skipped = True
|
||||
ct.skip_type = "cost_exhausted"
|
||||
ct.cost_window_usage = cost_totals.get(kid, 0)
|
||||
ct.cost_limit = self.config.cost_limit_per_key_tokens
|
||||
_attach_pool_extra(c, ct)
|
||||
trace.candidate_traces[kid] = ct
|
||||
continue
|
||||
|
||||
# Sticky hit?
|
||||
if sticky_key_id and kid == sticky_key_id:
|
||||
sticky_candidate = c
|
||||
ct.reason = "sticky"
|
||||
ct.sticky_hit = True
|
||||
trace.sticky_session_used = True
|
||||
else:
|
||||
available.append(c)
|
||||
if kid in custom_scores and self.config.scheduling_mode == "multi_score":
|
||||
ct.reason = "multi_score"
|
||||
else:
|
||||
ct.reason = "lru" if lru_scores.get(kid, 0) > 0 else "random"
|
||||
|
||||
ct.lru_score = lru_scores.get(kid, 0.0)
|
||||
ct.cost_window_usage = cost_totals.get(kid, 0)
|
||||
ct.cost_limit = self.config.cost_limit_per_key_tokens
|
||||
if kid in cost_soft:
|
||||
ct.cost_soft_threshold = True
|
||||
_attach_pool_extra(c, ct)
|
||||
trace.candidate_traces[kid] = ct
|
||||
|
||||
# --- 4. Sort available by LRU ---------------------------------
|
||||
if lru_scores and available:
|
||||
available.sort(key=lambda c: lru_scores.get(str(c.key.id), 0.0))
|
||||
|
||||
# Random tiebreak among candidates with the same LRU score
|
||||
if len(available) > 1 and lru_scores:
|
||||
_shuffle_same_score_groups(available, lru_scores)
|
||||
|
||||
# --- 5. Assemble final order ----------------------------------
|
||||
result: list[ProviderCandidate] = []
|
||||
if sticky_candidate is not None:
|
||||
result.append(sticky_candidate)
|
||||
result.extend(available)
|
||||
result.extend(skipped)
|
||||
|
||||
if sticky_candidate:
|
||||
logger.debug(
|
||||
"Pool[{}]: sticky hit key={}",
|
||||
pid[:8],
|
||||
sticky_key_id and sticky_key_id[:8],
|
||||
)
|
||||
|
||||
# --- Strategy: after_select -----------------------------------
|
||||
if result:
|
||||
first_kid = str(result[0].key.id)
|
||||
first_trace = trace.candidate_traces.get(first_kid)
|
||||
for strategy in strategies:
|
||||
if hasattr(strategy, "on_after_select") and first_trace:
|
||||
try:
|
||||
strategy.on_after_select(
|
||||
provider_id=pid,
|
||||
selected_key_id=first_kid,
|
||||
trace=first_trace,
|
||||
config=self.config,
|
||||
context=strategy_context,
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Attach the full trace to the first candidate for downstream use.
|
||||
if result:
|
||||
setattr(result[0], "_pool_scheduling_trace", trace)
|
||||
|
||||
return result
|
||||
|
||||
async def select_pool_keys(
|
||||
self,
|
||||
session_uuid: str | None,
|
||||
keys: list[ProviderAPIKey],
|
||||
*,
|
||||
availability_checker: (
|
||||
Callable[[ProviderAPIKey], tuple[bool, str | None, str | None]] | None
|
||||
) = None,
|
||||
page_size: int = 50,
|
||||
) -> tuple[list[ProviderAPIKey], PoolSchedulingTrace]:
|
||||
"""Select and order pool keys with trace output.
|
||||
|
||||
Reuses :meth:`reorder_candidates` logic by adapting keys to lightweight
|
||||
candidate-like wrappers, then propagates skip/trace metadata back onto
|
||||
each key object for downstream execution/recording.
|
||||
|
||||
When *availability_checker* is provided, post-sort availability checks
|
||||
are performed lazily: only the top *page_size* non-skipped keys are
|
||||
checked at a time; if all fail, the next page is checked, and so on.
|
||||
Keys beyond the last checked page are marked as ``deferred`` (skipped
|
||||
without checking) to avoid unnecessary CPU work on large pools.
|
||||
"""
|
||||
if not keys:
|
||||
return (
|
||||
[],
|
||||
PoolSchedulingTrace(
|
||||
provider_id=self.provider_id,
|
||||
total_keys=0,
|
||||
session_uuid=session_uuid[:8] if session_uuid else None,
|
||||
),
|
||||
)
|
||||
|
||||
class _KeyCandidate:
|
||||
__slots__ = (
|
||||
"key",
|
||||
"is_skipped",
|
||||
"skip_reason",
|
||||
"_pool_extra_data",
|
||||
"_pool_scheduling_trace",
|
||||
)
|
||||
|
||||
def __init__(self, key: ProviderAPIKey) -> None:
|
||||
self.key = key
|
||||
self.is_skipped = False
|
||||
self.skip_reason: str | None = None
|
||||
self._pool_extra_data: dict | None = None
|
||||
self._pool_scheduling_trace: PoolSchedulingTrace | None = None
|
||||
|
||||
wrappers = [_KeyCandidate(k) for k in keys]
|
||||
reordered_wrappers = await self.reorder_candidates(session_uuid, wrappers) # type: ignore[arg-type]
|
||||
|
||||
trace: PoolSchedulingTrace | None = None
|
||||
if reordered_wrappers:
|
||||
maybe_trace = getattr(reordered_wrappers[0], "_pool_scheduling_trace", None)
|
||||
if isinstance(maybe_trace, PoolSchedulingTrace):
|
||||
trace = maybe_trace
|
||||
if trace is None:
|
||||
trace = PoolSchedulingTrace(
|
||||
provider_id=self.provider_id,
|
||||
total_keys=len(keys),
|
||||
session_uuid=session_uuid[:8] if session_uuid else None,
|
||||
)
|
||||
|
||||
ordered_keys: list[ProviderAPIKey] = []
|
||||
for order_idx, wrapped in enumerate(reordered_wrappers):
|
||||
key = wrapped.key
|
||||
is_skipped = bool(getattr(wrapped, "is_skipped", False))
|
||||
skip_reason = str(getattr(wrapped, "skip_reason", "") or "")
|
||||
setattr(key, "_pool_skipped", is_skipped)
|
||||
setattr(key, "_pool_skip_reason", skip_reason if skip_reason else None)
|
||||
setattr(key, "_pool_order_index", order_idx)
|
||||
pool_extra = getattr(wrapped, "_pool_extra_data", None)
|
||||
setattr(
|
||||
key, "_pool_extra_data", dict(pool_extra) if isinstance(pool_extra, dict) else {}
|
||||
)
|
||||
ordered_keys.append(key)
|
||||
|
||||
# -- 分页可用性检查 --
|
||||
# 排序后对非 skipped key 分页调用 availability_checker,
|
||||
# 找到 page_size 个可用 key 后停止检查,剩余标记 deferred。
|
||||
if availability_checker is not None:
|
||||
available_count = 0
|
||||
found_enough = False
|
||||
for key in ordered_keys:
|
||||
if getattr(key, "_pool_skipped", False):
|
||||
continue
|
||||
if found_enough:
|
||||
setattr(key, "_pool_skipped", True)
|
||||
setattr(key, "_pool_skip_reason", "deferred")
|
||||
continue
|
||||
is_available, skip_reason_check, mapping_model = availability_checker(key)
|
||||
if not is_available:
|
||||
setattr(key, "_pool_skipped", True)
|
||||
setattr(key, "_pool_skip_reason", skip_reason_check)
|
||||
else:
|
||||
if mapping_model:
|
||||
setattr(key, "_pool_mapping_matched_model", mapping_model)
|
||||
available_count += 1
|
||||
if available_count >= page_size:
|
||||
found_enough = True
|
||||
|
||||
return ordered_keys, trace
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Post-request hooks
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def on_request_success(
|
||||
self,
|
||||
*,
|
||||
session_uuid: str | None,
|
||||
key_id: str,
|
||||
tokens_used: int = 0,
|
||||
ttfb_ms: int | None = None,
|
||||
) -> None:
|
||||
"""Called after a successful upstream request."""
|
||||
pid = self.provider_id
|
||||
|
||||
# Bind sticky session
|
||||
if session_uuid and self.config.sticky_session_ttl_seconds > 0:
|
||||
await redis_ops.set_sticky_binding(
|
||||
pid, session_uuid, key_id, self.config.sticky_session_ttl_seconds
|
||||
)
|
||||
|
||||
# Touch LRU -- needed for both plain LRU mode and multi_score dimensions
|
||||
# (e.g. cache_affinity) that rely on LRU timestamps.
|
||||
if self.config.lru_enabled or self.config.scheduling_mode == "multi_score":
|
||||
await redis_ops.touch_lru(pid, key_id)
|
||||
|
||||
# Record cost
|
||||
if tokens_used > 0 and self.config.cost_limit_per_key_tokens is not None:
|
||||
await redis_ops.add_cost_entry(
|
||||
pid, key_id, tokens_used, self.config.cost_window_seconds
|
||||
)
|
||||
|
||||
# Record latency sample for multi-score scheduling.
|
||||
if self.config.scheduling_mode == "multi_score" and ttfb_ms is not None and ttfb_ms >= 0:
|
||||
await redis_ops.record_latency(
|
||||
pid,
|
||||
key_id,
|
||||
ttfb_ms,
|
||||
self.config.latency_window_seconds,
|
||||
self.config.latency_sample_limit,
|
||||
)
|
||||
|
||||
async def on_request_error(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
status_code: int,
|
||||
error_body: str | None = None,
|
||||
response_headers: dict[str, str] | None = None,
|
||||
) -> None:
|
||||
"""Called after an upstream error. Delegates to health policy."""
|
||||
# Import lazily to avoid circular deps
|
||||
from src.services.provider.pool.health_policy import apply_health_policy
|
||||
|
||||
await apply_health_policy(
|
||||
provider_id=self.provider_id,
|
||||
key_id=key_id,
|
||||
status_code=status_code,
|
||||
error_body=error_body,
|
||||
response_headers=response_headers,
|
||||
config=self.config,
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Key schedulability check (used by candidate_builder)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def is_key_schedulable(self, key_id: str) -> tuple[bool, str | None]:
|
||||
"""Check if *key_id* is currently schedulable (not in cooldown, not
|
||||
cost-exhausted). Returns ``(True, None)`` or ``(False, reason)``.
|
||||
"""
|
||||
pid = self.provider_id
|
||||
|
||||
# Cooldown check
|
||||
cd = await redis_ops.get_cooldown(pid, key_id)
|
||||
if cd is not None:
|
||||
return False, f"pool cooldown: {cd}"
|
||||
|
||||
# Cost check
|
||||
if self.config.cost_limit_per_key_tokens is not None:
|
||||
total = await redis_ops.get_cost_window_total(
|
||||
pid, key_id, self.config.cost_window_seconds
|
||||
)
|
||||
if total >= self.config.cost_limit_per_key_tokens:
|
||||
return False, "pool cost limit reached"
|
||||
|
||||
return True, None
|
||||
|
||||
|
||||
# Backward-compatible alias
|
||||
ClaudeCodePoolManager = PoolManager
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_T = TypeVar("_T")
|
||||
|
||||
|
||||
def _attach_pool_extra(candidate: Any, ct: PoolCandidateTrace) -> None:
|
||||
"""Attach pool trace extra_data onto a candidate object."""
|
||||
existing = getattr(candidate, "_pool_extra_data", None) or {}
|
||||
existing.update(ct.to_extra_data())
|
||||
setattr(candidate, "_pool_extra_data", existing)
|
||||
|
||||
|
||||
def _get_active_strategies(config: PoolConfig) -> list[Any]:
|
||||
"""Get active strategies for the given config (lazy import)."""
|
||||
if not config.strategies:
|
||||
return []
|
||||
try:
|
||||
# Import triggers built-in strategy registration via module-level side effects.
|
||||
from src.services.provider.pool import strategies as _builtin_strategies # noqa: F401
|
||||
from src.services.provider.pool.strategy import get_active_strategies
|
||||
|
||||
return get_active_strategies(config.strategies)
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
||||
def _shuffle_same_score(
|
||||
items: list[_T],
|
||||
lru_scores: dict[str, float],
|
||||
key_fn: Callable[[_T], str],
|
||||
) -> None:
|
||||
"""In-place random shuffle within groups that share the same LRU score."""
|
||||
if len(items) <= 1:
|
||||
return
|
||||
|
||||
i = 0
|
||||
while i < len(items):
|
||||
score_i = lru_scores.get(key_fn(items[i]), 0.0)
|
||||
j = i + 1
|
||||
while j < len(items) and lru_scores.get(key_fn(items[j]), 0.0) == score_i:
|
||||
j += 1
|
||||
if j - i > 1:
|
||||
group = items[i:j]
|
||||
random.shuffle(group)
|
||||
items[i:j] = group
|
||||
i = j
|
||||
|
||||
|
||||
def _shuffle_same_score_groups(
|
||||
candidates: list[ProviderCandidate],
|
||||
lru_scores: dict[str, float],
|
||||
) -> None:
|
||||
_shuffle_same_score(candidates, lru_scores, lambda c: str(c.key.id))
|
||||
42
_deprecated_py_src/services/provider/pool/oauth_cache.py
Normal file
42
_deprecated_py_src/services/provider/pool/oauth_cache.py
Normal file
@@ -0,0 +1,42 @@
|
||||
"""OAuth token Redis cache for the Account Pool.
|
||||
|
||||
Additions over the base ``auth.py`` refresh flow:
|
||||
|
||||
- **Redis token cache**: Avoids repeated DB decryption for hot keys.
|
||||
Cache key: ``provider_oauth_token_cache:{key_id}``
|
||||
- **Configurable proactive refresh skew**: Default 180 s (3 min) instead
|
||||
of the base 120 s, configurable via ``PoolConfig.proactive_refresh_seconds``.
|
||||
- **401 immediate invalidation**: Clears the Redis cache so the next request
|
||||
triggers a fresh refresh.
|
||||
|
||||
This module does NOT replace ``auth.py``; it adds a caching layer that
|
||||
``auth.py`` can consult before decrypting from DB.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.services.provider.pool import redis_ops
|
||||
|
||||
|
||||
async def get_cached_token(key_id: str) -> str | None:
|
||||
"""Return cached access token from Redis, or None."""
|
||||
return await redis_ops.get_cached_oauth_token(key_id)
|
||||
|
||||
|
||||
async def cache_token(key_id: str, token: str, expires_in_seconds: int) -> None:
|
||||
"""Cache an access token in Redis.
|
||||
|
||||
*expires_in_seconds* is the remaining lifetime of the token. We shave
|
||||
off 60 s so the cache expires slightly before the token itself, giving
|
||||
the refresh flow time to act.
|
||||
"""
|
||||
ttl = max(1, expires_in_seconds - 60)
|
||||
await redis_ops.cache_oauth_token(key_id, token, ttl)
|
||||
logger.debug("Pool OAuth: cached token for key {} (TTL={}s)", key_id[:8], ttl)
|
||||
|
||||
|
||||
async def invalidate_token(key_id: str) -> None:
|
||||
"""Invalidate the cached token (e.g. after a 401)."""
|
||||
await redis_ops.invalidate_oauth_token_cache(key_id)
|
||||
logger.debug("Pool OAuth: invalidated token cache for key {}", key_id[:8])
|
||||
694
_deprecated_py_src/services/provider/pool/redis_ops.py
Normal file
694
_deprecated_py_src/services/provider/pool/redis_ops.py
Normal file
@@ -0,0 +1,694 @@
|
||||
"""Redis operations for the Account Pool (provider-agnostic).
|
||||
|
||||
All pool transient state is stored in Redis. This module centralises key
|
||||
naming, Lua scripts, and graceful fallbacks so that the rest of the pool
|
||||
layer is free of Redis specifics.
|
||||
|
||||
Key schema
|
||||
----------
|
||||
ap:{pid}:sticky:{session_uuid} STRING -> key_id (TTL: config)
|
||||
ap:{pid}:lru ZSET member=key_id, score=unix_ts
|
||||
ap:{pid}:cooldown:{key_id} STRING -> reason (TTL: error-specific)
|
||||
ap:{pid}:cost:{key_id} ZSET member=req_id, score=unix_ts
|
||||
ap:{pid}:latency:{key_id} ZSET member=req_id:ttfb_ms, score=unix_ts
|
||||
provider_oauth_token_cache:{key_id} STRING -> access_token (TTL: expires - 60)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
import uuid
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from src.clients.redis_client import get_redis_client
|
||||
from src.core.logger import logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import redis.asyncio as aioredis
|
||||
|
||||
PREFIX = "ap"
|
||||
|
||||
|
||||
def _sticky_key(provider_id: str, session_uuid: str) -> str:
|
||||
return f"{PREFIX}:{provider_id}:sticky:{session_uuid}"
|
||||
|
||||
|
||||
def _lru_key(provider_id: str) -> str:
|
||||
return f"{PREFIX}:{provider_id}:lru"
|
||||
|
||||
|
||||
def _cooldown_key(provider_id: str, key_id: str) -> str:
|
||||
return f"{PREFIX}:{provider_id}:cooldown:{key_id}"
|
||||
|
||||
|
||||
def _cooldown_index_key(provider_id: str) -> str:
|
||||
"""SET tracking which keys are in cooldown (for O(1) count queries)."""
|
||||
return f"{PREFIX}:{provider_id}:cooldown_idx"
|
||||
|
||||
|
||||
def _cost_key(provider_id: str, key_id: str) -> str:
|
||||
return f"{PREFIX}:{provider_id}:cost:{key_id}"
|
||||
|
||||
|
||||
def _latency_key(provider_id: str, key_id: str) -> str:
|
||||
return f"{PREFIX}:{provider_id}:latency:{key_id}"
|
||||
|
||||
|
||||
def _oauth_cache_key(key_id: str) -> str:
|
||||
return f"provider_oauth_token_cache:{key_id}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Lua scripts
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Sticky-select: GET binding, verify it's not in cooldown, refresh TTL.
|
||||
# KEYS[1] = sticky key, KEYS[2] = cooldown key prefix (ap:{pid}:cooldown:)
|
||||
# ARGV[1] = ttl
|
||||
# Returns: key_id or nil
|
||||
_STICKY_SELECT_LUA = """
|
||||
local binding = redis.call("GET", KEYS[1])
|
||||
if not binding then
|
||||
return nil
|
||||
end
|
||||
-- Check cooldown for the bound key
|
||||
local cooldown_key = KEYS[2] .. binding
|
||||
local in_cooldown = redis.call("EXISTS", cooldown_key)
|
||||
if in_cooldown == 1 then
|
||||
redis.call("DEL", KEYS[1])
|
||||
return nil
|
||||
end
|
||||
redis.call("EXPIRE", KEYS[1], tonumber(ARGV[1]))
|
||||
return binding
|
||||
"""
|
||||
|
||||
# Cost window sum (read-only, no cleanup on read path for performance).
|
||||
# KEYS[1] = cost zset key, ARGV[1] = window_start timestamp
|
||||
# Returns total token count within the window.
|
||||
_COST_WINDOW_SUM_LUA = """
|
||||
local key = KEYS[1]
|
||||
local window_start = tonumber(ARGV[1])
|
||||
local members = redis.call("ZRANGEBYSCORE", key, window_start, "+inf")
|
||||
local total = 0
|
||||
for _, m in ipairs(members) do
|
||||
local colon = string.find(m, ":", 1, true)
|
||||
if colon then
|
||||
local n = tonumber(string.sub(m, colon + 1))
|
||||
if n then total = total + n end
|
||||
end
|
||||
end
|
||||
return total
|
||||
"""
|
||||
|
||||
# Latency window average (read-only, no cleanup on read path for performance).
|
||||
# KEYS[1] = latency zset key, ARGV[1] = window_start timestamp
|
||||
# Returns nil when there are no samples, or avg(ms) as number.
|
||||
_LATENCY_WINDOW_AVG_LUA = """
|
||||
local key = KEYS[1]
|
||||
local window_start = tonumber(ARGV[1])
|
||||
local members = redis.call("ZRANGEBYSCORE", key, window_start, "+inf")
|
||||
local total = 0
|
||||
local count = 0
|
||||
for _, m in ipairs(members) do
|
||||
local colon = string.find(m, ":", 1, true)
|
||||
if colon then
|
||||
local n = tonumber(string.sub(m, colon + 1))
|
||||
if n then
|
||||
total = total + n
|
||||
count = count + 1
|
||||
end
|
||||
end
|
||||
end
|
||||
if count == 0 then
|
||||
return nil
|
||||
end
|
||||
return total / count
|
||||
"""
|
||||
|
||||
|
||||
async def _get_redis() -> "aioredis.Redis | None":
|
||||
return await get_redis_client(require_redis=False)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Sticky session
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def get_sticky_binding(provider_id: str, session_uuid: str, ttl: int) -> str | None:
|
||||
"""Get and refresh sticky session binding. Returns key_id or None."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return None
|
||||
try:
|
||||
result = await redis.eval(
|
||||
_STICKY_SELECT_LUA,
|
||||
2,
|
||||
_sticky_key(provider_id, session_uuid),
|
||||
f"{PREFIX}:{provider_id}:cooldown:",
|
||||
str(ttl),
|
||||
)
|
||||
if result:
|
||||
return result.decode() if isinstance(result, bytes) else str(result)
|
||||
return None
|
||||
except Exception:
|
||||
logger.debug("Pool: sticky GET failed for session {}", session_uuid[:8])
|
||||
return None
|
||||
|
||||
|
||||
async def set_sticky_binding(provider_id: str, session_uuid: str, key_id: str, ttl: int) -> None:
|
||||
"""Create or update sticky session binding."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
await redis.setex(_sticky_key(provider_id, session_uuid), ttl, key_id)
|
||||
except Exception:
|
||||
logger.debug("Pool: sticky SET failed for session {}", session_uuid[:8])
|
||||
|
||||
|
||||
async def delete_sticky_binding(provider_id: str, session_uuid: str) -> None:
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
await redis.delete(_sticky_key(provider_id, session_uuid))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# LRU
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def get_lru_scores(provider_id: str, key_ids: list[str]) -> dict[str, float]:
|
||||
"""Batch-fetch LRU timestamps. Missing keys get score 0 (highest priority)."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return {}
|
||||
try:
|
||||
lru_k = _lru_key(provider_id)
|
||||
scores = await redis.zmscore(lru_k, key_ids)
|
||||
result: dict[str, float] = {}
|
||||
for kid, score in zip(key_ids, scores):
|
||||
result[kid] = float(score) if score is not None else 0.0
|
||||
return result
|
||||
except Exception:
|
||||
logger.debug("Pool: LRU ZMSCORE failed for provider {}", provider_id[:8])
|
||||
return {}
|
||||
|
||||
|
||||
async def touch_lru(provider_id: str, key_id: str) -> None:
|
||||
"""Update last-used timestamp."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
await redis.zadd(_lru_key(provider_id), {key_id: time.time()})
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cooldown
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def set_cooldown(provider_id: str, key_id: str, reason: str, ttl: int) -> None:
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
pipe = redis.pipeline()
|
||||
pipe.setex(_cooldown_key(provider_id, key_id), ttl, reason)
|
||||
# Track in index set for O(1) count queries.
|
||||
idx_key = _cooldown_index_key(provider_id)
|
||||
pipe.sadd(idx_key, key_id)
|
||||
# Keep index alive at least as long as the longest cooldown entry.
|
||||
# Each set_cooldown call refreshes the TTL so the SET won't expire
|
||||
# while there are still active cooldowns.
|
||||
pipe.expire(idx_key, ttl + 60)
|
||||
await pipe.execute()
|
||||
logger.info(
|
||||
"Pool: key {} cooldown set: {} ({}s)",
|
||||
key_id[:8],
|
||||
reason,
|
||||
ttl,
|
||||
)
|
||||
except Exception:
|
||||
logger.debug("Pool: cooldown SET failed for key {}", key_id[:8])
|
||||
|
||||
|
||||
async def get_cooldown(provider_id: str, key_id: str) -> str | None:
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return None
|
||||
try:
|
||||
val = await redis.get(_cooldown_key(provider_id, key_id))
|
||||
if val:
|
||||
return val.decode() if isinstance(val, bytes) else str(val)
|
||||
return None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
async def clear_cooldown(provider_id: str, key_id: str) -> None:
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
pipe = redis.pipeline()
|
||||
pipe.delete(_cooldown_key(provider_id, key_id))
|
||||
pipe.srem(_cooldown_index_key(provider_id), key_id)
|
||||
await pipe.execute()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
async def batch_get_cooldowns(
|
||||
provider_id: str,
|
||||
key_ids: list[str],
|
||||
*,
|
||||
include_ttl: bool = False,
|
||||
) -> dict[str, str | None] | dict[str, tuple[str | None, int | None]]:
|
||||
"""Batch check cooldown status for multiple keys.
|
||||
|
||||
When *include_ttl* is ``True``, each value is a ``(reason, ttl_seconds)``
|
||||
tuple instead of a plain reason string. The TTL commands are batched in
|
||||
the same pipeline so there is no extra round-trip.
|
||||
"""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
if include_ttl:
|
||||
return {k: (None, None) for k in key_ids}
|
||||
return {k: None for k in key_ids}
|
||||
try:
|
||||
pipe = redis.pipeline()
|
||||
for kid in key_ids:
|
||||
ck = _cooldown_key(provider_id, kid)
|
||||
pipe.get(ck)
|
||||
if include_ttl:
|
||||
pipe.ttl(ck)
|
||||
results = await pipe.execute()
|
||||
|
||||
if include_ttl:
|
||||
out_ttl: dict[str, tuple[str | None, int | None]] = {}
|
||||
# results interleave GET/TTL: [val0, ttl0, val1, ttl1, ...]
|
||||
for i, kid in enumerate(key_ids):
|
||||
val = results[i * 2]
|
||||
ttl_val = results[i * 2 + 1]
|
||||
reason: str | None = None
|
||||
if val:
|
||||
reason = val.decode() if isinstance(val, bytes) else str(val)
|
||||
ttl_sec: int | None = None
|
||||
if isinstance(ttl_val, int) and ttl_val > 0:
|
||||
ttl_sec = ttl_val
|
||||
out_ttl[kid] = (reason, ttl_sec)
|
||||
return out_ttl
|
||||
|
||||
out: dict[str, str | None] = {}
|
||||
for kid, val in zip(key_ids, results):
|
||||
if val:
|
||||
out[kid] = val.decode() if isinstance(val, bytes) else str(val)
|
||||
else:
|
||||
out[kid] = None
|
||||
return out
|
||||
except Exception:
|
||||
if include_ttl:
|
||||
return {k: (None, None) for k in key_ids}
|
||||
return {k: None for k in key_ids}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cost tracking
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def add_cost_entry(provider_id: str, key_id: str, tokens: int, window_seconds: int) -> None:
|
||||
"""Record a cost entry (tokens used) with automatic window expiry."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
now = time.time()
|
||||
cost_k = _cost_key(provider_id, key_id)
|
||||
member = f"{uuid.uuid4().hex}:{tokens}"
|
||||
window_start = now - max(int(window_seconds), 1)
|
||||
pipe = redis.pipeline()
|
||||
pipe.zadd(cost_k, {member: now})
|
||||
# Prune expired entries on the write path (moved from read Lua script).
|
||||
pipe.zremrangebyscore(cost_k, "-inf", window_start)
|
||||
# Set a TTL slightly larger than the window to auto-clean abandoned keys.
|
||||
pipe.expire(cost_k, window_seconds + 600)
|
||||
await pipe.execute()
|
||||
except Exception:
|
||||
logger.debug("Pool: cost ADD failed for key {}", key_id[:8])
|
||||
|
||||
|
||||
async def get_cost_window_total(provider_id: str, key_id: str, window_seconds: int) -> int:
|
||||
"""Sum tokens used within the rolling window (single key)."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return 0
|
||||
try:
|
||||
now = time.time()
|
||||
window_start = now - window_seconds
|
||||
cost_k = _cost_key(provider_id, key_id)
|
||||
result = await redis.eval(_COST_WINDOW_SUM_LUA, 1, cost_k, str(window_start))
|
||||
return int(result) if result else 0
|
||||
except Exception:
|
||||
logger.debug("Pool: cost SUM failed for key {}", key_id[:8])
|
||||
return 0
|
||||
|
||||
|
||||
async def batch_get_cost_totals(
|
||||
provider_id: str, key_ids: list[str], window_seconds: int
|
||||
) -> dict[str, int]:
|
||||
"""Batch-fetch cost totals for multiple keys using pipeline + Lua."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return {k: 0 for k in key_ids}
|
||||
try:
|
||||
now = time.time()
|
||||
window_start = now - window_seconds
|
||||
pipe = redis.pipeline()
|
||||
for kid in key_ids:
|
||||
cost_k = _cost_key(provider_id, kid)
|
||||
pipe.eval(_COST_WINDOW_SUM_LUA, 1, cost_k, str(window_start))
|
||||
results = await pipe.execute()
|
||||
out: dict[str, int] = {}
|
||||
for kid, val in zip(key_ids, results):
|
||||
out[kid] = int(val) if val else 0
|
||||
return out
|
||||
except Exception:
|
||||
logger.debug("Pool: batch cost SUM failed for provider {}", provider_id[:8])
|
||||
return {k: 0 for k in key_ids}
|
||||
|
||||
|
||||
async def record_latency(
|
||||
provider_id: str,
|
||||
key_id: str,
|
||||
ttfb_ms: int,
|
||||
window_seconds: int,
|
||||
sample_limit: int,
|
||||
) -> None:
|
||||
"""Record one TTFB sample with rolling-window cleanup."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
now = time.time()
|
||||
latency_k = _latency_key(provider_id, key_id)
|
||||
sample = max(int(ttfb_ms), 0)
|
||||
member = f"{uuid.uuid4().hex}:{sample}"
|
||||
pipe = redis.pipeline()
|
||||
pipe.zadd(latency_k, {member: now})
|
||||
window_start = now - max(int(window_seconds), 1)
|
||||
pipe.zremrangebyscore(latency_k, "-inf", window_start)
|
||||
capped_limit = max(int(sample_limit), 1)
|
||||
pipe.zremrangebyrank(latency_k, 0, -(capped_limit + 1))
|
||||
pipe.expire(latency_k, max(int(window_seconds), 1) + 600)
|
||||
await pipe.execute()
|
||||
except Exception:
|
||||
logger.debug("Pool: latency ADD failed for key {}", key_id[:8])
|
||||
|
||||
|
||||
async def batch_get_latency_avgs(
|
||||
provider_id: str,
|
||||
key_ids: list[str],
|
||||
window_seconds: int,
|
||||
) -> dict[str, float]:
|
||||
"""Batch-fetch latency averages (ms) for keys in a rolling window."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return {}
|
||||
try:
|
||||
now = time.time()
|
||||
window_start = now - max(int(window_seconds), 1)
|
||||
pipe = redis.pipeline()
|
||||
for kid in key_ids:
|
||||
pipe.eval(_LATENCY_WINDOW_AVG_LUA, 1, _latency_key(provider_id, kid), str(window_start))
|
||||
results = await pipe.execute()
|
||||
out: dict[str, float] = {}
|
||||
for kid, val in zip(key_ids, results):
|
||||
if val is None:
|
||||
continue
|
||||
try:
|
||||
out[kid] = float(val)
|
||||
except Exception:
|
||||
continue
|
||||
return out
|
||||
except Exception:
|
||||
logger.debug("Pool: batch latency AVG failed for provider {}", provider_id[:8])
|
||||
return {}
|
||||
|
||||
|
||||
async def clear_cost(provider_id: str, key_id: str) -> None:
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
await redis.delete(_cost_key(provider_id, key_id))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# OAuth token cache
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def cache_oauth_token(key_id: str, token: str, ttl: int) -> None:
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
if ttl > 0:
|
||||
await redis.setex(_oauth_cache_key(key_id), ttl, token)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
async def get_cached_oauth_token(key_id: str) -> str | None:
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return None
|
||||
try:
|
||||
val = await redis.get(_oauth_cache_key(key_id))
|
||||
if val:
|
||||
return val.decode() if isinstance(val, bytes) else str(val)
|
||||
return None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
async def invalidate_oauth_token_cache(key_id: str) -> None:
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return
|
||||
try:
|
||||
await redis.delete(_oauth_cache_key(key_id))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pool status query (admin)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def get_sticky_session_count(provider_id: str) -> int:
|
||||
"""Approximate count of active sticky sessions (via SCAN, for admin display only)."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return 0
|
||||
try:
|
||||
pattern = f"{PREFIX}:{provider_id}:sticky:*"
|
||||
count = 0
|
||||
async for _ in redis.scan_iter(match=pattern, count=100):
|
||||
count += 1
|
||||
return count
|
||||
except Exception:
|
||||
return 0
|
||||
|
||||
|
||||
async def get_key_sticky_count(provider_id: str, key_id: str) -> int:
|
||||
"""Count sticky sessions bound to a specific key (admin only).
|
||||
|
||||
Uses batched SCAN + pipeline MGET to reduce Redis round-trips.
|
||||
"""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return 0
|
||||
try:
|
||||
pattern = f"{PREFIX}:{provider_id}:sticky:*"
|
||||
count = 0
|
||||
batch: list[bytes | str] = []
|
||||
async for k in redis.scan_iter(match=pattern, count=200):
|
||||
batch.append(k)
|
||||
if len(batch) >= 200:
|
||||
vals = await redis.mget(batch)
|
||||
for val in vals:
|
||||
if val:
|
||||
bound_id = val.decode() if isinstance(val, bytes) else str(val)
|
||||
if bound_id == key_id:
|
||||
count += 1
|
||||
batch.clear()
|
||||
if batch:
|
||||
vals = await redis.mget(batch)
|
||||
for val in vals:
|
||||
if val:
|
||||
bound_id = val.decode() if isinstance(val, bytes) else str(val)
|
||||
if bound_id == key_id:
|
||||
count += 1
|
||||
return count
|
||||
except Exception:
|
||||
return 0
|
||||
|
||||
|
||||
async def batch_get_key_sticky_counts(
|
||||
provider_id: str,
|
||||
key_ids: list[str],
|
||||
) -> dict[str, int]:
|
||||
"""Count sticky sessions for multiple keys in a single scan (admin only)."""
|
||||
if not key_ids:
|
||||
return {}
|
||||
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return {kid: 0 for kid in key_ids}
|
||||
|
||||
target_ids = set(key_ids)
|
||||
counts: dict[str, int] = {kid: 0 for kid in key_ids}
|
||||
|
||||
try:
|
||||
pattern = f"{PREFIX}:{provider_id}:sticky:*"
|
||||
batch: list[bytes | str] = []
|
||||
async for key in redis.scan_iter(match=pattern, count=200):
|
||||
batch.append(key)
|
||||
if len(batch) >= 200:
|
||||
vals = await redis.mget(batch)
|
||||
for val in vals:
|
||||
if not val:
|
||||
continue
|
||||
bound_id = val.decode() if isinstance(val, bytes) else str(val)
|
||||
if bound_id in target_ids:
|
||||
counts[bound_id] = counts.get(bound_id, 0) + 1
|
||||
batch.clear()
|
||||
|
||||
if batch:
|
||||
vals = await redis.mget(batch)
|
||||
for val in vals:
|
||||
if not val:
|
||||
continue
|
||||
bound_id = val.decode() if isinstance(val, bytes) else str(val)
|
||||
if bound_id in target_ids:
|
||||
counts[bound_id] = counts.get(bound_id, 0) + 1
|
||||
|
||||
return counts
|
||||
except Exception:
|
||||
return {kid: 0 for kid in key_ids}
|
||||
|
||||
|
||||
async def get_cooldown_ttl(provider_id: str, key_id: str) -> int | None:
|
||||
"""Get remaining cooldown TTL in seconds. None = no cooldown."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return None
|
||||
try:
|
||||
ttl = await redis.ttl(_cooldown_key(provider_id, key_id))
|
||||
return ttl if ttl > 0 else None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
async def batch_get_cooldown_ttls(provider_id: str, key_ids: list[str]) -> dict[str, int | None]:
|
||||
"""Batch-fetch cooldown TTLs for multiple keys using pipeline."""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return {k: None for k in key_ids}
|
||||
try:
|
||||
pipe = redis.pipeline()
|
||||
for kid in key_ids:
|
||||
pipe.ttl(_cooldown_key(provider_id, kid))
|
||||
results = await pipe.execute()
|
||||
out: dict[str, int | None] = {}
|
||||
for kid, ttl in zip(key_ids, results):
|
||||
out[kid] = int(ttl) if isinstance(ttl, int) and ttl > 0 else None
|
||||
return out
|
||||
except Exception:
|
||||
return {k: None for k in key_ids}
|
||||
|
||||
|
||||
async def batch_count_provider_cooldowns(provider_ids: list[str]) -> dict[str, int]:
|
||||
"""Count cooldown entries per provider using the cooldown index set.
|
||||
|
||||
Uses ``SCARD`` on the ``ap:{pid}:cooldown_idx`` set for O(1) count
|
||||
instead of scanning the key-space. The index set is maintained by
|
||||
:func:`set_cooldown` / :func:`clear_cooldown`.
|
||||
|
||||
Note: the index set may contain stale entries (expired cooldowns whose
|
||||
TTL elapsed before an explicit clear). This over-count is acceptable
|
||||
for admin display purposes -- precision is not critical here.
|
||||
"""
|
||||
if not provider_ids:
|
||||
return {}
|
||||
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return {pid: 0 for pid in provider_ids}
|
||||
|
||||
try:
|
||||
pipe = redis.pipeline()
|
||||
for pid in provider_ids:
|
||||
pipe.scard(_cooldown_index_key(pid))
|
||||
results = await pipe.execute()
|
||||
counts: dict[str, int] = {}
|
||||
for pid, val in zip(provider_ids, results):
|
||||
counts[pid] = max(int(val or 0), 0)
|
||||
return counts
|
||||
except Exception:
|
||||
return {pid: 0 for pid in provider_ids}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Stream timeout counter
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_STREAM_TIMEOUT_KEY_FMT = f"{PREFIX}:{{}}:stream_timeout:{{}}"
|
||||
|
||||
|
||||
def _stream_timeout_key(provider_id: str, key_id: str) -> str:
|
||||
return _STREAM_TIMEOUT_KEY_FMT.format(provider_id, key_id)
|
||||
|
||||
|
||||
async def incr_stream_timeout_count(
|
||||
provider_id: str,
|
||||
key_id: str,
|
||||
window_seconds: int,
|
||||
) -> int:
|
||||
"""Increment stream timeout counter and return count within the window.
|
||||
|
||||
Uses a ZSET with timestamps as scores. Old entries beyond the window
|
||||
are pruned on each call. Returns the count of timeouts in the window.
|
||||
"""
|
||||
redis = await _get_redis()
|
||||
if redis is None:
|
||||
return 0
|
||||
try:
|
||||
now = time.time()
|
||||
window_start = now - window_seconds
|
||||
key = _stream_timeout_key(provider_id, key_id)
|
||||
member = f"{uuid.uuid4().hex}"
|
||||
pipe = redis.pipeline()
|
||||
pipe.zremrangebyscore(key, "-inf", window_start)
|
||||
pipe.zadd(key, {member: now})
|
||||
pipe.zcard(key)
|
||||
pipe.expire(key, window_seconds + 60)
|
||||
results = await pipe.execute()
|
||||
count = int(results[2]) if results[2] else 0
|
||||
return count
|
||||
except Exception:
|
||||
logger.debug("Pool: stream timeout INCR failed for key {}", key_id[:8])
|
||||
return 0
|
||||
@@ -0,0 +1,466 @@
|
||||
"""Pool scheduling dimension registry and evaluation helpers.
|
||||
|
||||
This module keeps pool scheduling scoring isolated from API layer code.
|
||||
Callers build a :class:`PoolSchedulingSnapshot` and evaluate it against
|
||||
registered dimensions to obtain a normalized summary.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Protocol
|
||||
|
||||
PoolDimensionStatus = str # ok / degraded / blocked
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PoolSchedulingSnapshot:
|
||||
"""Point-in-time scheduling inputs for one key."""
|
||||
|
||||
is_active: bool
|
||||
cooldown_reason: str | None
|
||||
cooldown_ttl_seconds: int | None
|
||||
circuit_breaker_open: bool
|
||||
cost_window_usage: int
|
||||
cost_limit: int | None
|
||||
cost_soft_threshold_percent: int = 80
|
||||
health_score: float = 1.0
|
||||
latency_avg_ms: float | None = None
|
||||
account_blocked: bool = False
|
||||
account_block_label: str | None = None
|
||||
account_block_reason: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PoolSchedulingDimensionResult:
|
||||
"""Evaluation output for one scheduling dimension."""
|
||||
|
||||
code: str
|
||||
label: str
|
||||
status: PoolDimensionStatus = "ok"
|
||||
blocking: bool = False
|
||||
source: str = "pool"
|
||||
weight: int = 1
|
||||
score: float = 1.0
|
||||
detail: str | None = None
|
||||
ttl_seconds: int | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PoolSchedulingSummary:
|
||||
"""Merged scheduling state across all dimensions."""
|
||||
|
||||
status: str # available / degraded / blocked
|
||||
reason: str
|
||||
label: str
|
||||
candidate_eligible: bool
|
||||
blocked_count: int
|
||||
degraded_count: int
|
||||
|
||||
|
||||
class PoolSchedulingDimension(Protocol):
|
||||
"""Dimension evaluator protocol."""
|
||||
|
||||
code: str
|
||||
label: str
|
||||
source: str
|
||||
weight: int
|
||||
|
||||
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
|
||||
"""Evaluate one dimension from snapshot."""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _AccountStateDimension:
|
||||
code: str = "account_state"
|
||||
label: str = "账号状态"
|
||||
source: str = "policy"
|
||||
weight: int = 10
|
||||
|
||||
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
|
||||
if not snapshot.account_blocked:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="ok",
|
||||
score=1.0,
|
||||
)
|
||||
|
||||
blocked_label = snapshot.account_block_label or "账号异常"
|
||||
if blocked_label == "账号封禁":
|
||||
blocked_code = "account_banned"
|
||||
elif blocked_label == "工作区停用":
|
||||
blocked_code = "workspace_deactivated"
|
||||
elif blocked_label == "账号停用":
|
||||
blocked_code = "account_disabled"
|
||||
elif blocked_label == "需要验证":
|
||||
blocked_code = "account_verification"
|
||||
elif blocked_label == "访问受限":
|
||||
blocked_code = "account_forbidden"
|
||||
else:
|
||||
blocked_code = "account_blocked"
|
||||
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=blocked_code,
|
||||
label=blocked_label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="blocked",
|
||||
blocking=True,
|
||||
score=0.0,
|
||||
detail=snapshot.account_block_reason or blocked_label,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ManualEnableDimension:
|
||||
code: str = "manual_disabled"
|
||||
label: str = "禁用"
|
||||
source: str = "manual"
|
||||
weight: int = 8
|
||||
|
||||
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
|
||||
if snapshot.is_active:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="ok",
|
||||
score=1.0,
|
||||
)
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="blocked",
|
||||
blocking=True,
|
||||
score=0.0,
|
||||
detail="账号被手动禁用",
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _CooldownDimension:
|
||||
code: str = "cooldown"
|
||||
label: str = "冷却中"
|
||||
source: str = "pool"
|
||||
weight: int = 7
|
||||
|
||||
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
|
||||
if not snapshot.cooldown_reason:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="ok",
|
||||
score=1.0,
|
||||
)
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="blocked",
|
||||
blocking=True,
|
||||
score=0.0,
|
||||
detail=snapshot.cooldown_reason,
|
||||
ttl_seconds=snapshot.cooldown_ttl_seconds,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _CircuitBreakerDimension:
|
||||
code: str = "circuit_open"
|
||||
label: str = "熔断中"
|
||||
source: str = "health"
|
||||
weight: int = 6
|
||||
|
||||
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
|
||||
if not snapshot.circuit_breaker_open:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="ok",
|
||||
score=1.0,
|
||||
)
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="blocked",
|
||||
blocking=True,
|
||||
score=0.0,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _CostDimension:
|
||||
code: str = "cost"
|
||||
label: str = "成本"
|
||||
source: str = "pool"
|
||||
weight: int = 5
|
||||
|
||||
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
|
||||
limit = snapshot.cost_limit
|
||||
usage = max(snapshot.cost_window_usage, 0)
|
||||
if limit is None or limit <= 0:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="ok",
|
||||
score=1.0,
|
||||
detail=f"{usage}/-",
|
||||
)
|
||||
|
||||
ratio = usage / limit
|
||||
detail = f"{usage}/{limit}"
|
||||
if ratio >= 1.0:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code="cost_exhausted",
|
||||
label="成本超限",
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="blocked",
|
||||
blocking=True,
|
||||
score=0.0,
|
||||
detail=detail,
|
||||
)
|
||||
|
||||
soft_threshold = max(1, min(snapshot.cost_soft_threshold_percent, 100))
|
||||
if ratio * 100 >= soft_threshold:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code="cost_soft",
|
||||
label="成本接近上限",
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="degraded",
|
||||
score=0.45,
|
||||
detail=detail,
|
||||
)
|
||||
|
||||
if ratio >= 0.6:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="degraded",
|
||||
score=0.72,
|
||||
detail=detail,
|
||||
)
|
||||
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="ok",
|
||||
score=1.0,
|
||||
detail=detail,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _HealthDimension:
|
||||
code: str = "health"
|
||||
label: str = "健康度"
|
||||
source: str = "health"
|
||||
weight: int = 4
|
||||
|
||||
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
|
||||
score = max(0.0, min(snapshot.health_score, 1.0))
|
||||
detail = f"{score:.2f}"
|
||||
if score < 0.5:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code="health_low",
|
||||
label="健康度过低",
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="degraded",
|
||||
score=0.3,
|
||||
detail=detail,
|
||||
)
|
||||
if score < 0.8:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code="health_degraded",
|
||||
label="健康度下降",
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="degraded",
|
||||
score=0.65,
|
||||
detail=detail,
|
||||
)
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="ok",
|
||||
score=1.0,
|
||||
detail=detail,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _LatencyDimension:
|
||||
code: str = "latency"
|
||||
label: str = "延迟"
|
||||
source: str = "runtime"
|
||||
weight: int = 3
|
||||
|
||||
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
|
||||
latency = snapshot.latency_avg_ms
|
||||
if latency is None:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="ok",
|
||||
score=1.0,
|
||||
detail="-",
|
||||
)
|
||||
|
||||
value = max(float(latency), 0.0)
|
||||
detail = f"{value:.0f}ms"
|
||||
if value >= 3000:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code="latency_high",
|
||||
label="延迟偏高",
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="degraded",
|
||||
score=0.5,
|
||||
detail=detail,
|
||||
)
|
||||
if value >= 1200:
|
||||
return PoolSchedulingDimensionResult(
|
||||
code="latency_slow",
|
||||
label="延迟较慢",
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="degraded",
|
||||
score=0.72,
|
||||
detail=detail,
|
||||
)
|
||||
return PoolSchedulingDimensionResult(
|
||||
code=self.code,
|
||||
label=self.label,
|
||||
source=self.source,
|
||||
weight=self.weight,
|
||||
status="ok",
|
||||
score=1.0,
|
||||
detail=detail,
|
||||
)
|
||||
|
||||
|
||||
_POOL_DIMENSION_REGISTRY: dict[str, PoolSchedulingDimension] = {}
|
||||
_POOL_DIMENSION_ORDER: list[str] = []
|
||||
|
||||
|
||||
def register_pool_scheduling_dimension(name: str, dimension: PoolSchedulingDimension) -> None:
|
||||
"""Register a dimension evaluator by name."""
|
||||
normalized = name.strip()
|
||||
if not normalized:
|
||||
return
|
||||
if normalized not in _POOL_DIMENSION_ORDER:
|
||||
_POOL_DIMENSION_ORDER.append(normalized)
|
||||
_POOL_DIMENSION_REGISTRY[normalized] = dimension
|
||||
|
||||
|
||||
def get_pool_scheduling_dimension(name: str) -> PoolSchedulingDimension | None:
|
||||
"""Fetch a registered dimension evaluator."""
|
||||
return _POOL_DIMENSION_REGISTRY.get(name.strip())
|
||||
|
||||
|
||||
def list_pool_scheduling_dimensions() -> tuple[str, ...]:
|
||||
"""List registered dimension names in evaluation order."""
|
||||
return tuple(_POOL_DIMENSION_ORDER)
|
||||
|
||||
|
||||
def evaluate_pool_scheduling_dimensions(
|
||||
snapshot: PoolSchedulingSnapshot,
|
||||
*,
|
||||
dimension_names: tuple[str, ...] | None = None,
|
||||
) -> list[PoolSchedulingDimensionResult]:
|
||||
"""Evaluate snapshot across all registered dimensions."""
|
||||
names = dimension_names or list_pool_scheduling_dimensions()
|
||||
results: list[PoolSchedulingDimensionResult] = []
|
||||
for name in names:
|
||||
dimension = get_pool_scheduling_dimension(name)
|
||||
if dimension is None:
|
||||
continue
|
||||
results.append(dimension.evaluate(snapshot))
|
||||
return results
|
||||
|
||||
|
||||
def summarize_pool_scheduling_dimensions(
|
||||
dimensions: list[PoolSchedulingDimensionResult],
|
||||
) -> PoolSchedulingSummary:
|
||||
"""Summarize dimension outputs into a unified scheduling state."""
|
||||
if not dimensions:
|
||||
return PoolSchedulingSummary(
|
||||
status="available",
|
||||
reason="available",
|
||||
label="可用",
|
||||
candidate_eligible=True,
|
||||
blocked_count=0,
|
||||
degraded_count=0,
|
||||
)
|
||||
|
||||
blocked = [item for item in dimensions if item.status == "blocked" or item.blocking]
|
||||
degraded = [item for item in dimensions if item.status == "degraded"]
|
||||
|
||||
if blocked:
|
||||
primary = blocked[0]
|
||||
return PoolSchedulingSummary(
|
||||
status="blocked",
|
||||
reason=primary.code,
|
||||
label=primary.label,
|
||||
candidate_eligible=False,
|
||||
blocked_count=len(blocked),
|
||||
degraded_count=len(degraded),
|
||||
)
|
||||
|
||||
if degraded:
|
||||
primary = degraded[0]
|
||||
return PoolSchedulingSummary(
|
||||
status="degraded",
|
||||
reason=primary.code,
|
||||
label=primary.label,
|
||||
candidate_eligible=True,
|
||||
blocked_count=0,
|
||||
degraded_count=len(degraded),
|
||||
)
|
||||
|
||||
return PoolSchedulingSummary(
|
||||
status="available",
|
||||
reason="available",
|
||||
label="可用",
|
||||
candidate_eligible=True,
|
||||
blocked_count=0,
|
||||
degraded_count=0,
|
||||
)
|
||||
|
||||
|
||||
def _register_default_dimensions() -> None:
|
||||
register_pool_scheduling_dimension("account_state", _AccountStateDimension())
|
||||
register_pool_scheduling_dimension("manual", _ManualEnableDimension())
|
||||
register_pool_scheduling_dimension("cooldown", _CooldownDimension())
|
||||
register_pool_scheduling_dimension("circuit", _CircuitBreakerDimension())
|
||||
register_pool_scheduling_dimension("cost", _CostDimension())
|
||||
register_pool_scheduling_dimension("latency", _LatencyDimension())
|
||||
register_pool_scheduling_dimension("health", _HealthDimension())
|
||||
|
||||
|
||||
_register_default_dimensions()
|
||||
@@ -0,0 +1,8 @@
|
||||
"""Built-in pool strategies."""
|
||||
|
||||
# Import side effects: register built-in strategies.
|
||||
import src.services.provider.pool.dimensions # noqa: F401
|
||||
|
||||
from . import multi_score # noqa: F401
|
||||
|
||||
__all__ = ["multi_score"]
|
||||
@@ -0,0 +1,264 @@
|
||||
"""Multi-dimension pool scoring strategy."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import src.services.provider.pool.dimensions # noqa: F401
|
||||
from src.services.provider.pool.dimensions import get_preset_dimension, get_preset_names
|
||||
from src.services.provider.pool.dimensions._helpers import rank_ascending, safe_float
|
||||
from src.services.provider.pool.strategy import register_pool_strategy
|
||||
|
||||
|
||||
def _normalize_mutex_group(value: Any) -> str | None:
|
||||
if not isinstance(value, str):
|
||||
return None
|
||||
normalized = value.strip().lower()
|
||||
return normalized or None
|
||||
|
||||
|
||||
def _get_preset_mutex_group(preset_name: str) -> str | None:
|
||||
# LRU is a built-in preset (not in registry) but shares the distribution mutex group.
|
||||
if preset_name == "lru":
|
||||
return "distribution_mode"
|
||||
dim = get_preset_dimension(preset_name)
|
||||
if dim is None:
|
||||
return None
|
||||
return _normalize_mutex_group(getattr(dim, "mutex_group", None))
|
||||
|
||||
|
||||
def _normalize_presets_from_config(
|
||||
config: Any,
|
||||
*,
|
||||
provider_type: str | None = None,
|
||||
) -> tuple[tuple[str, str | None], ...]:
|
||||
"""Extract enabled (preset_name, mode) tuples from config.scheduling_presets.
|
||||
|
||||
Supports both new SchedulingPreset objects and legacy string lists.
|
||||
Excludes ``lru`` from output (LRU is a final tie-breaker only).
|
||||
For mutex groups, enabled members inherit the group's first appearance index
|
||||
so the selected member keeps the group's visible priority slot.
|
||||
"""
|
||||
|
||||
raw = getattr(config, "scheduling_presets", ())
|
||||
if not isinstance(raw, (list, tuple)):
|
||||
return ()
|
||||
|
||||
normalized_provider_type = str(provider_type or "").strip().lower()
|
||||
allowed = get_preset_names() | {"lru"}
|
||||
entries: list[tuple[int, str, bool, str | None]] = []
|
||||
seen: set[str] = set()
|
||||
for idx, item in enumerate(raw):
|
||||
preset_name: str | None = None
|
||||
enabled = True
|
||||
mode: str | None = None
|
||||
|
||||
if hasattr(item, "preset"):
|
||||
preset_name = str(getattr(item, "preset", "")).strip().lower()
|
||||
enabled = bool(getattr(item, "enabled", True))
|
||||
raw_mode = getattr(item, "mode", None)
|
||||
if isinstance(raw_mode, str):
|
||||
mode = raw_mode.strip().lower() or None
|
||||
elif isinstance(item, str):
|
||||
preset_name = item.strip().lower()
|
||||
else:
|
||||
continue
|
||||
|
||||
if not preset_name or preset_name not in allowed or preset_name in seen:
|
||||
continue
|
||||
seen.add(preset_name)
|
||||
entries.append((idx, preset_name, enabled, mode))
|
||||
|
||||
# Codex 默认启用额度刷新优先维度(除非显式配置了 recent_refresh)。
|
||||
if (
|
||||
normalized_provider_type == "codex"
|
||||
and entries
|
||||
and "recent_refresh" not in {name for _idx, name, _enabled, _mode in entries}
|
||||
and "recent_refresh" in allowed
|
||||
):
|
||||
entries.append((len(entries), "recent_refresh", True, None))
|
||||
|
||||
if not entries:
|
||||
return ()
|
||||
|
||||
group_anchor_index: dict[str, int] = {}
|
||||
for idx, preset_name, _enabled, _mode in entries:
|
||||
mutex_group = _get_preset_mutex_group(preset_name)
|
||||
if mutex_group and mutex_group not in group_anchor_index:
|
||||
group_anchor_index[mutex_group] = idx
|
||||
|
||||
ordered_enabled: list[tuple[int, int, str, str | None]] = []
|
||||
group_enabled: dict[str, tuple[int, int, str, str | None]] = {}
|
||||
for idx, preset_name, enabled, mode in entries:
|
||||
if not enabled or preset_name == "lru":
|
||||
continue
|
||||
mutex_group = _get_preset_mutex_group(preset_name)
|
||||
if not mutex_group:
|
||||
ordered_enabled.append((idx, idx, preset_name, mode))
|
||||
continue
|
||||
|
||||
anchor = group_anchor_index.get(mutex_group, idx)
|
||||
existing = group_enabled.get(mutex_group)
|
||||
if existing is None or idx < existing[1]:
|
||||
group_enabled[mutex_group] = (anchor, idx, preset_name, mode)
|
||||
|
||||
ordered_enabled.extend(group_enabled.values())
|
||||
ordered_enabled.sort(key=lambda item: (item[0], item[1]))
|
||||
return tuple((preset_name, mode) for _anchor, _idx, preset_name, mode in ordered_enabled)
|
||||
|
||||
|
||||
class MultiScoreStrategy:
|
||||
name = "multi_score"
|
||||
|
||||
def compute_score(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
config: Any,
|
||||
context: dict[str, Any],
|
||||
) -> float | None:
|
||||
mode = str(getattr(config, "scheduling_mode", "lru") or "lru").strip().lower()
|
||||
if mode != "multi_score":
|
||||
return None
|
||||
|
||||
all_key_ids = [str(k) for k in (context.get("all_key_ids") or []) if str(k)]
|
||||
if not all_key_ids:
|
||||
return None
|
||||
|
||||
lru_scores = context.get("lru_scores", {})
|
||||
if not isinstance(lru_scores, dict):
|
||||
lru_scores = {}
|
||||
latency_avgs = context.get("latency_avgs", {})
|
||||
if not isinstance(latency_avgs, dict):
|
||||
latency_avgs = {}
|
||||
health_scores = context.get("health_scores", {})
|
||||
if not isinstance(health_scores, dict):
|
||||
health_scores = {}
|
||||
cost_totals = context.get("cost_totals", {})
|
||||
if not isinstance(cost_totals, dict):
|
||||
cost_totals = {}
|
||||
keys_by_id = context.get("keys_by_id", {})
|
||||
if not isinstance(keys_by_id, dict):
|
||||
keys_by_id = {}
|
||||
|
||||
presets = _normalize_presets_from_config(
|
||||
config,
|
||||
provider_type=context.get("provider_type"),
|
||||
)
|
||||
lru_enabled = bool(getattr(config, "lru_enabled", True))
|
||||
if presets:
|
||||
return self._compute_preset_score(
|
||||
key_id=key_id,
|
||||
all_key_ids=all_key_ids,
|
||||
presets=presets,
|
||||
lru_enabled=lru_enabled,
|
||||
lru_scores=lru_scores,
|
||||
keys_by_id=keys_by_id,
|
||||
context=context,
|
||||
)
|
||||
|
||||
weights = getattr(config, "scoring_weights", None)
|
||||
w_lru = safe_float(getattr(weights, "lru", 0.3)) or 0.0
|
||||
w_latency = safe_float(getattr(weights, "latency", 0.25)) or 0.0
|
||||
w_health = safe_float(getattr(weights, "health", 0.2)) or 0.0
|
||||
w_cost = safe_float(getattr(weights, "cost_remaining", 0.25)) or 0.0
|
||||
|
||||
lru_rank = rank_ascending(key_id, lru_scores, all_key_ids)
|
||||
latency_rank = rank_ascending(key_id, latency_avgs, all_key_ids)
|
||||
|
||||
health_raw = safe_float(health_scores.get(key_id))
|
||||
if health_raw is None:
|
||||
health_raw = 1.0
|
||||
health_norm = 1.0 - max(0.0, min(health_raw, 1.0))
|
||||
|
||||
cost_limit = getattr(config, "cost_limit_per_key_tokens", None)
|
||||
used = safe_float(cost_totals.get(key_id)) or 0.0
|
||||
if cost_limit is None or int(cost_limit) <= 0:
|
||||
cost_norm = 0.0
|
||||
else:
|
||||
cost_norm = max(0.0, min(used / float(cost_limit), 1.0))
|
||||
|
||||
return (
|
||||
w_lru * lru_rank
|
||||
+ w_latency * latency_rank
|
||||
+ w_health * health_norm
|
||||
+ w_cost * cost_norm
|
||||
)
|
||||
|
||||
def _compute_preset_score(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
all_key_ids: list[str],
|
||||
presets: tuple[tuple[str, str | None], ...],
|
||||
lru_enabled: bool,
|
||||
lru_scores: dict[str, Any],
|
||||
keys_by_id: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
) -> float:
|
||||
cache_signature = (tuple(all_key_ids), presets, bool(lru_enabled))
|
||||
cache = context.get("_preset_hard_order_cache")
|
||||
if (
|
||||
isinstance(cache, dict)
|
||||
and cache.get("signature") == cache_signature
|
||||
and isinstance(cache.get("ranks"), dict)
|
||||
):
|
||||
cached_rank = safe_float(cache["ranks"].get(key_id))
|
||||
if cached_rank is not None:
|
||||
return max(0.0, min(cached_rank, 1.0))
|
||||
|
||||
# Hard-priority semantics:
|
||||
# 1) Compare by preset[0] metric first;
|
||||
# 2) only if tied, compare preset[1], preset[2], ...
|
||||
# 3) if all preset metrics tie and LRU is enabled, use LRU as final tiebreak.
|
||||
metric_vectors: dict[str, tuple[float, ...]] = {}
|
||||
for kid in all_key_ids:
|
||||
vector_parts: list[float] = []
|
||||
for preset_name, mode in presets:
|
||||
metric = 0.5
|
||||
dim = get_preset_dimension(preset_name)
|
||||
if dim is not None:
|
||||
metric = dim.compute_metric(
|
||||
key_id=kid,
|
||||
all_key_ids=all_key_ids,
|
||||
keys_by_id=keys_by_id,
|
||||
lru_scores=lru_scores,
|
||||
context=context,
|
||||
mode=mode,
|
||||
)
|
||||
metric_value = safe_float(metric)
|
||||
vector_parts.append(
|
||||
max(0.0, min(metric_value, 1.0)) if metric_value is not None else 0.5
|
||||
)
|
||||
|
||||
if lru_enabled:
|
||||
vector_parts.append(rank_ascending(kid, lru_scores, all_key_ids))
|
||||
|
||||
metric_vectors[kid] = tuple(vector_parts)
|
||||
|
||||
decorated = [
|
||||
(metric_vectors.get(kid, (0.5,)), idx, kid) for idx, kid in enumerate(all_key_ids)
|
||||
]
|
||||
decorated.sort(key=lambda item: (item[0], item[1]))
|
||||
|
||||
total = len(decorated)
|
||||
ranks: dict[str, float] = {}
|
||||
for rank_idx, (_vec, _idx, kid) in enumerate(decorated):
|
||||
ranks[kid] = 0.0 if total <= 1 else rank_idx / float(total - 1)
|
||||
|
||||
context["_preset_hard_order_cache"] = {
|
||||
"signature": cache_signature,
|
||||
"ranks": ranks,
|
||||
}
|
||||
rank = safe_float(ranks.get(key_id))
|
||||
if rank is None:
|
||||
return 0.5
|
||||
return max(0.0, min(rank, 1.0))
|
||||
|
||||
|
||||
register_pool_strategy("multi_score", MultiScoreStrategy())
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MultiScoreStrategy",
|
||||
]
|
||||
107
_deprecated_py_src/services/provider/pool/strategy.py
Normal file
107
_deprecated_py_src/services/provider/pool/strategy.py
Normal file
@@ -0,0 +1,107 @@
|
||||
"""Pluggable pool scheduling strategies.
|
||||
|
||||
Strategies allow customising pool-level candidate selection without
|
||||
modifying the core :class:`PoolManager`. Each strategy is an object
|
||||
that implements one or more optional methods defined by the
|
||||
:class:`PoolSchedulingStrategy` protocol.
|
||||
|
||||
Registration uses a thread-safe global registry (same pattern as
|
||||
:mod:`~src.services.provider.pool.hooks`).
|
||||
|
||||
Usage::
|
||||
|
||||
from src.services.provider.pool.strategy import register_pool_strategy
|
||||
|
||||
class MyStrategy:
|
||||
name = "usage_weight"
|
||||
|
||||
def compute_score(self, *, key_id, config, context):
|
||||
...
|
||||
|
||||
register_pool_strategy("usage_weight", MyStrategy())
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from src.services.provider.pool.config import PoolConfig
|
||||
from src.services.provider.pool.trace import PoolCandidateTrace
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class PoolSchedulingStrategy(Protocol):
|
||||
"""Pluggable pool scheduling strategy.
|
||||
|
||||
All methods are optional -- callers check via ``hasattr``.
|
||||
Strategies are activated per-provider through ``PoolConfig.strategies``.
|
||||
"""
|
||||
|
||||
name: str
|
||||
|
||||
def on_before_select(
|
||||
self,
|
||||
*,
|
||||
provider_id: str,
|
||||
key_ids: list[str],
|
||||
config: PoolConfig,
|
||||
context: dict[str, Any],
|
||||
) -> list[str] | None:
|
||||
"""Filter / reorder *key_ids* before selection.
|
||||
|
||||
Return ``None`` to leave the list unchanged.
|
||||
"""
|
||||
...
|
||||
|
||||
def on_after_select(
|
||||
self,
|
||||
*,
|
||||
provider_id: str,
|
||||
selected_key_id: str,
|
||||
trace: PoolCandidateTrace,
|
||||
config: PoolConfig,
|
||||
context: dict[str, Any],
|
||||
) -> None:
|
||||
"""Called after a key has been selected (for logging / metrics)."""
|
||||
...
|
||||
|
||||
def compute_score(
|
||||
self,
|
||||
*,
|
||||
key_id: str,
|
||||
config: PoolConfig,
|
||||
context: dict[str, Any],
|
||||
) -> float | None:
|
||||
"""Return a custom sort score. ``None`` means "do not override"."""
|
||||
...
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Registry
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_strategy_registry: dict[str, PoolSchedulingStrategy] = {}
|
||||
_strategy_lock = threading.Lock()
|
||||
|
||||
|
||||
def register_pool_strategy(name: str, strategy: PoolSchedulingStrategy) -> None:
|
||||
"""Register a pool scheduling strategy globally."""
|
||||
with _strategy_lock:
|
||||
_strategy_registry[name] = strategy
|
||||
|
||||
|
||||
def get_pool_strategy(name: str) -> PoolSchedulingStrategy | None:
|
||||
"""Return a registered strategy by *name*, or ``None``."""
|
||||
return _strategy_registry.get(name)
|
||||
|
||||
|
||||
def get_active_strategies(names: tuple[str, ...] | list[str]) -> list[PoolSchedulingStrategy]:
|
||||
"""Return registered strategies whose names appear in *names*."""
|
||||
result: list[PoolSchedulingStrategy] = []
|
||||
for n in names:
|
||||
s = _strategy_registry.get(n)
|
||||
if s is not None:
|
||||
result.append(s)
|
||||
return result
|
||||
137
_deprecated_py_src/services/provider/pool/trace.py
Normal file
137
_deprecated_py_src/services/provider/pool/trace.py
Normal file
@@ -0,0 +1,137 @@
|
||||
"""Pool scheduling trace -- per-candidate decision records.
|
||||
|
||||
Collects scheduling decisions made during pool-level candidate selection
|
||||
without adding any extra Redis round-trips. Trace data is later written
|
||||
to ``RequestCandidate.extra_data`` and ``Usage.request_metadata``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class PoolCandidateTrace:
|
||||
"""Single candidate scheduling decision in a pool context."""
|
||||
|
||||
key_id: str
|
||||
reason: str = "" # sticky / lru / random / tiebreak
|
||||
sticky_hit: bool = False
|
||||
lru_score: float = 0.0
|
||||
cost_window_usage: int = 0
|
||||
cost_limit: int | None = None
|
||||
cost_soft_threshold: bool = False
|
||||
skipped: bool = False
|
||||
skip_type: str | None = None # cooldown / cost_exhausted / account_blocked / upstream
|
||||
cooldown_reason: str | None = None
|
||||
cooldown_ttl: int | None = None
|
||||
account_block_code: str | None = None
|
||||
account_block_label: str | None = None
|
||||
account_block_reason: str | None = None
|
||||
latency_avg_ms: float = 0.0
|
||||
health_score: float = 1.0
|
||||
composite_score: float = 0.0
|
||||
scoring_mode: str = "lru"
|
||||
|
||||
def to_extra_data(self) -> dict[str, Any]:
|
||||
"""Build dict to merge into ``RequestCandidate.extra_data``."""
|
||||
if self.skipped:
|
||||
skip_info: dict[str, Any] = {"type": self.skip_type}
|
||||
if self.cooldown_reason is not None:
|
||||
skip_info["cooldown_reason"] = self.cooldown_reason
|
||||
if self.cooldown_ttl is not None:
|
||||
skip_info["cooldown_ttl"] = self.cooldown_ttl
|
||||
if self.account_block_code is not None:
|
||||
skip_info["account_block_code"] = self.account_block_code
|
||||
if self.account_block_label is not None:
|
||||
skip_info["account_block_label"] = self.account_block_label
|
||||
if self.account_block_reason is not None:
|
||||
skip_info["account_block_reason"] = self.account_block_reason
|
||||
if self.cost_window_usage:
|
||||
skip_info["cost_window_usage"] = self.cost_window_usage
|
||||
if self.scoring_mode:
|
||||
skip_info["scoring_mode"] = self.scoring_mode
|
||||
return {"pool_skip": skip_info}
|
||||
|
||||
sel: dict[str, Any] = {"reason": self.reason}
|
||||
if self.sticky_hit:
|
||||
sel["sticky_hit"] = True
|
||||
if self.lru_score:
|
||||
sel["lru_score"] = self.lru_score
|
||||
if self.cost_window_usage:
|
||||
sel["cost_window_usage"] = self.cost_window_usage
|
||||
if self.cost_limit is not None:
|
||||
sel["cost_limit"] = self.cost_limit
|
||||
if self.cost_soft_threshold:
|
||||
sel["cost_soft_threshold"] = True
|
||||
if self.latency_avg_ms > 0:
|
||||
sel["latency_avg_ms"] = round(self.latency_avg_ms, 2)
|
||||
if self.health_score < 1.0:
|
||||
sel["health_score"] = round(self.health_score, 4)
|
||||
if self.reason == "multi_score":
|
||||
sel["composite_score"] = round(self.composite_score, 6)
|
||||
if self.scoring_mode:
|
||||
sel["scoring_mode"] = self.scoring_mode
|
||||
return {"pool_selection": sel}
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class PoolSchedulingTrace:
|
||||
"""Aggregated scheduling trace for one pool-provider dispatch."""
|
||||
|
||||
provider_id: str
|
||||
total_keys: int = 0
|
||||
sticky_session_used: bool = False
|
||||
session_uuid: str | None = None
|
||||
candidate_traces: dict[str, PoolCandidateTrace] = field(default_factory=dict)
|
||||
|
||||
def build_summary(
|
||||
self,
|
||||
success_key_id: str | None = None,
|
||||
*,
|
||||
attempted_key_ids: set[str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Build compact dict for ``Usage.request_metadata["pool_summary"]``."""
|
||||
skipped_cooldown = 0
|
||||
skipped_cost = 0
|
||||
skipped_account_blocked = 0
|
||||
attempted = 0
|
||||
for t in self.candidate_traces.values():
|
||||
if t.skipped:
|
||||
if t.skip_type == "cooldown":
|
||||
skipped_cooldown += 1
|
||||
elif t.skip_type == "cost_exhausted":
|
||||
skipped_cost += 1
|
||||
elif t.skip_type == "account_blocked":
|
||||
skipped_account_blocked += 1
|
||||
|
||||
if attempted_key_ids is None:
|
||||
# Backward-compatible behavior: count all schedulable keys.
|
||||
attempted = sum(1 for t in self.candidate_traces.values() if not t.skipped)
|
||||
else:
|
||||
# Preferred behavior: count only keys that were actually executed.
|
||||
attempted = sum(
|
||||
1
|
||||
for kid in attempted_key_ids
|
||||
if kid in self.candidate_traces and not self.candidate_traces[kid].skipped
|
||||
)
|
||||
|
||||
success_reason: str | None = None
|
||||
if success_key_id and success_key_id in self.candidate_traces:
|
||||
success_reason = self.candidate_traces[success_key_id].reason
|
||||
|
||||
summary: dict[str, Any] = {
|
||||
"enabled": True,
|
||||
"total_keys": self.total_keys,
|
||||
"attempted": attempted,
|
||||
"skipped_cooldown": skipped_cooldown,
|
||||
"skipped_cost": skipped_cost,
|
||||
"skipped_account_blocked": skipped_account_blocked,
|
||||
"sticky_session": self.sticky_session_used,
|
||||
}
|
||||
if success_key_id:
|
||||
summary["success_key_id"] = success_key_id[:8]
|
||||
if success_reason:
|
||||
summary["success_reason"] = success_reason
|
||||
return summary
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user