From 5add230d4c53a178a50e726f318004e23e834694 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torbj=C3=B8rn=20Lindahl?= Date: Sun, 1 Feb 2026 17:55:26 +0100 Subject: [PATCH] feat(providers): add Antigravity provider for unified Claude/Gemini access Implements a new provider that uses Google's Antigravity unified gateway API to access Claude, Gemini, and other models through a single OAuth2-authenticated endpoint. Features: - OAuth2 token management with automatic refresh - Multi-account rotation for rate limit distribution - Support for Claude Opus/Sonnet 4.5 (with/without thinking) - Support for Gemini 2.5/3 models (Pro/Flash variants) - Thinking mode support with configurable tokens - Image processing support - Dual quota pool tracking (antigravity vs gemini-cli) - Gemini-style API request format Authentication: - Reads from ANTIGRAVITY_REFRESH_TOKEN env var (priority) - Falls back to ~/.config/opencode/antigravity-accounts.json - Automatic token refresh with retry logic - Rate limit tracking per account and quota pool Files added: - providers/antigravity.py - Main provider implementation - providers/antigravity_auth.py - OAuth token manager - providers/registries/antigravity.py - Model registry - conf/antigravity_models.json - Model definitions (11 models) - docs/antigravity_provider.md - Setup and usage docs - tests/test_antigravity_provider.py - Unit tests (14 pass) Integration: - Added to provider priority order after ZEN - Registered in server.py with auto-detection - ToS warning logged on first use --- conf/antigravity_models.json | 140 +++++++ docs/antigravity_provider.md | 189 +++++++++ providers/__init__.py | 2 + providers/antigravity.py | 576 ++++++++++++++++++++++++++++ providers/antigravity_auth.py | 299 +++++++++++++++ providers/registries/antigravity.py | 56 +++ providers/registry.py | 15 + providers/shared/provider_type.py | 1 + server.py | 24 +- tests/test_antigravity_provider.py | 311 +++++++++++++++ 10 files changed, 1612 insertions(+), 1 deletion(-) create mode 100644 conf/antigravity_models.json create mode 100644 docs/antigravity_provider.md create mode 100644 providers/antigravity.py create mode 100644 providers/antigravity_auth.py create mode 100644 providers/registries/antigravity.py create mode 100644 tests/test_antigravity_provider.py diff --git a/conf/antigravity_models.json b/conf/antigravity_models.json new file mode 100644 index 0000000..87879d9 --- /dev/null +++ b/conf/antigravity_models.json @@ -0,0 +1,140 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "description": "Antigravity model configurations for the PAL MCP server", + "models": [ + { + "model_name": "claude-opus-4-5-thinking", + "friendly_name": "Claude Opus 4.5 Thinking (Antigravity)", + "aliases": ["antigravity-claude-opus-4-5-thinking", "ag-opus-thinking"], + "context_window": 200000, + "max_output_tokens": 64000, + "supports_extended_thinking": true, + "max_thinking_tokens": 32768, + "supports_images": true, + "intelligence_score": 20, + "temperature_constraint": "range" + }, + { + "model_name": "claude-sonnet-4-5-thinking", + "friendly_name": "Claude Sonnet 4.5 Thinking (Antigravity)", + "aliases": ["antigravity-claude-sonnet-4-5-thinking", "ag-sonnet-thinking"], + "context_window": 200000, + "max_output_tokens": 64000, + "supports_extended_thinking": true, + "max_thinking_tokens": 32768, + "supports_images": true, + "intelligence_score": 19, + "temperature_constraint": "range" + }, + { + "model_name": "claude-sonnet-4-5", + "friendly_name": "Claude Sonnet 4.5 (Antigravity)", + "aliases": ["antigravity-claude-sonnet-4-5", "ag-sonnet"], + "context_window": 200000, + "max_output_tokens": 64000, + "supports_extended_thinking": false, + "supports_images": true, + "intelligence_score": 18, + "temperature_constraint": "range" + }, + { + "model_name": "gemini-3-pro-high", + "friendly_name": "Gemini 3 Pro High (Antigravity)", + "aliases": ["antigravity-gemini-3-pro", "ag-gemini-3-pro"], + "context_window": 1048576, + "max_output_tokens": 65535, + "supports_extended_thinking": true, + "max_thinking_tokens": 65535, + "supports_images": true, + "intelligence_score": 20, + "temperature_constraint": "range" + }, + { + "model_name": "gemini-3-pro-low", + "friendly_name": "Gemini 3 Pro Low (Antigravity)", + "aliases": ["ag-gemini-3-pro-low"], + "context_window": 1048576, + "max_output_tokens": 65535, + "supports_extended_thinking": true, + "max_thinking_tokens": 16384, + "supports_images": true, + "intelligence_score": 19, + "temperature_constraint": "range" + }, + { + "model_name": "gemini-3-flash", + "friendly_name": "Gemini 3 Flash (Antigravity)", + "aliases": ["antigravity-gemini-3-flash", "ag-gemini-3-flash"], + "context_window": 1048576, + "max_output_tokens": 65536, + "supports_extended_thinking": true, + "max_thinking_tokens": 32768, + "supports_images": true, + "intelligence_score": 17, + "temperature_constraint": "range" + }, + { + "model_name": "gemini-3-flash-preview", + "friendly_name": "Gemini 3 Flash Preview (Gemini CLI)", + "aliases": ["ag-gemini-3-flash-preview"], + "context_window": 1048576, + "max_output_tokens": 65536, + "supports_extended_thinking": true, + "max_thinking_tokens": 32768, + "supports_images": true, + "intelligence_score": 17, + "temperature_constraint": "range", + "quota_pool": "gemini-cli" + }, + { + "model_name": "gemini-3-pro-preview", + "friendly_name": "Gemini 3 Pro Preview (Gemini CLI)", + "aliases": ["ag-gemini-3-pro-preview"], + "context_window": 1048576, + "max_output_tokens": 65535, + "supports_extended_thinking": true, + "max_thinking_tokens": 65535, + "supports_images": true, + "intelligence_score": 19, + "temperature_constraint": "range", + "quota_pool": "gemini-cli" + }, + { + "model_name": "gemini-2.5-pro", + "friendly_name": "Gemini 2.5 Pro (Gemini CLI)", + "aliases": ["ag-gemini-2.5-pro"], + "context_window": 1048576, + "max_output_tokens": 65536, + "supports_extended_thinking": true, + "max_thinking_tokens": 32768, + "supports_images": true, + "intelligence_score": 18, + "temperature_constraint": "range", + "quota_pool": "gemini-cli" + }, + { + "model_name": "gemini-2.5-flash", + "friendly_name": "Gemini 2.5 Flash (Gemini CLI)", + "aliases": ["ag-gemini-2.5-flash"], + "context_window": 1048576, + "max_output_tokens": 65536, + "supports_extended_thinking": true, + "max_thinking_tokens": 16384, + "supports_images": true, + "intelligence_score": 16, + "temperature_constraint": "range", + "quota_pool": "gemini-cli" + }, + { + "model_name": "gpt-oss-120b-medium", + "friendly_name": "GPT-OSS 120B Medium (Antigravity)", + "aliases": ["ag-gpt-oss"], + "context_window": 128000, + "max_output_tokens": 16384, + "supports_extended_thinking": false, + "supports_images": false, + "intelligence_score": 15, + "temperature_constraint": "range" + } + ] +} diff --git a/docs/antigravity_provider.md b/docs/antigravity_provider.md new file mode 100644 index 0000000..1d8db02 --- /dev/null +++ b/docs/antigravity_provider.md @@ -0,0 +1,189 @@ +# Antigravity Provider + +The Antigravity provider enables access to Claude, Gemini, and other models through Google's unified gateway API. This provides an alternative way to access high-quality models using Google OAuth2 authentication instead of individual API keys. + +## Important Warning + +> **Terms of Service Risk**: Using the Antigravity provider may violate Google's Terms of Service. Users have reported account bans when using this approach. +> +> **High-Risk Scenarios**: +> - Fresh Google accounts have a very high chance of getting banned +> - New accounts with Pro/Ultra subscriptions are frequently flagged +> +> **By using this provider, you acknowledge and accept all associated risks.** + +## Available Models + +### Antigravity Quota Pool + +These models use the Antigravity unified gateway quota: + +| Model ID | Description | Thinking Support | +|----------|-------------|------------------| +| `claude-opus-4-5-thinking` | Claude Opus 4.5 with extended thinking | Yes (up to 32K tokens) | +| `claude-sonnet-4-5-thinking` | Claude Sonnet 4.5 with extended thinking | Yes (up to 32K tokens) | +| `claude-sonnet-4-5` | Claude Sonnet 4.5 standard | No | +| `gemini-3-pro-high` | Gemini 3 Pro with high thinking budget | Yes (up to 65K tokens) | +| `gemini-3-pro-low` | Gemini 3 Pro with low thinking budget | Yes (up to 16K tokens) | +| `gemini-3-flash` | Gemini 3 Flash | Yes (up to 32K tokens) | +| `gpt-oss-120b-medium` | GPT-OSS 120B | No | + +### Gemini CLI Quota Pool + +These models use a separate Gemini CLI quota: + +| Model ID | Description | Thinking Support | +|----------|-------------|------------------| +| `gemini-3-flash-preview` | Gemini 3 Flash Preview | Yes | +| `gemini-3-pro-preview` | Gemini 3 Pro Preview | Yes | +| `gemini-2.5-pro` | Gemini 2.5 Pro | Yes | +| `gemini-2.5-flash` | Gemini 2.5 Flash | Yes | + +## Setup + +### Option 1: OpenCode Integration (Recommended) + +If you use OpenCode with the `opencode-antigravity-auth` plugin, your accounts are automatically detected from: + +``` +~/.config/opencode/antigravity-accounts.json +``` + +No additional configuration is needed. The PAL MCP server will automatically discover and use these accounts. + +### Option 2: Manual Refresh Token + +Set the refresh token directly via environment variable: + +```bash +export ANTIGRAVITY_REFRESH_TOKEN="your_oauth_refresh_token" +export ANTIGRAVITY_PROJECT_ID="your_gcp_project_id" # Optional +``` + +### Option 3: Custom Accounts File + +Create `~/.config/opencode/antigravity-accounts.json` manually: + +```json +{ + "version": 3, + "accounts": [ + { + "email": "your.email@gmail.com", + "refreshToken": "1//your_refresh_token_here", + "projectId": "your-gcp-project-id", + "enabled": true + } + ], + "activeIndex": 0 +} +``` + +## Environment Variables + +| Variable | Required | Description | +|----------|----------|-------------| +| `ANTIGRAVITY_REFRESH_TOKEN` | No* | OAuth2 refresh token (overrides accounts file) | +| `ANTIGRAVITY_PROJECT_ID` | No | Google Cloud project ID (defaults to `rising-fact-p41fc`) | +| `ANTIGRAVITY_BASE_URL` | No | API endpoint (defaults to production `cloudcode-pa.googleapis.com`) | + +*Either `ANTIGRAVITY_REFRESH_TOKEN` or a valid accounts file at `~/.config/opencode/antigravity-accounts.json` is required. + +## Features + +### Multi-Account Rotation + +When multiple accounts are configured, the provider automatically: +- Rotates between accounts when one is rate-limited +- Tracks per-model rate limits separately +- Preserves account preferences for cache efficiency + +### Extended Thinking + +For models that support thinking (Claude and Gemini 3), you can control the thinking budget: + +```python +# Thinking modes: minimal, low, medium, high, max +response = provider.generate_content( + prompt="Explain quantum computing", + model_name="claude-opus-4-5-thinking", + thinking_mode="high" # Uses 67% of max thinking budget +) +``` + +| Mode | Budget Percentage | +|------|-------------------| +| minimal | 0.5% | +| low | 8% | +| medium | 33% (default) | +| high | 67% | +| max | 100% | + +### Image Support + +Models with vision capability support image inputs: + +```python +response = provider.generate_content( + prompt="What's in this image?", + model_name="gemini-3-flash", + images=["/path/to/image.png"] +) +``` + +## Troubleshooting + +### "No Antigravity accounts configured" + +Ensure either: +1. `ANTIGRAVITY_REFRESH_TOKEN` environment variable is set, or +2. `~/.config/opencode/antigravity-accounts.json` exists with valid accounts + +### "All Antigravity accounts are rate limited" + +The provider has exhausted all available accounts. Either: +1. Wait for rate limits to reset (usually 60 seconds to a few minutes) +2. Add more accounts to the accounts file + +### "Token refresh failed" + +Your refresh token may have expired. Re-authenticate: +1. If using OpenCode: `opencode auth login` +2. If manual: Obtain a new refresh token via OAuth2 flow + +### 403 Permission Denied + +Your Google Cloud project may not have the required APIs enabled: +1. Go to [Google Cloud Console](https://console.cloud.google.com/) +2. Enable the **Gemini for Google Cloud API** (`cloudaicompanion.googleapis.com`) + +## API Format + +The Antigravity API uses Gemini-style request format internally: + +```json +{ + "project": "your-project-id", + "model": "claude-opus-4-5-thinking", + "request": { + "contents": [ + {"role": "user", "parts": [{"text": "Your prompt"}]} + ], + "generationConfig": { + "temperature": 0.7, + "maxOutputTokens": 8192, + "thinkingConfig": { + "thinkingBudget": 8000, + "includeThoughts": true + } + } + } +} +``` + +This transformation is handled automatically by the provider. + +## Related Resources + +- [opencode-antigravity-auth](https://github.com/NoeFabris/opencode-antigravity-auth) - OpenCode plugin for Antigravity authentication +- [Antigravity API Specification](https://github.com/NoeFabris/opencode-antigravity-auth/blob/dev/docs/ANTIGRAVITY_API_SPEC.md) - Detailed API documentation diff --git a/providers/__init__.py b/providers/__init__.py index b56f69d..c7fc30e 100644 --- a/providers/__init__.py +++ b/providers/__init__.py @@ -1,5 +1,6 @@ """Model provider abstractions for supporting multiple AI providers.""" +from .antigravity import AntigravityProvider from .azure_openai import AzureOpenAIProvider from .base import ModelProvider from .gemini import GeminiModelProvider @@ -15,6 +16,7 @@ __all__ = [ "ModelResponse", "ModelCapabilities", "ModelProviderRegistry", + "AntigravityProvider", "AzureOpenAIProvider", "GeminiModelProvider", "OpenAIModelProvider", diff --git a/providers/antigravity.py b/providers/antigravity.py new file mode 100644 index 0000000..2c40663 --- /dev/null +++ b/providers/antigravity.py @@ -0,0 +1,576 @@ +"""Antigravity model provider implementation. + +Antigravity is Google's unified gateway API for accessing multiple AI models +(Claude, Gemini, GPT-OSS) through a single Gemini-style interface. + +WARNING: Using this provider may violate Google's Terms of Service. +See docs/antigravity_provider.md for important information about risks. +""" + +from __future__ import annotations + +import base64 +import json +import logging +import re +import uuid +from typing import TYPE_CHECKING, Any, ClassVar + +import httpx + +from utils.env import get_env +from utils.image_utils import validate_image + +from .antigravity_auth import AntigravityTokenManager +from .base import ModelProvider +from .registries.antigravity import AntigravityModelRegistry +from .shared import ModelCapabilities, ModelResponse, ProviderType + +if TYPE_CHECKING: + from tools.models import ToolModelCategory + +logger = logging.getLogger(__name__) + +# Antigravity API endpoints +PRODUCTION_ENDPOINT = "https://cloudcode-pa.googleapis.com" +DAILY_ENDPOINT = "https://daily-cloudcode-pa.sandbox.googleapis.com" + +# API path for content generation +GENERATE_CONTENT_PATH = "/v1internal:generateContent" +STREAM_GENERATE_CONTENT_PATH = "/v1internal:streamGenerateContent" + +# ToS warning message +TOS_WARNING = """ +================================================================================ +ANTIGRAVITY PROVIDER WARNING + +Using the Antigravity provider may violate Google's Terms of Service. +Users have reported account bans when using this approach. + +HIGH-RISK SCENARIOS: +- Fresh Google accounts have a very high chance of getting banned +- New accounts with Pro/Ultra subscriptions are frequently flagged + +By using this provider, you acknowledge and accept all associated risks. +See docs/antigravity_provider.md for more information. +================================================================================ +""" + + +class AntigravityProvider(ModelProvider): + """Provider for Google's Antigravity unified gateway API. + + Antigravity provides access to Claude, Gemini, and other models through a + single Gemini-style API interface. This provider handles: + + * OAuth2 authentication with refresh token management + * Multi-account rotation on rate limits + * Request transformation to Gemini-style format + * Extended thinking support for compatible models + * Image processing for vision requests + + Configuration: + ANTIGRAVITY_REFRESH_TOKEN: OAuth2 refresh token (env var override) + ANTIGRAVITY_PROJECT_ID: Google Cloud project ID + ANTIGRAVITY_BASE_URL: Custom endpoint URL (defaults to production) + """ + + # Model registry for capability metadata + _registry: ClassVar[AntigravityModelRegistry | None] = None + MODEL_CAPABILITIES: ClassVar[dict[str, ModelCapabilities]] = {} + + # Thinking mode budgets (percentage of max_thinking_tokens) + THINKING_BUDGETS = { + "minimal": 0.005, + "low": 0.08, + "medium": 0.33, + "high": 0.67, + "max": 1.0, + } + + # Track whether warning has been shown this session + _warning_shown: ClassVar[bool] = False + + def __init__(self, api_key: str = "", **kwargs) -> None: + """Initialize Antigravity provider. + + Args: + api_key: Not used (authentication via OAuth2 refresh tokens) + **kwargs: Additional configuration options + """ + # Initialize registry + if AntigravityProvider._registry is None: + AntigravityProvider._registry = AntigravityModelRegistry() + AntigravityProvider.MODEL_CAPABILITIES = dict(self._registry.model_map) + logger.info("Antigravity provider loaded %d models", len(self._registry.model_map)) + + super().__init__(api_key, **kwargs) + + # Initialize token manager + self._token_manager = AntigravityTokenManager() + + # HTTP client for API requests + self._http_client: httpx.Client | None = None + + # Base URL configuration + self._base_url = get_env("ANTIGRAVITY_BASE_URL") or PRODUCTION_ENDPOINT + + # Show warning on first use + if not AntigravityProvider._warning_shown: + logger.warning(TOS_WARNING) + AntigravityProvider._warning_shown = True + + self._invalidate_capability_cache() + + @property + def http_client(self) -> httpx.Client: + """Lazy initialization of HTTP client.""" + if self._http_client is None: + timeout = httpx.Timeout( + connect=30.0, + read=600.0, # 10 minutes for long responses + write=60.0, + pool=600.0, + ) + self._http_client = httpx.Client(timeout=timeout, follow_redirects=True) + return self._http_client + + # ------------------------------------------------------------------ + # Provider identity + # ------------------------------------------------------------------ + + def get_provider_type(self) -> ProviderType: + """Return the provider type.""" + return ProviderType.ANTIGRAVITY + + # ------------------------------------------------------------------ + # Capability surface + # ------------------------------------------------------------------ + + def _lookup_capabilities( + self, + canonical_name: str, + requested_name: str | None = None, + ) -> ModelCapabilities | None: + """Fetch capabilities from the registry.""" + if self._registry: + return self._registry.get_capabilities(canonical_name) + return None + + def get_all_model_capabilities(self) -> dict[str, ModelCapabilities]: + """Return registry-backed capabilities.""" + if not self._registry: + return {} + return dict(self._registry.model_map) + + def list_models( + self, + *, + respect_restrictions: bool = True, + include_aliases: bool = True, + lowercase: bool = False, + unique: bool = False, + ) -> list[str]: + """Return available Antigravity model names.""" + if not self._registry: + return [] + + from utils.model_restrictions import get_restriction_service + + restriction_service = get_restriction_service() if respect_restrictions else None + allowed_configs: dict[str, ModelCapabilities] = {} + + for model_name in self._registry.list_models(): + config = self._registry.resolve(model_name) + if not config: + continue + + if restriction_service: + if not restriction_service.is_allowed(self.get_provider_type(), model_name): + continue + + allowed_configs[model_name] = config + + if not allowed_configs: + return [] + + return ModelCapabilities.collect_model_names( + allowed_configs, + include_aliases=include_aliases, + lowercase=lowercase, + unique=unique, + ) + + def _resolve_model_name(self, model_name: str) -> str: + """Resolve aliases to canonical model names.""" + if self._registry: + config = self._registry.resolve(model_name) + if config and config.model_name != model_name: + logger.debug("Resolved Antigravity alias '%s' to '%s'", model_name, config.model_name) + return config.model_name + return model_name + + # ------------------------------------------------------------------ + # Content generation + # ------------------------------------------------------------------ + + def generate_content( + self, + prompt: str, + model_name: str, + system_prompt: str | None = None, + temperature: float = 0.7, + max_output_tokens: int | None = None, + thinking_mode: str = "medium", + images: list[str] | None = None, + **kwargs, + ) -> ModelResponse: + """Generate content using the Antigravity API. + + Args: + prompt: User prompt to send to the model + model_name: Model identifier (e.g., 'claude-opus-4-5-thinking') + system_prompt: Optional system instructions + temperature: Sampling temperature (0.0-2.0) + max_output_tokens: Maximum tokens in response + thinking_mode: Thinking budget level for thinking models + images: Optional list of image paths or data URLs + **kwargs: Additional parameters + + Returns: + ModelResponse with generated content and metadata + """ + # Validate and resolve model name + self.validate_parameters(model_name, temperature) + capabilities = self.get_capabilities(model_name) + resolved_model = self._resolve_model_name(model_name) + + # Get authentication + access_token, project_id, fingerprint_headers = self._token_manager.get_access_token(model=resolved_model) + + # Build request + request_body = self._build_request( + prompt=prompt, + model_name=resolved_model, + project_id=project_id, + system_prompt=system_prompt, + temperature=temperature, + max_output_tokens=max_output_tokens, + thinking_mode=thinking_mode, + images=images, + capabilities=capabilities, + ) + + # Build headers + headers = { + "Authorization": f"Bearer {access_token}", + "Content-Type": "application/json", + **fingerprint_headers, + } + + # Retry logic + max_retries = 4 + retry_delays = [1.0, 3.0, 5.0, 8.0] + attempt_counter = {"value": 0} + + def _attempt() -> ModelResponse: + attempt_counter["value"] += 1 + + url = f"{self._base_url}{GENERATE_CONTENT_PATH}" + logger.debug("Antigravity request to %s for model %s", url, resolved_model) + + response = self.http_client.post( + url, + json=request_body, + headers=headers, + ) + + # Handle errors + if response.status_code != 200: + self._handle_error_response(response, resolved_model) + + return self._parse_response(response.json(), resolved_model, capabilities, thinking_mode) + + try: + return self._run_with_retries( + operation=_attempt, + max_attempts=max_retries, + delays=retry_delays, + log_prefix=f"Antigravity API ({resolved_model})", + ) + except Exception as exc: + attempts = max(attempt_counter["value"], 1) + error_msg = f"Antigravity API error for {resolved_model} after {attempts} attempt(s): {exc}" + raise RuntimeError(error_msg) from exc + + def _build_request( + self, + prompt: str, + model_name: str, + project_id: str, + system_prompt: str | None = None, + temperature: float = 0.7, + max_output_tokens: int | None = None, + thinking_mode: str = "medium", + images: list[str] | None = None, + capabilities: ModelCapabilities | None = None, + ) -> dict[str, Any]: + """Build the Antigravity API request body.""" + # Build contents array (Gemini-style format) + contents = [] + + # Build user message parts + user_parts = [] + + # Add text prompt + user_parts.append({"text": prompt}) + + # Add images if provided + if images and capabilities and capabilities.supports_images: + for image_path in images: + try: + image_part = self._process_image(image_path) + if image_part: + user_parts.append(image_part) + except Exception as e: + logger.warning("Failed to process image %s: %s", image_path, e) + elif images: + logger.warning("Model %s does not support images, ignoring %d image(s)", model_name, len(images)) + + contents.append({"role": "user", "parts": user_parts}) + + # Build generation config + generation_config: dict[str, Any] = { + "temperature": temperature, + } + + if max_output_tokens: + generation_config["maxOutputTokens"] = max_output_tokens + + # Add thinking config for thinking-capable models + if capabilities and capabilities.supports_extended_thinking and thinking_mode in self.THINKING_BUDGETS: + max_thinking = capabilities.max_thinking_tokens or 8192 + thinking_budget = int(max_thinking * self.THINKING_BUDGETS[thinking_mode]) + generation_config["thinkingConfig"] = { + "thinkingBudget": thinking_budget, + "includeThoughts": True, + } + # Ensure maxOutputTokens > thinkingBudget + if not max_output_tokens or max_output_tokens <= thinking_budget: + generation_config["maxOutputTokens"] = thinking_budget + 10000 + + # Build system instruction (must be object with parts, not string) + system_instruction = None + if system_prompt: + system_instruction = {"parts": [{"text": system_prompt}]} + + # Build full request + request: dict[str, Any] = { + "project": project_id, + "model": model_name, + "request": { + "contents": contents, + "generationConfig": generation_config, + }, + "userAgent": "antigravity", + "requestId": f"pal-{uuid.uuid4().hex[:12]}", + } + + if system_instruction: + request["request"]["systemInstruction"] = system_instruction + + return request + + def _parse_response( + self, + data: dict[str, Any], + model_name: str, + capabilities: ModelCapabilities | None, + thinking_mode: str, + ) -> ModelResponse: + """Parse the Antigravity API response.""" + response_data = data.get("response", data) + + # Extract content from candidates + content = "" + thinking_content = "" + candidates = response_data.get("candidates", []) + + if candidates: + candidate = candidates[0] + candidate_content = candidate.get("content", {}) + parts = candidate_content.get("parts", []) + + for part in parts: + if part.get("thought"): + # This is thinking content + thinking_content += part.get("text", "") + elif "thoughtSignature" in part and "text" in part: + # Gemini thinking with signature + thinking_content += part.get("text", "") + elif "text" in part: + content += part.get("text", "") + + # Extract usage metadata + usage_metadata = response_data.get("usageMetadata", {}) + usage = { + "input_tokens": usage_metadata.get("promptTokenCount", 0), + "output_tokens": usage_metadata.get("candidatesTokenCount", 0), + "total_tokens": usage_metadata.get("totalTokenCount", 0), + } + + # Add thinking tokens if available + if "thoughtsTokenCount" in usage_metadata: + usage["thinking_tokens"] = usage_metadata["thoughtsTokenCount"] + + # Extract finish reason + finish_reason = "UNKNOWN" + if candidates: + finish_reason = candidates[0].get("finishReason", "STOP") + + return ModelResponse( + content=content, + usage=usage, + model_name=model_name, + friendly_name="Antigravity", + provider=ProviderType.ANTIGRAVITY, + metadata={ + "thinking_mode": thinking_mode if capabilities and capabilities.supports_extended_thinking else None, + "thinking_content": thinking_content if thinking_content else None, + "finish_reason": finish_reason, + "model_version": response_data.get("modelVersion"), + "response_id": response_data.get("responseId"), + "trace_id": data.get("traceId"), + }, + ) + + def _handle_error_response(self, response: httpx.Response, model_name: str) -> None: + """Handle error responses from the API.""" + try: + error_data = response.json() + error = error_data.get("error", {}) + error_code = error.get("code", response.status_code) + error_message = error.get("message", response.text) + error_status = error.get("status", "UNKNOWN") + + # Handle rate limiting + if error_code == 429: + # Extract retry delay from response + retry_delay = 60.0 + details = error.get("details", []) + for detail in details: + if detail.get("@type", "").endswith("RetryInfo"): + delay_str = detail.get("retryDelay", "60s") + # Parse delay string (e.g., "3.957525076s") + match = re.match(r"([\d.]+)s?", delay_str) + if match: + retry_delay = float(match.group(1)) + break + + # Mark current account as rate limited + self._token_manager.mark_rate_limited(model_name, retry_delay) + + raise RuntimeError(f"Rate limited (429): {error_message}. Retry after {retry_delay}s") + + # Re-raise with formatted message + raise RuntimeError(f"Antigravity API error {error_code} ({error_status}): {error_message}") + + except json.JSONDecodeError: + raise RuntimeError(f"Antigravity API error {response.status_code}: {response.text}") + + def _process_image(self, image_path: str) -> dict[str, Any] | None: + """Process an image for the Antigravity API.""" + try: + image_bytes, mime_type = validate_image(image_path) + + if image_path.startswith("data:"): + # Extract base64 data from data URL + _, data = image_path.split(",", 1) + return {"inlineData": {"mimeType": mime_type, "data": data}} + else: + # Encode file bytes + image_data = base64.b64encode(image_bytes).decode() + return {"inlineData": {"mimeType": mime_type, "data": image_data}} + + except ValueError as e: + logger.warning("Image validation failed: %s", e) + return None + except Exception as e: + logger.error("Error processing image %s: %s", image_path, e) + return None + + def _is_error_retryable(self, error: Exception) -> bool: + """Determine if an error should be retried.""" + error_str = str(error).lower() + + # Rate limits should be retried (we'll rotate accounts) + if "429" in error_str or "rate limit" in error_str: + return True + + # Check for retryable indicators + retryable_indicators = [ + "timeout", + "connection", + "network", + "temporary", + "unavailable", + "retry", + "500", + "502", + "503", + "504", + ] + + return any(indicator in error_str for indicator in retryable_indicators) + + # ------------------------------------------------------------------ + # Preferred model selection + # ------------------------------------------------------------------ + + def get_preferred_model(self, category: ToolModelCategory, allowed_models: list[str]) -> str | None: + """Get preferred model for a category from allowed models.""" + from tools.models import ToolModelCategory + + if not allowed_models: + return None + + capability_map = self.get_all_model_capabilities() + + def find_best(candidates: list[str]) -> str | None: + return sorted(candidates, reverse=True)[0] if candidates else None + + if category == ToolModelCategory.EXTENDED_REASONING: + # Prefer thinking-capable models + thinking_models = [ + m for m in allowed_models if m in capability_map and capability_map[m].supports_extended_thinking + ] + if thinking_models: + # Prefer Claude Opus, then Claude Sonnet, then Gemini + opus = [m for m in thinking_models if "opus" in m] + if opus: + return find_best(opus) + sonnet = [m for m in thinking_models if "sonnet" in m] + if sonnet: + return find_best(sonnet) + return find_best(thinking_models) + + elif category == ToolModelCategory.FAST_RESPONSE: + # Prefer Flash models + flash_models = [m for m in allowed_models if "flash" in m] + if flash_models: + return find_best(flash_models) + + # Default: prefer models by capability rank + return find_best(allowed_models) + + # ------------------------------------------------------------------ + # Resource cleanup + # ------------------------------------------------------------------ + + def close(self) -> None: + """Clean up resources.""" + if self._http_client: + self._http_client.close() + self._http_client = None + if self._token_manager: + self._token_manager.close() diff --git a/providers/antigravity_auth.py b/providers/antigravity_auth.py new file mode 100644 index 0000000..5e85c7a --- /dev/null +++ b/providers/antigravity_auth.py @@ -0,0 +1,299 @@ +"""OAuth2 token manager for Antigravity provider. + +This module handles authentication with Google's Antigravity unified gateway API, +including refresh token management and multi-account rotation. +""" + +from __future__ import annotations + +import json +import logging +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import httpx + +from utils.env import get_env + +logger = logging.getLogger(__name__) + +# OAuth2 endpoints +TOKEN_URL = "https://oauth2.googleapis.com/token" + +# Google OAuth2 client credentials (from Antigravity/Gemini CLI) +# These are publicly documented OAuth client credentials +OAUTH_CLIENT_ID = "764086051850-6qr4p6gpi6hn506pt8ejuq83di341hur.apps.googleusercontent.com" +OAUTH_CLIENT_SECRET = "d-FL95Q19q7MQmFpd7hHD0Ty" + +# Required OAuth2 scopes for Antigravity API access +OAUTH_SCOPES = [ + "https://www.googleapis.com/auth/cloud-platform", + "https://www.googleapis.com/auth/userinfo.email", + "https://www.googleapis.com/auth/userinfo.profile", + "https://www.googleapis.com/auth/cclog", + "https://www.googleapis.com/auth/experimentsandconfigs", +] + +# Default accounts file location +DEFAULT_ACCOUNTS_FILE = Path.home() / ".config" / "opencode" / "antigravity-accounts.json" + + +@dataclass +class AntigravityAccount: + """Represents a single Antigravity OAuth account.""" + + email: str + refresh_token: str + project_id: str + enabled: bool = True + access_token: str | None = None + token_expiry: float = 0.0 + rate_limit_reset_times: dict[str, float] = field(default_factory=dict) + fingerprint: dict[str, Any] = field(default_factory=dict) + + def is_token_valid(self) -> bool: + """Check if the current access token is still valid.""" + if not self.access_token: + return False + # Refresh 60 seconds before expiry + return time.time() < (self.token_expiry - 60) + + def is_rate_limited(self, model: str | None = None) -> bool: + """Check if this account is rate limited for a specific model or globally.""" + now = time.time() * 1000 # Convert to milliseconds + if model: + reset_time = self.rate_limit_reset_times.get(model, 0) + if reset_time > now: + return True + # Check for any active rate limits + return any(reset > now for reset in self.rate_limit_reset_times.values()) + + def set_rate_limited(self, model: str, reset_delay_seconds: float) -> None: + """Mark this account as rate limited for a model.""" + reset_time = (time.time() + reset_delay_seconds) * 1000 + self.rate_limit_reset_times[model] = reset_time + logger.info("Account %s rate limited for %s until %.0f", self.email, model, reset_time) + + +class AntigravityTokenManager: + """Manages OAuth2 tokens for Antigravity API access. + + Supports: + - Loading accounts from ~/.config/opencode/antigravity-accounts.json + - Environment variable override via ANTIGRAVITY_REFRESH_TOKEN + - Token refresh with automatic expiry handling + - Multi-account rotation on rate limits + """ + + def __init__(self, accounts_file: Path | str | None = None) -> None: + self._accounts: list[AntigravityAccount] = [] + self._current_account_index: int = 0 + self._http_client: httpx.Client | None = None + + # Try environment variable first + env_token = get_env("ANTIGRAVITY_REFRESH_TOKEN") + env_project = get_env("ANTIGRAVITY_PROJECT_ID") + + if env_token: + logger.info("Using Antigravity refresh token from environment variable") + self._accounts.append( + AntigravityAccount( + email="env-configured", + refresh_token=env_token, + project_id=env_project or "rising-fact-p41fc", # Default project + ) + ) + else: + # Load from accounts file + accounts_path = Path(accounts_file) if accounts_file else DEFAULT_ACCOUNTS_FILE + self._load_accounts(accounts_path) + + if self._accounts: + logger.info("Antigravity token manager initialized with %d account(s)", len(self._accounts)) + else: + logger.warning("No Antigravity accounts configured") + + @property + def http_client(self) -> httpx.Client: + """Lazy-init HTTP client.""" + if self._http_client is None: + self._http_client = httpx.Client(timeout=30.0) + return self._http_client + + def _load_accounts(self, accounts_path: Path) -> None: + """Load accounts from the antigravity-accounts.json file.""" + if not accounts_path.exists(): + logger.debug("Antigravity accounts file not found at %s", accounts_path) + return + + try: + with open(accounts_path, encoding="utf-8") as f: + data = json.load(f) + + version = data.get("version", 1) + if version < 3: + logger.warning( + "Antigravity accounts file has old format (version %d), may have limited support", version + ) + + for account_data in data.get("accounts", []): + if not account_data.get("enabled", True): + continue + + email = account_data.get("email", "unknown") + refresh_token = account_data.get("refreshToken") + project_id = account_data.get("projectId") or account_data.get("managedProjectId") + + if not refresh_token: + logger.warning("Skipping account %s: no refresh token", email) + continue + + account = AntigravityAccount( + email=email, + refresh_token=refresh_token, + project_id=project_id or "rising-fact-p41fc", + enabled=account_data.get("enabled", True), + rate_limit_reset_times=account_data.get("rateLimitResetTimes", {}), + fingerprint=account_data.get("fingerprint", {}), + ) + self._accounts.append(account) + logger.debug("Loaded Antigravity account: %s (project: %s)", email, project_id) + + except json.JSONDecodeError as e: + logger.error("Failed to parse Antigravity accounts file: %s", e) + except Exception as e: + logger.error("Error loading Antigravity accounts: %s", e) + + def has_accounts(self) -> bool: + """Check if any accounts are configured.""" + return len(self._accounts) > 0 + + def get_account_count(self) -> int: + """Return the number of configured accounts.""" + return len(self._accounts) + + def get_access_token(self, model: str | None = None) -> tuple[str, str, dict[str, str]]: + """Get a valid access token, refreshing if necessary. + + Args: + model: Optional model name for rate limit checking + + Returns: + Tuple of (access_token, project_id, fingerprint_headers) + + Raises: + RuntimeError: If no valid token can be obtained + """ + if not self._accounts: + raise RuntimeError("No Antigravity accounts configured") + + # Try each account starting from current index + attempts = 0 + original_index = self._current_account_index + + while attempts < len(self._accounts): + account = self._accounts[self._current_account_index] + + # Skip rate-limited accounts for this model + if model and account.is_rate_limited(model): + logger.debug("Skipping rate-limited account %s for model %s", account.email, model) + self._rotate_account() + attempts += 1 + continue + + # Refresh token if needed + if not account.is_token_valid(): + try: + self._refresh_token(account) + except Exception as e: + logger.warning("Failed to refresh token for %s: %s", account.email, e) + self._rotate_account() + attempts += 1 + continue + + # Build fingerprint headers + fingerprint_headers = self._build_fingerprint_headers(account) + + return account.access_token, account.project_id, fingerprint_headers + + # All accounts exhausted + self._current_account_index = original_index + raise RuntimeError("All Antigravity accounts are rate limited or have invalid tokens") + + def mark_rate_limited(self, model: str, retry_delay: float = 60.0) -> None: + """Mark the current account as rate limited and rotate to next.""" + if self._accounts: + account = self._accounts[self._current_account_index] + account.set_rate_limited(model, retry_delay) + self._rotate_account() + + def _rotate_account(self) -> None: + """Rotate to the next available account.""" + if len(self._accounts) > 1: + self._current_account_index = (self._current_account_index + 1) % len(self._accounts) + logger.debug("Rotated to account index %d", self._current_account_index) + + def _refresh_token(self, account: AntigravityAccount) -> None: + """Refresh the OAuth2 access token for an account.""" + logger.debug("Refreshing access token for %s", account.email) + + response = self.http_client.post( + TOKEN_URL, + data={ + "client_id": OAUTH_CLIENT_ID, + "client_secret": OAUTH_CLIENT_SECRET, + "refresh_token": account.refresh_token, + "grant_type": "refresh_token", + }, + ) + + if response.status_code != 200: + error_msg = f"Token refresh failed: {response.status_code} - {response.text}" + logger.error(error_msg) + raise RuntimeError(error_msg) + + data = response.json() + account.access_token = data["access_token"] + # Token typically expires in 3600 seconds + expires_in = data.get("expires_in", 3600) + account.token_expiry = time.time() + expires_in + + logger.debug("Token refreshed for %s, expires in %d seconds", account.email, expires_in) + + def _build_fingerprint_headers(self, account: AntigravityAccount) -> dict[str, str]: + """Build fingerprint headers for API requests.""" + fingerprint = account.fingerprint + headers = {} + + if fingerprint.get("userAgent"): + headers["User-Agent"] = fingerprint["userAgent"] + else: + headers["User-Agent"] = "antigravity/1.15.8 linux/amd64" + + if fingerprint.get("apiClient"): + headers["X-Goog-Api-Client"] = fingerprint["apiClient"] + else: + headers["X-Goog-Api-Client"] = "google-cloud-sdk intellij/2024.1" + + # Build client metadata + client_metadata = fingerprint.get("clientMetadata", {}) + metadata = { + "ideType": client_metadata.get("ideType", "CLOUD_SHELL_EDITOR"), + "platform": client_metadata.get("platform", "LINUX"), + "pluginType": client_metadata.get("pluginType", "GEMINI"), + } + headers["Client-Metadata"] = json.dumps(metadata) + + # Quota user for rate limiting + if fingerprint.get("quotaUser"): + headers["X-Goog-Quota-User"] = fingerprint["quotaUser"] + + return headers + + def close(self) -> None: + """Close the HTTP client.""" + if self._http_client: + self._http_client.close() + self._http_client = None diff --git a/providers/registries/antigravity.py b/providers/registries/antigravity.py new file mode 100644 index 0000000..ca26729 --- /dev/null +++ b/providers/registries/antigravity.py @@ -0,0 +1,56 @@ +"""Antigravity model registry for managing model configurations and aliases.""" + +from __future__ import annotations + +from ..shared import ModelCapabilities, ProviderType +from .base import CAPABILITY_FIELD_NAMES, CapabilityModelRegistry + + +class AntigravityModelRegistry(CapabilityModelRegistry): + """Capability registry backed by ``conf/antigravity_models.json``. + + The Antigravity provider accesses Claude, Gemini, and other models through + Google's unified gateway API. Models are split across two quota pools: + + * **antigravity**: Claude models and Gemini 3 variants + * **gemini-cli**: Gemini 2.5 and preview models using Gemini CLI quota + """ + + # Extra fields beyond standard ModelCapabilities + EXTRA_KEYS = {"quota_pool"} + + def __init__(self, config_path: str | None = None) -> None: + super().__init__( + env_var_name="ANTIGRAVITY_MODELS_CONFIG_PATH", + default_filename="antigravity_models.json", + provider=ProviderType.ANTIGRAVITY, + friendly_prefix="Antigravity ({model})", + config_path=config_path, + ) + + def _extra_keys(self) -> set[str]: + """Include quota_pool as a valid entry key.""" + return self.EXTRA_KEYS + + def _finalise_entry(self, entry: dict) -> tuple[ModelCapabilities, dict]: + """Build ModelCapabilities and extract extras like quota_pool.""" + # Extract extras before building capabilities + extras = {} + quota_pool = entry.pop("quota_pool", None) + if quota_pool: + extras["quota_pool"] = quota_pool + + # Build capabilities from remaining fields + filtered = {k: v for k, v in entry.items() if k in CAPABILITY_FIELD_NAMES} + filtered.setdefault("provider", ProviderType.ANTIGRAVITY) + filtered.setdefault("friendly_name", f"Antigravity ({entry.get('model_name', 'unknown')})") + capability = ModelCapabilities(**filtered) + + return capability, extras + + def get_quota_pool(self, model_name: str) -> str: + """Return the quota pool for a model ('antigravity' or 'gemini-cli').""" + extras = self.get_entry(model_name) + if extras and "quota_pool" in extras: + return extras["quota_pool"] + return "antigravity" # Default to antigravity quota pool diff --git a/providers/registry.py b/providers/registry.py index 9e9a2ed..d3a7636 100644 --- a/providers/registry.py +++ b/providers/registry.py @@ -41,6 +41,7 @@ class ModelProviderRegistry: ProviderType.AZURE, # Azure-hosted OpenAI deployments ProviderType.XAI, # Direct X.AI GROK access ProviderType.ZEN, # OpenCode Zen curated models + ProviderType.ANTIGRAVITY, # Antigravity unified gateway (Claude/Gemini via Google OAuth) ProviderType.DIAL, # DIAL unified API access ProviderType.CUSTOM, # Local/self-hosted models ProviderType.OPENROUTER, # Catch-all for cloud models @@ -140,6 +141,19 @@ class ModelProviderRegistry: azure_endpoint=azure_endpoint, api_version=azure_version, ) + elif provider_type == ProviderType.ANTIGRAVITY: + # Antigravity uses OAuth tokens from file or env var + # The provider handles token loading internally + from pathlib import Path + + # Check if refresh token is available from env or accounts file + accounts_file = Path.home() / ".config" / "opencode" / "antigravity-accounts.json" + if not api_key and not accounts_file.exists(): + logging.debug("Antigravity: No refresh token in env and no accounts file found") + return None + + # Initialize provider (it will load tokens internally) + provider = provider_class(api_key=api_key or "") else: if not api_key: return None @@ -341,6 +355,7 @@ class ModelProviderRegistry: ProviderType.OPENROUTER: "OPENROUTER_API_KEY", ProviderType.CUSTOM: "CUSTOM_API_KEY", # Can be empty for providers that don't need auth ProviderType.DIAL: "DIAL_API_KEY", + ProviderType.ANTIGRAVITY: "ANTIGRAVITY_REFRESH_TOKEN", # Uses OAuth refresh token, not API key } env_var = key_mapping.get(provider_type) diff --git a/providers/shared/provider_type.py b/providers/shared/provider_type.py index bb9a533..2cc197e 100644 --- a/providers/shared/provider_type.py +++ b/providers/shared/provider_type.py @@ -16,3 +16,4 @@ class ProviderType(Enum): CUSTOM = "custom" DIAL = "dial" ZEN = "zen" + ANTIGRAVITY = "antigravity" diff --git a/server.py b/server.py index 9a04174..b5f6aca 100644 --- a/server.py +++ b/server.py @@ -392,6 +392,7 @@ def configure_providers(): value = get_env(key) logger.debug(f" {key}: {'[PRESENT]' if value else '[MISSING]'}") from providers import ModelProviderRegistry + from providers.antigravity import AntigravityProvider from providers.azure_openai import AzureOpenAIProvider from providers.custom import CustomProvider from providers.dial import DIALModelProvider @@ -490,6 +491,21 @@ def configure_providers(): else: logger.debug("OpenCode Zen API key is placeholder value") + # Check for Antigravity OAuth accounts + antigravity_token = get_env("ANTIGRAVITY_REFRESH_TOKEN") + antigravity_accounts_file = Path.home() / ".config" / "opencode" / "antigravity-accounts.json" + has_antigravity = False + if antigravity_token: + valid_providers.append("Antigravity (OAuth)") + has_antigravity = True + logger.info("Antigravity refresh token found - Claude/Gemini models via unified gateway") + elif antigravity_accounts_file.exists(): + valid_providers.append("Antigravity (OAuth)") + has_antigravity = True + logger.info("Antigravity accounts file found - Claude/Gemini models via unified gateway") + else: + logger.debug("Antigravity: No refresh token or accounts file found") + # Check for custom API endpoint (Ollama, vLLM, etc.) custom_url = get_env("CUSTOM_API_URL") if custom_url: @@ -551,7 +567,13 @@ def configure_providers(): registered_providers.append(ProviderType.ZEN.value) logger.debug(f"Registered provider: {ProviderType.ZEN.value}") - # 4. OpenRouter last (catch-all for everything else) + # 4. Antigravity (Google unified gateway for Claude/Gemini) + if has_antigravity: + ModelProviderRegistry.register_provider(ProviderType.ANTIGRAVITY, AntigravityProvider) + registered_providers.append(ProviderType.ANTIGRAVITY.value) + logger.debug(f"Registered provider: {ProviderType.ANTIGRAVITY.value}") + + # 5. OpenRouter last (catch-all for everything else) if has_openrouter: ModelProviderRegistry.register_provider(ProviderType.OPENROUTER, OpenRouterProvider) registered_providers.append(ProviderType.OPENROUTER.value) diff --git a/tests/test_antigravity_provider.py b/tests/test_antigravity_provider.py new file mode 100644 index 0000000..b129ee6 --- /dev/null +++ b/tests/test_antigravity_provider.py @@ -0,0 +1,311 @@ +"""Unit tests for the Antigravity provider.""" + +import json +import os +import tempfile +from pathlib import Path +from unittest import mock + +import pytest + +from providers.antigravity import AntigravityProvider +from providers.antigravity_auth import AntigravityAccount, AntigravityTokenManager +from providers.registries.antigravity import AntigravityModelRegistry +from providers.shared import ProviderType + + +class TestAntigravityModelRegistry: + """Tests for the Antigravity model registry.""" + + def test_registry_loads_models(self): + """Verify registry loads models from config file.""" + registry = AntigravityModelRegistry() + models = registry.list_models() + + assert len(models) > 0 + assert "claude-opus-4-5-thinking" in models + assert "claude-sonnet-4-5" in models + assert "gemini-3-pro-high" in models + + def test_registry_aliases(self): + """Verify alias resolution works correctly.""" + registry = AntigravityModelRegistry() + + # Test alias resolution + capabilities = registry.resolve("antigravity-claude-opus-4-5-thinking") + assert capabilities is not None + assert capabilities.model_name == "claude-opus-4-5-thinking" + + # Test short alias + capabilities = registry.resolve("ag-opus-thinking") + assert capabilities is not None + assert capabilities.model_name == "claude-opus-4-5-thinking" + + def test_model_capabilities(self): + """Verify model capabilities are correct.""" + registry = AntigravityModelRegistry() + + opus = registry.get_capabilities("claude-opus-4-5-thinking") + assert opus is not None + assert opus.supports_extended_thinking is True + assert opus.max_thinking_tokens == 32768 + assert opus.supports_images is True + + sonnet = registry.get_capabilities("claude-sonnet-4-5") + assert sonnet is not None + assert sonnet.supports_extended_thinking is False + + def test_quota_pool(self): + """Verify quota pool assignment.""" + registry = AntigravityModelRegistry() + + # Antigravity quota models + assert registry.get_quota_pool("claude-opus-4-5-thinking") == "antigravity" + assert registry.get_quota_pool("gemini-3-pro-high") == "antigravity" + + # Gemini CLI quota models + assert registry.get_quota_pool("gemini-2.5-pro") == "gemini-cli" + assert registry.get_quota_pool("gemini-3-flash-preview") == "gemini-cli" + + +class TestAntigravityTokenManager: + """Tests for the OAuth2 token manager.""" + + def test_empty_initialization(self): + """Verify manager handles no accounts gracefully.""" + # Create temp dir without accounts file + with tempfile.TemporaryDirectory() as tmpdir: + fake_path = Path(tmpdir) / "nonexistent.json" + manager = AntigravityTokenManager(accounts_file=fake_path) + assert not manager.has_accounts() + assert manager.get_account_count() == 0 + + def test_env_var_override(self, monkeypatch): + """Verify environment variable takes precedence.""" + monkeypatch.setenv("ANTIGRAVITY_REFRESH_TOKEN", "test_token_123") + monkeypatch.setenv("ANTIGRAVITY_PROJECT_ID", "test-project") + + manager = AntigravityTokenManager() + assert manager.has_accounts() + assert manager.get_account_count() == 1 + + def test_load_accounts_from_file(self): + """Verify loading accounts from JSON file.""" + accounts_data = { + "version": 3, + "accounts": [ + { + "email": "test@example.com", + "refreshToken": "token_1", + "projectId": "project-1", + "enabled": True, + }, + { + "email": "test2@example.com", + "refreshToken": "token_2", + "projectId": "project-2", + "enabled": True, + }, + { + "email": "disabled@example.com", + "refreshToken": "token_3", + "projectId": "project-3", + "enabled": False, # Should be skipped + }, + ], + } + + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: + json.dump(accounts_data, f) + temp_path = f.name + + try: + # Clear env to ensure file is used + with mock.patch.dict(os.environ, {}, clear=True): + manager = AntigravityTokenManager(accounts_file=temp_path) + assert manager.has_accounts() + assert manager.get_account_count() == 2 # Disabled account skipped + finally: + os.unlink(temp_path) + + +class TestAntigravityAccount: + """Tests for the AntigravityAccount dataclass.""" + + def test_token_validity(self): + """Test token validity checking.""" + import time + + account = AntigravityAccount( + email="test@example.com", + refresh_token="token", + project_id="project", + ) + + # No access token + assert not account.is_token_valid() + + # Set valid token + account.access_token = "access_token" + account.token_expiry = time.time() + 3600 + assert account.is_token_valid() + + # Expired token + account.token_expiry = time.time() - 100 + assert not account.is_token_valid() + + def test_rate_limit_tracking(self): + """Test rate limit tracking.""" + import time + + account = AntigravityAccount( + email="test@example.com", + refresh_token="token", + project_id="project", + ) + + # No rate limits initially + assert not account.is_rate_limited() + assert not account.is_rate_limited("claude-opus-4-5-thinking") + + # Set rate limit + account.set_rate_limited("claude-opus-4-5-thinking", 60.0) + assert account.is_rate_limited("claude-opus-4-5-thinking") + assert account.is_rate_limited() # Any rate limit active + + +class TestAntigravityProvider: + """Tests for the Antigravity provider.""" + + def test_provider_type(self): + """Verify provider type is correct.""" + # Mock token manager to avoid needing real credentials + with mock.patch.object(AntigravityTokenManager, "__init__", return_value=None): + with mock.patch.object(AntigravityTokenManager, "has_accounts", return_value=True): + provider = AntigravityProvider() + assert provider.get_provider_type() == ProviderType.ANTIGRAVITY + + def test_model_capabilities(self): + """Verify model capabilities are accessible.""" + with mock.patch.object(AntigravityTokenManager, "__init__", return_value=None): + with mock.patch.object(AntigravityTokenManager, "has_accounts", return_value=True): + provider = AntigravityProvider() + caps = provider.get_all_model_capabilities() + assert len(caps) > 0 + assert "claude-opus-4-5-thinking" in caps + + def test_model_resolution(self): + """Verify alias resolution works through provider.""" + with mock.patch.object(AntigravityTokenManager, "__init__", return_value=None): + with mock.patch.object(AntigravityTokenManager, "has_accounts", return_value=True): + provider = AntigravityProvider() + + # Test alias resolution + resolved = provider._resolve_model_name("ag-opus-thinking") + assert resolved == "claude-opus-4-5-thinking" + + # Test canonical name + resolved = provider._resolve_model_name("claude-sonnet-4-5") + assert resolved == "claude-sonnet-4-5" + + def test_request_building(self): + """Verify request body is built correctly.""" + with mock.patch.object(AntigravityTokenManager, "__init__", return_value=None): + with mock.patch.object(AntigravityTokenManager, "has_accounts", return_value=True): + provider = AntigravityProvider() + + capabilities = provider.get_capabilities("claude-opus-4-5-thinking") + + request = provider._build_request( + prompt="Hello, world!", + model_name="claude-opus-4-5-thinking", + project_id="test-project", + system_prompt="You are a helpful assistant.", + temperature=0.7, + max_output_tokens=1000, + thinking_mode="medium", + capabilities=capabilities, + ) + + assert request["project"] == "test-project" + assert request["model"] == "claude-opus-4-5-thinking" + assert "contents" in request["request"] + assert "systemInstruction" in request["request"] + assert "generationConfig" in request["request"] + + # Check thinking config for thinking models + config = request["request"]["generationConfig"] + assert "thinkingConfig" in config + assert config["thinkingConfig"]["thinkingBudget"] > 0 + + def test_list_models(self): + """Verify list_models returns expected models.""" + with mock.patch.object(AntigravityTokenManager, "__init__", return_value=None): + with mock.patch.object(AntigravityTokenManager, "has_accounts", return_value=True): + provider = AntigravityProvider() + models = provider.list_models(respect_restrictions=False) + + assert len(models) > 0 + # Check that both canonical names and aliases are present + assert any("claude" in m.lower() for m in models) + assert any("gemini" in m.lower() for m in models) + + +@pytest.mark.integration +class TestAntigravityIntegration: + """Integration tests (require actual credentials and valid tokens). + + These tests require: + - Valid Antigravity credentials (env var or accounts file) + - Working OAuth tokens that can be refreshed + + Run with: pytest -m integration tests/test_antigravity_provider.py + Skip with: pytest -m "not integration" tests/test_antigravity_provider.py + """ + + @pytest.mark.skipif( + not os.getenv("ANTIGRAVITY_REFRESH_TOKEN") + and not Path.home().joinpath(".config/opencode/antigravity-accounts.json").exists(), + reason="Antigravity credentials not available", + ) + def test_token_refresh(self): + """Test actual token refresh with real credentials.""" + manager = AntigravityTokenManager() + assert manager.has_accounts() + + # This will trigger a token refresh - may fail if tokens are expired + try: + token, project_id, headers = manager.get_access_token() + assert token is not None + assert len(token) > 0 + assert project_id is not None + except RuntimeError as e: + if "rate limited or have invalid tokens" in str(e): + pytest.skip("Antigravity tokens are expired or invalid") + raise + + @pytest.mark.skipif( + not os.getenv("ANTIGRAVITY_REFRESH_TOKEN") + and not Path.home().joinpath(".config/opencode/antigravity-accounts.json").exists(), + reason="Antigravity credentials not available", + ) + def test_simple_generation(self): + """Test actual content generation with Antigravity API.""" + provider = AntigravityProvider() + + try: + response = provider.generate_content( + prompt="Say 'hello' and nothing else.", + model_name="claude-sonnet-4-5", + temperature=0.0, + max_output_tokens=100, + ) + except RuntimeError as e: + if "rate limited or have invalid tokens" in str(e): + pytest.skip("Antigravity tokens are expired or invalid") + raise + + assert response is not None + assert response.content is not None + assert "hello" in response.content.lower() + assert response.provider == ProviderType.ANTIGRAVITY