feat(providers): add Antigravity provider for unified Claude/Gemini access
Some checks failed
Semantic Release / release (push) Has been cancelled
Some checks failed
Semantic Release / release (push) Has been cancelled
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
This commit is contained in:
140
conf/antigravity_models.json
Normal file
140
conf/antigravity_models.json
Normal file
@@ -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"
|
||||||
|
}
|
||||||
|
]
|
||||||
|
}
|
||||||
189
docs/antigravity_provider.md
Normal file
189
docs/antigravity_provider.md
Normal file
@@ -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
|
||||||
@@ -1,5 +1,6 @@
|
|||||||
"""Model provider abstractions for supporting multiple AI providers."""
|
"""Model provider abstractions for supporting multiple AI providers."""
|
||||||
|
|
||||||
|
from .antigravity import AntigravityProvider
|
||||||
from .azure_openai import AzureOpenAIProvider
|
from .azure_openai import AzureOpenAIProvider
|
||||||
from .base import ModelProvider
|
from .base import ModelProvider
|
||||||
from .gemini import GeminiModelProvider
|
from .gemini import GeminiModelProvider
|
||||||
@@ -15,6 +16,7 @@ __all__ = [
|
|||||||
"ModelResponse",
|
"ModelResponse",
|
||||||
"ModelCapabilities",
|
"ModelCapabilities",
|
||||||
"ModelProviderRegistry",
|
"ModelProviderRegistry",
|
||||||
|
"AntigravityProvider",
|
||||||
"AzureOpenAIProvider",
|
"AzureOpenAIProvider",
|
||||||
"GeminiModelProvider",
|
"GeminiModelProvider",
|
||||||
"OpenAIModelProvider",
|
"OpenAIModelProvider",
|
||||||
|
|||||||
576
providers/antigravity.py
Normal file
576
providers/antigravity.py
Normal file
@@ -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()
|
||||||
299
providers/antigravity_auth.py
Normal file
299
providers/antigravity_auth.py
Normal file
@@ -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
|
||||||
56
providers/registries/antigravity.py
Normal file
56
providers/registries/antigravity.py
Normal file
@@ -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
|
||||||
@@ -41,6 +41,7 @@ class ModelProviderRegistry:
|
|||||||
ProviderType.AZURE, # Azure-hosted OpenAI deployments
|
ProviderType.AZURE, # Azure-hosted OpenAI deployments
|
||||||
ProviderType.XAI, # Direct X.AI GROK access
|
ProviderType.XAI, # Direct X.AI GROK access
|
||||||
ProviderType.ZEN, # OpenCode Zen curated models
|
ProviderType.ZEN, # OpenCode Zen curated models
|
||||||
|
ProviderType.ANTIGRAVITY, # Antigravity unified gateway (Claude/Gemini via Google OAuth)
|
||||||
ProviderType.DIAL, # DIAL unified API access
|
ProviderType.DIAL, # DIAL unified API access
|
||||||
ProviderType.CUSTOM, # Local/self-hosted models
|
ProviderType.CUSTOM, # Local/self-hosted models
|
||||||
ProviderType.OPENROUTER, # Catch-all for cloud models
|
ProviderType.OPENROUTER, # Catch-all for cloud models
|
||||||
@@ -140,6 +141,19 @@ class ModelProviderRegistry:
|
|||||||
azure_endpoint=azure_endpoint,
|
azure_endpoint=azure_endpoint,
|
||||||
api_version=azure_version,
|
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:
|
else:
|
||||||
if not api_key:
|
if not api_key:
|
||||||
return None
|
return None
|
||||||
@@ -341,6 +355,7 @@ class ModelProviderRegistry:
|
|||||||
ProviderType.OPENROUTER: "OPENROUTER_API_KEY",
|
ProviderType.OPENROUTER: "OPENROUTER_API_KEY",
|
||||||
ProviderType.CUSTOM: "CUSTOM_API_KEY", # Can be empty for providers that don't need auth
|
ProviderType.CUSTOM: "CUSTOM_API_KEY", # Can be empty for providers that don't need auth
|
||||||
ProviderType.DIAL: "DIAL_API_KEY",
|
ProviderType.DIAL: "DIAL_API_KEY",
|
||||||
|
ProviderType.ANTIGRAVITY: "ANTIGRAVITY_REFRESH_TOKEN", # Uses OAuth refresh token, not API key
|
||||||
}
|
}
|
||||||
|
|
||||||
env_var = key_mapping.get(provider_type)
|
env_var = key_mapping.get(provider_type)
|
||||||
|
|||||||
@@ -16,3 +16,4 @@ class ProviderType(Enum):
|
|||||||
CUSTOM = "custom"
|
CUSTOM = "custom"
|
||||||
DIAL = "dial"
|
DIAL = "dial"
|
||||||
ZEN = "zen"
|
ZEN = "zen"
|
||||||
|
ANTIGRAVITY = "antigravity"
|
||||||
|
|||||||
24
server.py
24
server.py
@@ -392,6 +392,7 @@ def configure_providers():
|
|||||||
value = get_env(key)
|
value = get_env(key)
|
||||||
logger.debug(f" {key}: {'[PRESENT]' if value else '[MISSING]'}")
|
logger.debug(f" {key}: {'[PRESENT]' if value else '[MISSING]'}")
|
||||||
from providers import ModelProviderRegistry
|
from providers import ModelProviderRegistry
|
||||||
|
from providers.antigravity import AntigravityProvider
|
||||||
from providers.azure_openai import AzureOpenAIProvider
|
from providers.azure_openai import AzureOpenAIProvider
|
||||||
from providers.custom import CustomProvider
|
from providers.custom import CustomProvider
|
||||||
from providers.dial import DIALModelProvider
|
from providers.dial import DIALModelProvider
|
||||||
@@ -490,6 +491,21 @@ def configure_providers():
|
|||||||
else:
|
else:
|
||||||
logger.debug("OpenCode Zen API key is placeholder value")
|
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.)
|
# Check for custom API endpoint (Ollama, vLLM, etc.)
|
||||||
custom_url = get_env("CUSTOM_API_URL")
|
custom_url = get_env("CUSTOM_API_URL")
|
||||||
if custom_url:
|
if custom_url:
|
||||||
@@ -551,7 +567,13 @@ def configure_providers():
|
|||||||
registered_providers.append(ProviderType.ZEN.value)
|
registered_providers.append(ProviderType.ZEN.value)
|
||||||
logger.debug(f"Registered provider: {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:
|
if has_openrouter:
|
||||||
ModelProviderRegistry.register_provider(ProviderType.OPENROUTER, OpenRouterProvider)
|
ModelProviderRegistry.register_provider(ProviderType.OPENROUTER, OpenRouterProvider)
|
||||||
registered_providers.append(ProviderType.OPENROUTER.value)
|
registered_providers.append(ProviderType.OPENROUTER.value)
|
||||||
|
|||||||
311
tests/test_antigravity_provider.py
Normal file
311
tests/test_antigravity_provider.py
Normal file
@@ -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
|
||||||
Reference in New Issue
Block a user