refactor: 移除 Python 后端源码,全面迁移至 Rust gateway 架构

- 删除全部 Python 源码 (src/) 及 Alembic 迁移脚本,归档至 _deprecated_py_src/
- 重构 Rust gateway ai_pipeline: 拆分 planner/finalize 模块,新增 contracts/adaptation 层
- 重组 handlers 模块为 admin/public/proxy/internal/shared 子模块结构
- 新增 executor 模块,引入 Rust 原生数据库迁移 (aether-data/migrations)
- 简化 CI/Docker 构建流程,移除 base image 二级构建,统一为单一 app image
- 移除 Python 相关基础设施文件 (entrypoint.sh, gunicorn_conf.py, Dockerfile.base)
This commit is contained in:
fawney19
2026-04-03 16:26:16 +08:00
parent 8f26e1a31f
commit 1d9c77522a
868 changed files with 1735 additions and 2433 deletions

View 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}")

View File

@@ -0,0 +1 @@
"""Provider-specific adapters (antigravity, codex, etc.)."""

View File

@@ -0,0 +1 @@
"""Antigravity integration package."""

View File

@@ -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.rsSandbox 优先,避免 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 的 4xx400/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": {...}}, ...}}
对齐 AMSandbox 优先 + 正确的 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",
]

View File

@@ -0,0 +1,279 @@
"""Antigravity 全局常量定义。
注意:这里的 PROVIDER_TYPE 指的是 Provider.provider_type用于路由与特判
不是 endpoint signaturefamily: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

View File

@@ -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优先 paidTierfallback 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优先 paidTierfallback 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)

View File

@@ -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"]

View File

@@ -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"]

View File

@@ -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"]

View File

@@ -0,0 +1,3 @@
"""Claude Code provider adapter."""
__all__ = []

View File

@@ -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",
]

View File

@@ -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",
]

View File

@@ -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",
]

View File

@@ -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``
异步接管)时仅做 maskingmasking 始终在会话限制检查之后执行,避免
用被伪装后的 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",
]

View File

@@ -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"]

View File

@@ -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"]

View File

@@ -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"]

View File

@@ -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"]

View File

@@ -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()

View File

@@ -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"]

View File

@@ -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"]

View File

@@ -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,
)

View File

@@ -0,0 +1,3 @@
"""Codex provider integration package."""
__all__ = []

View File

@@ -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",
]

View 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,
)

View File

@@ -0,0 +1,5 @@
"""GeminiCLI provider adapter package."""
from .plugin import register_all
__all__ = ["register_all"]

View File

@@ -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",
]

View File

@@ -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",
]

View File

@@ -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",
]

View File

@@ -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",
]

View File

@@ -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",
]

View File

@@ -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"]

View File

@@ -0,0 +1,3 @@
"""Kiro provider adapter."""
__all__ = []

View File

@@ -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",
]

View File

@@ -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",
]

View 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",
]

View 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"]

View File

@@ -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",
]

View File

@@ -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",
]

View 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",
]

View File

@@ -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",
]

View File

@@ -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 regionIdC 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"]

View File

@@ -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",
]

View File

@@ -0,0 +1,6 @@
"""AWS Event Stream parser for Kiro."""
from .decoder import EventStreamDecoder
from .frame import Frame
__all__ = ["EventStreamDecoder", "Frame"]

View File

@@ -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"]

View File

@@ -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",
]

View File

@@ -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",
]

View File

@@ -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",
]

View File

@@ -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",
]

View 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"]

View 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",
]

View 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"]

View File

@@ -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",
]

View 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 均视为封禁)
# 对于 Kiro403 表示账户权限被拒绝(无论是封禁、权限不足还是其他原因),
# 都应该标记为异常状态,以便管理员及时处理
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",
]

View File

@@ -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")

View File

@@ -0,0 +1,45 @@
"""Vertex AI 常量配置。
从 transport.py 迁移,集中管理 Vertex AI 模型格式映射和 region 配置。
"""
from __future__ import annotations
# Vertex AI 模型前缀到 API 格式的映射
# 用于 provider_type=vertex_ai 时,根据模型名动态确定实际的请求/响应格式
# 格式:前缀 -> endpoint signaturefamily: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"

View File

@@ -0,0 +1,451 @@
"""Vertex AI provider plugin — 统一注册入口。
注册 Vertex AI 对各通用 registry / capability registry 的 hooks
- Transport Hook (URL 构建Gemini 走 Express modeClaude 走 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"]

View File

@@ -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"
# 构建 URLglobal 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

View 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 失效(不停用 keyaccess 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

View 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"]

View 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),
}

View 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,
)

View 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",
]

View 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)

View 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",
]

View 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

View 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",
]

View 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",
]

View 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",
]

View 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),)

View 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

View File

@@ -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",
]

View 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",
]

View File

@@ -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())

View File

@@ -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())

View File

@@ -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_typeFree 账号优先调度)"
@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())

View File

@@ -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_typeoauth_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())

View File

@@ -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())

View File

@@ -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())

View File

@@ -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())

View File

@@ -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_typePlus/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())

View File

@@ -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())

View File

@@ -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())

View File

@@ -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())

View 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",
]

View File

@@ -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())

View File

@@ -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_typeTeam 账号优先调度)"
@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())

View 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",
]

View 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),
)

View 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)

View 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))

View 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])

View 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

View File

@@ -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()

View File

@@ -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"]

View File

@@ -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",
]

View 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

View 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