"""Antigravity provider plugin — 统一注册入口。 将 Antigravity 对各通用 registry / capability registry 的注册集中在一个文件中： - Envelope (v1internal 信封) - Transport Hook (URL 构建) - Auth Enricher (OAuth enrichment) - Model Fetcher (模型获取) - Provider Format Capability（跨格式变体） 新增 provider 时参照此文件创建对应的 plugin.py 即可。 """ from __future__ import annotations import time from typing import Any from urllib.parse import urlencode from src.core.logger import logger from src.services.provider.adapters.antigravity.constants import V1INTERNAL_PATH_TEMPLATE from src.services.provider.adapters.antigravity.url_availability import url_availability from src.services.provider.request_context import set_selected_base_url # --------------------------------------------------------------------------- # Transport Hook # --------------------------------------------------------------------------- def build_antigravity_url( endpoint: Any, *, is_stream: bool, effective_query_params: dict[str, Any], **_kwargs: Any, ) -> str: """构建 Antigravity v1internal URL。 使用 url_availability 选择最优端点，构建 v1internal:generateContent 或 v1internal:streamGenerateContent URL。 """ ordered_urls = url_availability.get_ordered_urls(prefer_daily=True) base_url = ordered_urls[0] if ordered_urls else endpoint.base_url # 存入 contextvars（供后续 Handler 层 envelope 获取） set_selected_base_url(str(base_url) if base_url is not None else None) action = "streamGenerateContent" if is_stream else "generateContent" path = V1INTERNAL_PATH_TEMPLATE.format(action=action) query_params = dict(effective_query_params) # v1internal 流式请求同样支持 ?alt=sse if is_stream: query_params.setdefault("alt", "sse") # 移除 v1internal 不支持的查询参数 query_params.pop("beta", None) url = f"{str(base_url).rstrip('/')}{path}" if query_params: query_string = urlencode(query_params, doseq=True) if query_string: url = f"{url}?{query_string}" return url # --------------------------------------------------------------------------- # Auth Enricher # --------------------------------------------------------------------------- async def enrich_antigravity( auth_config: dict[str, Any], token_response: dict[str, Any], access_token: str, proxy_config: dict[str, Any] | None, ) -> dict[str, Any]: """Antigravity auth_config enrichment. 1. 通过 Google userinfo API 获取 email 2. 通过 loadCodeAssist 获取 project_id / tier 3. 未激活账号尝试 onboardUser 4. fallback 到随机 project_id """ from src.core.provider_oauth_utils import fetch_google_email from src.services.provider.adapters.antigravity.client import ( extract_project_id, extract_tier_id, generate_fallback_project_id, load_code_assist, onboard_user, ) # Email（仅在缺失时获取） if not auth_config.get("email"): email = await fetch_google_email( access_token, proxy_config=proxy_config, timeout_seconds=10.0, ) if email: auth_config["email"] = email # Project ID + Tier（需要 loadCodeAssist 时一起获取） need_project = not auth_config.get("project_id") # Tier 每次 enrich 都重新获取（确保归一化为 Free/Pro/Ultra） need_tier = True if need_project or need_tier: project_id = auth_config.get("project_id", "") try: code_assist = await load_code_assist(access_token, proxy_config=proxy_config) # 提取 tier 信息（对齐 sub2api：优先 paidTier，fallback currentTier） tier_str = _extract_tier_from_code_assist(code_assist) if tier_str: auth_config["tier"] = tier_str logger.info("[enrich] Antigravity tier: {}", tier_str) if need_project: project_id = extract_project_id(code_assist) # 未激活：尝试 onboardUser if not project_id and code_assist.get("allowedTiers"): tier_id = extract_tier_id(code_assist) logger.info("[enrich] Antigravity onboardUser tier={}", tier_id) project_id = await onboard_user( access_token, tier_id=tier_id, proxy_config=proxy_config, ) except Exception as e: logger.warning("[enrich] Antigravity loadCodeAssist/onboardUser 失败: {}", e) if need_project: if project_id: auth_config["project_id"] = project_id logger.info("[enrich] Antigravity project_id: {}", project_id[:8] + "...") else: fallback = generate_fallback_project_id() auth_config["project_id"] = fallback logger.info("[enrich] Antigravity 随机 project_id fallback: {}", fallback) return auth_config def _extract_tier_from_code_assist(code_assist: dict[str, Any]) -> str: """从 loadCodeAssist 响应中提取用户层级，返回 Free/Pro/Ultra。 对齐 sub2api/CLIProxyAPI：优先 paidTier，fallback currentTier。 兼容 tier 为字符串或 {"id": "...", "tierType": "..."} 两种格式。 """ # 优先 paidTier（付费订阅级别） paid_tier = code_assist.get("paidTier") tier = _normalize_tier(_extract_tier_raw(paid_tier)) if tier: return tier # fallback currentTier current_tier = code_assist.get("currentTier") tier = _normalize_tier(_extract_tier_raw(current_tier)) if tier: return tier return "Free" def _extract_tier_raw(tier_obj: Any) -> str: """从 tier 对象中提取原始标识，兼容字符串和字典两种格式。""" if isinstance(tier_obj, str) and tier_obj.strip(): return tier_obj.strip() if isinstance(tier_obj, dict): for key in ("id", "tierType"): val = tier_obj.get(key) if isinstance(val, str) and val.strip(): return val.strip() return "" def _normalize_tier(raw: str) -> str: """将上游 tier 标识归一化为 Free/Pro/Ultra。 上游格式示例： - id: "free-tier", "g1-pro-tier", "g1-ultra-tier" - tierType: "FREE", "PAID" """ if not raw: return "" lower = raw.lower() if "ultra" in lower: return "Ultra" if "pro" in lower or "paid" in lower: return "Pro" if "free" in lower or "legacy" in lower: return "Free" return raw # --------------------------------------------------------------------------- # Model Fetcher # --------------------------------------------------------------------------- # Antigravity 内部测试/调试模型，不应暴露给用户 BLOCKED_MODELS: frozenset[str] = frozenset({"chat_23310", "chat_20706"}) async def fetch_models_antigravity( ctx: Any, timeout_seconds: float, ) -> tuple[list[dict], list[str], bool, dict[str, Any] | None]: """Antigravity 模型获取策略。 调用 v1internal:fetchAvailableModels 获取可用模型列表， 解析配额信息，过滤黑名单模型。 """ from src.services.provider.adapters.antigravity.client import ( AntigravityAccountForbiddenException, fetch_available_models, ) auth_config = ctx.auth_config or {} project_id = auth_config.get("project_id") if not isinstance(project_id, str) or not project_id.strip(): return [], ["antigravity: missing auth_config.project_id (please re-auth)"], False, None try: data = await fetch_available_models( ctx.api_key_value, project_id=project_id.strip(), proxy_config=ctx.proxy_config, timeout_seconds=timeout_seconds, ) except AntigravityAccountForbiddenException: # 重新抛出，由上层（keys.py）处理自动停用逻辑 raise except Exception as e: return [], [f"antigravity: fetchAvailableModels error: {e}"], False, None raw_models = data.get("models") if not isinstance(raw_models, dict): return [], ["antigravity: invalid response (missing models)"], False, None models: list[dict] = [] quota_by_model: dict[str, dict[str, Any]] = {} for model_id, model_data in raw_models.items(): if not isinstance(model_id, str) or not model_id.strip(): continue if model_id.strip() in BLOCKED_MODELS: continue if not isinstance(model_data, dict): model_data = {} display_name = model_data.get("displayName") if not isinstance(display_name, str) or not display_name: display_name = model_id models.append( { "id": model_id, "owned_by": "antigravity", "display_name": display_name, "api_format": "gemini:chat", } ) quota_info = model_data.get("quotaInfo") if not isinstance(quota_info, dict): # 没有 quotaInfo 视为配额耗尽 quota_by_model[model_id] = { "remaining_fraction": 0.0, "used_percent": 100.0, } continue remaining = quota_info.get("remainingFraction") reset_time = quota_info.get("resetTime") remaining_fraction: float | None = None try: if remaining is not None: remaining_fraction = float(remaining) except Exception: remaining_fraction = None if remaining_fraction is None: # remainingFraction 缺失视为配额耗尽 payload: dict[str, Any] = { "remaining_fraction": 0.0, "used_percent": 100.0, } if isinstance(reset_time, str) and reset_time.strip(): payload["reset_time"] = reset_time.strip() quota_by_model[model_id] = payload continue used_percent = (1.0 - remaining_fraction) * 100.0 if used_percent < 0: used_percent = 0.0 if used_percent > 100: # noqa: PLR2004 used_percent = 100.0 payload: dict[str, Any] = { "remaining_fraction": remaining_fraction, "used_percent": used_percent, } if isinstance(reset_time, str) and reset_time.strip(): payload["reset_time"] = reset_time.strip() quota_by_model[model_id] = payload upstream_metadata: dict[str, Any] | None = None if quota_by_model: upstream_metadata = { "antigravity": { "updated_at": int(time.time()), "quota_by_model": quota_by_model, } } return models, [], True, upstream_metadata # --------------------------------------------------------------------------- # Export builder # --------------------------------------------------------------------------- _AG_SKIP_KEYS = frozenset( { "access_token", "expires_at", "updated_at", "token_type", "scope", } ) def antigravity_export_builder( auth_config: dict[str, Any], upstream_metadata: dict[str, Any] | None, ) -> dict[str, Any]: """Antigravity 导出：保留 refresh_token / email / project_id / tier。""" return { k: v for k, v in auth_config.items() if k not in _AG_SKIP_KEYS and v is not None and v != "" } # --------------------------------------------------------------------------- # Unified Registration # --------------------------------------------------------------------------- def register_all() -> None: """一次性注册 Antigravity 的所有 hooks 到各通用 registry。""" from src.core.api_format.capabilities import register_provider_behavior_variant from src.core.provider_oauth_utils import register_auth_enricher from src.services.model.upstream_fetcher import UpstreamModelsFetcherRegistry from src.services.provider.adapters.antigravity.envelope import antigravity_v1internal_envelope from src.services.provider.envelope import register_envelope from src.services.provider.export import register_export_builder from src.services.provider.transport import register_transport_hook # Envelope register_envelope("antigravity", "gemini:chat", antigravity_v1internal_envelope) # Backward compat: allow existing endpoints that still use the old signature. register_envelope("antigravity", "gemini:cli", antigravity_v1internal_envelope) register_envelope("antigravity", "", antigravity_v1internal_envelope) # Transport register_transport_hook("antigravity", "gemini:chat", build_antigravity_url) # Backward compat: allow existing endpoints that still use the old signature. register_transport_hook("antigravity", "gemini:cli", build_antigravity_url) # Auth register_auth_enricher("antigravity", enrich_antigravity) # Export register_export_builder("antigravity", antigravity_export_builder) # Model Fetcher UpstreamModelsFetcherRegistry.register( provider_types=["antigravity"], fetcher=fetch_models_antigravity, ) # Provider Format Capability register_provider_behavior_variant("antigravity", cross_format=True)