feat(providers): add Antigravity provider for unified Claude/Gemini access
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:
2026-02-01 17:55:26 +01:00
parent c71a535f16
commit 5add230d4c
10 changed files with 1612 additions and 1 deletions

View 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