Files
Aether/_deprecated_py_src/services/provider/adapters/antigravity/client.py
fawney19 1d9c77522a 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)
2026-04-03 16:26:16 +08:00

628 lines
22 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""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",
]