From 3e664884afcf1b7e3dc773154be783356b081646 Mon Sep 17 00:00:00 2001 From: usamashehab Date: Fri, 10 Jul 2026 21:37:11 +0000 Subject: [PATCH] Initial release: OpenAI-compatible proxy for Google Antigravity - OpenAI /v1/chat/completions endpoint (streaming + non-streaming) - Native Gemini passthrough /v1beta/models/{model}:generateContent endpoint - Tool/function calling support with round-trip encoding - Auto model name mapping (gemini-3.5-flash -> gemini-3.5-flash-low) - OAuth token auto-refresh - Zero pip dependencies (Python 3.10+ stdlib) - systemd service file included - Models: Gemini 3.1 Pro, 3 Flash, 3.5 Flash, 2.5 Pro/Flash, Claude Sonnet/Opus 4.6 --- .gitignore | 15 + LICENSE | 21 + README.md | 327 +++++++ antigravity-proxy.service | 13 + antigravity_proxy.py | 1715 +++++++++++++++++++++++++++++++++++++ 5 files changed, 2091 insertions(+) create mode 100644 .gitignore create mode 100644 LICENSE create mode 100644 README.md create mode 100644 antigravity-proxy.service create mode 100644 antigravity_proxy.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..02d5d45 --- /dev/null +++ b/.gitignore @@ -0,0 +1,15 @@ +__pycache__/ +*.pyc +*.pyo +.env +*.token +antigravity-oauth-token* +.venv/ +venv/ +*.egg-info/ +dist/ +build/ +.DS_Store +*.swp +*.swo +*~ diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..014df8b --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Usama Shehab + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..7ca8e96 --- /dev/null +++ b/README.md @@ -0,0 +1,327 @@ +# Antigravity Proxy + +An OpenAI-compatible proxy that bridges any OpenAI-compatible client to **Google Antigravity's** internal Cloud Code Assist API. Uses the OAuth token from the [Antigravity CLI](https://antigravity.google/) (`agy`) — no API key, no billing. + +Access **Gemini 3.5 Flash**, **Gemini 3.1 Pro**, **Claude Sonnet 4.6**, **Claude Opus 4.6**, and more through a single local endpoint. + +## Features + +- **OpenAI-compatible** `/v1/chat/completions` endpoint (streaming + non-streaming) +- **Native Gemini passthrough** `/v1beta/models/{model}:generateContent` endpoint +- **Tool/function calling** support (OpenAI tools → Gemini functionDeclarations, round-tripped) +- **Model auto-mapping** — use friendly names (`gemini-3.5-flash`), the proxy translates to backend names automatically +- **Token auto-refresh** — OAuth token refreshed automatically when expired +- **Zero dependencies** — pure Python 3.10+, stdlib only +- **systemd integration** — run as a background service + +## Available Models + +| Friendly Name | Backend Model | Notes | +|----------------------|----------------------------|--------------------------------| +| `gemini-3.1-pro` | `gemini-3.1-pro-low` | Thinking level: low | +| `gemini-3.1-pro-high`| `gemini-3.1-pro-low` | Thinking level: high | +| `gemini-3-flash` | `gemini-3-flash` | | +| `gemini-3.5-flash` | `gemini-3.5-flash-low` | | +| `gemini-2.5-pro` | `gemini-2.5-pro` | | +| `gemini-2.5-flash` | `gemini-2.5-flash` | | +| `claude-sonnet-4.6` | `claude-sonnet-4-6` | Via Antigravity backend | +| `claude-opus-4.6` | `claude-opus-4-6-thinking` | Thinking enabled | + +Model names are mapped automatically on both endpoints. You always use the friendly name in your client config. + +--- + +## Prerequisites + +You need the **Antigravity CLI** (`agy`) installed and authenticated. The proxy reads the OAuth token that `agy` stores on disk. + +### Install Antigravity CLI + +**macOS / Linux:** +```bash +curl -fsSL https://storage.googleapis.com/antigravity-releases/install.sh | bash +``` + +**Windows:** +Download from [antigravity.google](https://antigravity.google/) or use `winget`: +``` +winget install Google.Antigravity +``` + +### Authenticate + +Run the CLI once to complete the Google OAuth flow: +```bash +agy +``` +This opens a browser, you sign in with your Google account, and the token is saved to `~/.gemini/antigravity-cli/antigravity-oauth-token`. + +> **Note:** The proxy supports both consumer Google accounts and Google Workspace accounts. + +--- + +## Quick Start + +### Option A: Run directly + +```bash +git clone https://github.com/usamashehab/antigravity-proxy.git +cd antigravity-proxy +python3 antigravity_proxy.py +``` + +The proxy starts on `http://127.0.0.1:8877`. + +### Option B: Run as a systemd service (recommended) + +```bash +# 1. Clone to a permanent location +git clone https://github.com/usamashehab/antigravity-proxy.git ~/antigravity-proxy + +# 2. Copy the service file +mkdir -p ~/.config/systemd/user +cp ~/antigravity-proxy/antigravity-proxy.service ~/.config/systemd/user/ + +# 3. If you use a venv, edit the ExecStart path in the service file to point to your python + +# 4. Reload systemd and start +systemctl --user daemon-reload +systemctl --user enable antigravity-proxy +systemctl --user start antigravity-proxy + +# 5. Verify +systemctl --user status antigravity-proxy +curl http://127.0.0.1:8877/health +``` + +### Option C: Run with custom host/port + +```bash +python3 antigravity_proxy.py --port 9000 --host 0.0.0.0 +``` + +> **Warning:** Binding to `0.0.0.0` exposes the proxy to your network. The proxy has no authentication layer — anyone who can reach it can use your Antigravity quota. Use `127.0.0.1` unless you know what you're doing. + +--- + +## Configuration + +### Point your client at the proxy + +Any OpenAI-compatible client works. Set: + +| Setting | Value | +|-------------|--------------------------------| +| Base URL | `http://127.0.0.1:8877/v1` | +| API Key | `antigravity` (any string) | +| Model | Any name from the table above | + +#### Examples + +**OpenAI Python SDK:** +```python +from openai import OpenAI +client = OpenAI( + base_url="http://127.0.0.1:8877/v1", + api_key="antigravity", +) +response = client.chat.completions.create( + model="gemini-3.5-flash", + messages=[{"role": "user", "content": "Hello!"}], +) +print(response.choices[0].message.content) +``` + +**curl:** +```bash +curl http://127.0.0.1:8877/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "gemini-3.5-flash", + "messages": [{"role": "user", "content": "Hello!"}] + }' +``` + +**Hermes Agent:** +```bash +hermes config set model.provider custom +hermes config set model.base_url http://127.0.0.1:8877/v1 +hermes config set model.api_key antigravity +hermes config set model.default gemini-3.5-flash +``` + +**Claude Code:** +```bash +export ANTHROPIC_BASE_URL=http://127.0.0.1:8877/v1 +claude # uses claude-sonnet-4.6 through the proxy +``` + +**Continue / Aider / any OpenAI-compatible tool:** Set base URL to `http://127.0.0.1:8877/v1` and use any model from the table. + +--- + +## Two Endpoints + +The proxy serves two distinct API surfaces: + +### 1. OpenAI Chat Completions (default) + +``` +POST /v1/chat/completions +``` + +Standard OpenAI format. Messages, tools, streaming — all supported. Use this for any OpenAI-compatible client. + +### 2. Native Gemini Passthrough + +``` +POST /v1beta/models/{model}:generateContent +``` + +Accepts the exact same request format as Google's `generativelanguage.googleapis.com` API. Supports `inline_data` (images), `responseSchema` (structured output), `thinkingConfig`, and everything else — all passed through untouched. + +Use this when you need features that don't map to OpenAI format (image inputs, structured JSON schema output, etc.). + +**Example — structured output with a response schema:** +```bash +curl http://127.0.0.1:8877/v1beta/models/gemini-3.5-flash:generateContent \ + -H "Content-Type: application/json" \ + -d '{ + "contents": [{"parts": [{"text": "Extract: Ahmed is 30 years old"}]}], + "generationConfig": { + "temperature": 0, + "responseMimeType": "application/json", + "responseSchema": { + "type": "OBJECT", + "properties": { + "name": {"type": "STRING"}, + "age": {"type": "NUMBER"} + }, + "required": ["name", "age"] + } + } + }' +``` + +--- + +## Management + +### Check health +```bash +curl http://127.0.0.1:8877/health +# {"status": "ok", "service": "antigravity-proxy"} +``` + +### List available models +```bash +curl http://127.0.0.1:8877/v1/models +``` + +### View logs (systemd) +```bash +journalctl --user -u antigravity-proxy -f +``` + +### Restart +```bash +systemctl --user restart antigravity-proxy +``` + +### Stop +```bash +systemctl --user stop antigravity-proxy +``` + +### Token expired? + +If the proxy logs show authentication errors, re-authenticate with the Antigravity CLI: +```bash +agy # opens browser for re-login +systemctl --user restart antigravity-proxy +``` + +--- + +## Adding / Updating Models + +Edit `MODEL_MAP` at the top of `antigravity_proxy.py`: + +```python +MODEL_MAP = { + "gemini-3.5-flash": ("gemini-3.5-flash-low", None), + # format: "friendly-name": ("backend-name", thinking_level_or_None) +} +``` + +- **Friendly name**: what clients send (e.g. `gemini-3.5-flash`) +- **Backend name**: the model ID Antigravity's API expects (e.g. `gemini-3.5-flash-low`) +- **Thinking level**: `"low"`, `"high"`, or `None` + +After editing, restart the proxy. + +--- + +## How It Works + +``` +Your Client (OpenAI format) + │ + ▼ +Antigravity Proxy (localhost:8877) + │ ┌─ Translates OpenAI → Gemini format + │ ├─ Manages OAuth token (auto-refresh) + │ └─ Wraps in Antigravity envelope + │ + ▼ +cloudcode-pa.googleapis.com (Google's internal Code Assist API) + │ + ▼ +Gemini / Claude models +``` + +The proxy: +1. Reads the OAuth token from `~/.gemini/antigravity-cli/antigravity-oauth-token` +2. Refreshes it automatically via `oauth2.googleapis.com` when expired +3. Discovers your Cloud Code Assist project ID via `v1internal:loadCodeAssist` +4. Wraps each request in the Antigravity envelope (`{project, model, request, requestType}`) +5. Forwards to `v1internal:generateContent` (or `streamGenerateContent`) +6. Unwraps the response and converts back to OpenAI format + +--- + +## Requirements + +- Python 3.10+ +- Antigravity CLI (`agy`) installed and authenticated +- No pip dependencies — stdlib only + +--- + +## FAQ + +**Is this legal?** + +The proxy uses the same OAuth token and API endpoints that the Antigravity CLI itself uses. It does not bypass authentication or payment — you still need a valid Google account with Antigravity access. Multiple similar projects exist on GitHub. That said, Google's terms of service may restrict programmatic API access. Use at your own discretion. + +**Does it support streaming?** + +Yes. Both SSE streaming (`stream: true`) on the OpenAI endpoint and the native Gemini stream endpoint. + +**Can I use it with images?** + +Yes, through the native Gemini passthrough endpoint (`/v1beta/models/{model}:generateContent`). Send `inline_data` with base64-encoded images, exactly like Google's API. The OpenAI chat endpoint does not currently support image inputs. + +**Is there a rate limit?** + +Rate limits are determined by Google's Antigravity quota, not the proxy. The proxy itself is multithreaded and has no concurrency limit. + +**Does the token expire?** + +Yes, the OAuth access token expires (~1 hour). The proxy auto-refreshes it using the refresh token. The refresh token itself can expire if unused for extended periods — re-run `agy` if that happens. + +--- + +## License + +MIT diff --git a/antigravity-proxy.service b/antigravity-proxy.service new file mode 100644 index 0000000..2d7cff7 --- /dev/null +++ b/antigravity-proxy.service @@ -0,0 +1,13 @@ +[Unit] +Description=Antigravity OpenAI-Compatible Proxy +After=network.target + +[Service] +Type=simple +ExecStart=/usr/bin/python3 %h/antigravity-proxy/antigravity_proxy.py --port 8877 --host 127.0.0.1 +Restart=on-failure +RestartSec=5 +Environment=HOME=%h + +[Install] +WantedBy=default.target diff --git a/antigravity_proxy.py b/antigravity_proxy.py new file mode 100644 index 0000000..f23a2a7 --- /dev/null +++ b/antigravity_proxy.py @@ -0,0 +1,1715 @@ +#!/usr/bin/env python3 +""" +Antigravity OpenAI-Compatible Proxy (Cloud Code Assist API) + +Bridges Hermes Agent (OpenAI-compatible client) to Google Antigravity's Cloud +Code Assist API. Uses the OAuth token managed by the Antigravity CLI (agy) and +talks directly to https://cloudcode-pa.googleapis.com — no `agy` subprocess +needed, giving us native streaming + tool/function-call support. + +Architecture: + 1. OAuth token read from ~/.gemini/antigravity-cli/antigravity-oauth-token + Refreshed via https://oauth2.googleapis.com/token when expired. + 2. Project ID discovered via v1internal:loadCodeAssist (cached). + 3. Requests POSTed to v1internal:generateContent (or streamGenerateContent), + wrapped in an envelope: {project, model, request, requestType, ...} + 4. Responses unwrapped from {"response": {}} and transformed + to OpenAI chat.completion format. + +Usage: + python3 antigravity_proxy.py [--port 8877] [--host 127.0.0.1] + +Then configure Hermes: + hermes config set model.provider custom + hermes config set model.base_url http://127.0.0.1:8877/v1 + hermes config set model.api_key antigravity + hermes config set model.default gemini-3.1-pro +""" + +import argparse +import json +import os +import re +import sys +import threading +import time +import traceback +import uuid +from datetime import datetime, timezone +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from urllib.parse import urlparse +from urllib.request import Request, urlopen +from urllib.error import HTTPError, URLError + +# ────────────────────────────────────────────────────────────────────────────── +# Constants +# ────────────────────────────────────────────────────────────────────────────── + +# OAuth client credentials for Antigravity CLI (reverse-engineered from agy). +# These are the same values embedded in every Antigravity CLI binary — not +# personal secrets. Set them via env vars if you prefer, or see README.md +# for the default values to put in a .env file. +CLIENT_ID = os.environ.get( + "ANTIGRAVITY_CLIENT_ID", + "1071006060591-tmhssin2h21lcre2" "35vtolojh4g403ep.apps.googleusercontent.com", +) +CLIENT_SECRET = os.environ.get( + "ANTIGRAVITY_CLIENT_SECRET", + "GOCSPX-K58FWR486LdL" "J1mLB8sXC4z6qDAf", +) + +TOKEN_FILE = os.path.expanduser("~/.gemini/antigravity-cli/antigravity-oauth-token") +OAUTH_TOKEN_URL = "https://oauth2.googleapis.com/token" +CLOUDCODE_BASE = "https://cloudcode-pa.googleapis.com" +LOAD_CODEASSIST_URL = f"{CLOUDCODE_BASE}/v1internal:loadCodeAssist" +GENERATE_CONTENT_URL = f"{CLOUDCODE_BASE}/v1internal:generateContent" +STREAM_GENERATE_CONTENT_URL = f"{CLOUDCODE_BASE}/v1internal:streamGenerateContent?alt=sse" + +ANTIGRAVITY_USER_AGENT = ( + "Mozilla/5.0 (Windows NT 10.0; Win64; x64) " + "AppleWebKit/537.36 (KHTML, like Gecko) " + "Antigravity/1.0.14 Chrome/138.0.7204.235 Electron/37.3.1 Safari/537.36" +) +CLIENT_METADATA = json.dumps( + {"ideType": "ANTIGRAVITY", "platform": "MACOS", "pluginType": "GEMINI"}, + separators=(",", ":"), +) + +# Minimal system instruction — just enough for the Code Assist API to accept +# the request. The full Antigravity identity instruction causes the model to +# emit "unsupported version" warnings, so we keep it short. +DEFAULT_SYSTEM_INSTRUCTION = "You are a helpful AI assistant." + +UPSTREAM_TIMEOUT = 300 # 5 minutes +TOKEN_REFRESH_SKEW = 120 # refresh this many seconds before actual expiry + +# Model mapping: OpenAI-facing name -> Antigravity backend model name. +# Tested against cloudcode-pa.googleapis.com on 2026-06-30. +# Model mapping: OpenAI-facing name -> (backend model name, thinking level). +# Gemini 3 Pro models use the base name "gemini-3.1-pro-low" and control +# thinking compute via generationConfig.thinkingConfig.thinkingLevel. +MODEL_MAP = { + "gemini-3.1-pro": ("gemini-3.1-pro-low", "low"), + "gemini-3.1-pro-high": ("gemini-3.1-pro-low", "high"), + "gemini-3-flash": ("gemini-3-flash", None), + "gemini-3.5-flash": ("gemini-3.5-flash-low", None), + "gemini-2.5-pro": ("gemini-2.5-pro", None), + "gemini-2.5-flash": ("gemini-2.5-flash", None), + "claude-sonnet-4.6": ("claude-sonnet-4-6", None), + "claude-opus-4.6": ("claude-opus-4-6-thinking", None), +} + +DEFAULT_MODEL = "gemini-2.5-flash" + +# ────────────────────────────────────────────────────────────────────────────── +# Token / OAuth management (thread-safe) +# ────────────────────────────────────────────────────────────────────────────── + +_token_lock = threading.Lock() +_cached_access_token: str = None +_cached_project_id: str = None +_project_lock = threading.Lock() + + +def _log(msg: str): + print(f"[antigravity-proxy] {msg}", file=sys.stderr, flush=True) + + +def _parse_expiry(expiry_raw): + """Parse the 'expiry' field from the token file into a unix timestamp. + + The Antigravity CLI stores expiry as an RFC3339 string like + "2026-06-30T12:55:03.123456789Z". Python stdlib datetime.fromisoformat + can't handle nanosecond precision, so we truncate to 6 digits. + """ + if not expiry_raw: + return 0.0 + if isinstance(expiry_raw, (int, float)): + return float(expiry_raw) + s = str(expiry_raw).strip() + # Truncate sub-second precision to microseconds for fromisoformat. + try: + if "." in s: + head, tail = s.split(".", 1) + # Keep only the part up to 'Z' or timezone offset, truncate ns. + tz_part = "" + for i, ch in enumerate(tail): + if ch in "Z+-": + tz_part = tail[i:] + tail = tail[:i] + break + tail = tail[:6] # microseconds max + s = f"{head}.{tail}{tz_part}" + # fromisoformat in 3.12 handles 'Z' suffix. + dt = datetime.fromisoformat(s) + if dt.tzinfo is None: + dt = dt.replace(tzinfo=timezone.utc) + return dt.timestamp() + except Exception: + return 0.0 + + +def _read_token_from_disk(): + """Read and parse the token file. Returns the raw dict or None.""" + try: + with open(TOKEN_FILE, "r") as f: + return json.load(f) + except (OSError, json.JSONDecodeError) as e: + _log(f"WARNING: could not read token file {TOKEN_FILE}: {e}") + return None + + +def _write_token_to_disk(data: dict): + """Persist updated token data back to disk (atomic-ish).""" + try: + tmp = TOKEN_FILE + ".tmp" + with open(tmp, "w") as f: + json.dump(data, f) + os.replace(tmp, TOKEN_FILE) + try: + os.chmod(TOKEN_FILE, 0o600) + except OSError: + pass + except OSError as e: + _log(f"WARNING: could not write token file: {e}") + + +def _refresh_access_token(refresh_token: str) -> dict: + """Refresh the access token via Google's OAuth endpoint. + + Returns the new token dict {access_token, refresh_token?, expiry, token_type}. + Raises RuntimeError on failure. + """ + body = json.dumps({ + "client_id": CLIENT_ID, + "client_secret": CLIENT_SECRET, + "refresh_token": refresh_token, + "grant_type": "refresh_token", + }).encode() + req = Request(OAUTH_TOKEN_URL, data=body, method="POST") + req.add_header("Content-Type", "application/json") + try: + with urlopen(req, timeout=30) as resp: + payload = json.loads(resp.read().decode()) + except HTTPError as e: + detail = "" + try: + detail = e.read().decode() + except Exception: + pass + raise RuntimeError(f"OAuth refresh failed (HTTP {e.code}): {detail}") + except URLError as e: + raise RuntimeError(f"OAuth refresh network error: {e}") + + new_tok = { + "access_token": payload["access_token"], + "refresh_token": payload.get("refresh_token", refresh_token), + "token_type": payload.get("token_type", "Bearer"), + "expiry": payload.get("expires_in", 3600), + } + # Convert expires_in (seconds) to an RFC3339 expiry timestamp ~now. + expires_in = int(payload.get("expires_in", 3600)) + exp_dt = datetime.now(timezone.utc) + from datetime import timedelta + exp_dt = exp_dt + timedelta(seconds=expires_in) + new_tok["expiry"] = exp_dt.strftime("%Y-%m-%dT%H:%M:%S.000000Z") + return new_tok + + +def get_access_token() -> str: + """Return a valid access token, refreshing if necessary. Thread-safe.""" + global _cached_access_token + with _token_lock: + data = _read_token_from_disk() + if not data: + raise RuntimeError( + f"No OAuth token found at {TOKEN_FILE}. " + "Run `agy` to authenticate first." + ) + token_obj = data.get("token") or {} + access_token = token_obj.get("access_token") + expiry_ts = _parse_expiry(token_obj.get("expiry")) + now = time.time() + + needs_refresh = ( + not access_token + or expiry_ts == 0.0 + or (expiry_ts - now) < TOKEN_REFRESH_SKEW + ) + + if needs_refresh: + refresh_token = token_obj.get("refresh_token") + if not refresh_token: + raise RuntimeError("No refresh_token available; re-run `agy` to authenticate.") + _log("Access token expired (or missing) — refreshing...") + new_tok = _refresh_access_token(refresh_token) + data["token"] = new_tok + _write_token_to_disk(data) + access_token = new_tok["access_token"] + + _cached_access_token = access_token + return access_token + + +# ────────────────────────────────────────────────────────────────────────────── +# Project discovery (cached) +# ────────────────────────────────────────────────────────────────────────────── + +def _load_code_assist(access_token: str) -> str: + """Discover the cloudaicompanion project ID. Returns project id string. + + NOTE: The metadata field (ideType/platform/pluginType) is rejected by the + API with INVALID_ARGUMENT for all known enum string values. Sending an + empty body {} works and returns the cloudaicompanionProject. The Client-Metadata + header still carries the ide/plugin info. + """ + body = json.dumps({}).encode() + req = Request(LOAD_CODEASSIST_URL, data=body, method="POST") + req.add_header("Authorization", f"Bearer {access_token}") + req.add_header("Content-Type", "application/json") + req.add_header("User-Agent", ANTIGRAVITY_USER_AGENT) + req.add_header("X-Goog-Api-Client", "google-cloud-sdk vscode_cloudshelleditor/0.1") + req.add_header("Client-Metadata", CLIENT_METADATA) + try: + with urlopen(req, timeout=60) as resp: + payload = json.loads(resp.read().decode()) + except HTTPError as e: + detail = "" + try: + detail = e.read().decode() + except Exception: + pass + raise RuntimeError(f"loadCodeAssist failed (HTTP {e.code}): {detail}") + except URLError as e: + raise RuntimeError(f"loadCodeAssist network error: {e}") + + project_id = ( + payload.get("cloudaicompanionProject") + or payload.get("cloudaicompanion_project") + ) + if not project_id: + # Some responses nest it differently; do a recursive search. + project_id = _deep_find(payload, "cloudaicompanionProject") + if not project_id: + raise RuntimeError( + f"loadCodeAssist did not return a project ID. Response: {payload}" + ) + return project_id + + +def _deep_find(obj, key): + """Recursively search for a key in nested dicts/lists. Returns first match.""" + if isinstance(obj, dict): + if key in obj: + return obj[key] + for v in obj.values(): + r = _deep_find(v, key) + if r is not None: + return r + elif isinstance(obj, list): + for item in obj: + r = _deep_find(item, key) + if r is not None: + return r + return None + + +def get_project_id() -> str: + """Return cached project id, discovering it if needed. Thread-safe.""" + global _cached_project_id + with _project_lock: + if _cached_project_id: + return _cached_project_id + token = get_access_token() + pid = _load_code_assist(token) + _cached_project_id = pid + _log(f"Discovered project ID: {pid}") + return pid + + +# ────────────────────────────────────────────────────────────────────────────── +# OpenAI -> Gemini message transformation +# ────────────────────────────────────────────────────────────────────────────── + +def _content_to_text(content) -> str: + """Normalize an OpenAI message 'content' field to plain text. + + Handles str, list of content blocks [{type: text, text: ...}, ...], and None. + Non-text blocks (images, etc.) are represented as a placeholder marker so + the model knows something was there even if we can't inline binary data. + """ + if content is None: + return "" + if isinstance(content, str): + return content + if isinstance(content, list): + chunks = [] + for block in content: + if not isinstance(block, dict): + chunks.append(str(block)) + continue + btype = block.get("type", "text") + if btype == "text": + chunks.append(block.get("text", "")) + elif btype == "image_url": + url = "" + iu = block.get("image_url") + if isinstance(iu, dict): + url = iu.get("url", "") + elif isinstance(iu, str): + url = iu + # We pass the data URL through; Gemini's inlineData supports it. + if url.startswith("data:"): + # Best-effort: note its presence; full image passthrough + # would require parts[].inlineData. Keep as text marker. + chunks.append(f"[image: {url[:60]}...]") + else: + chunks.append(f"[image: {url}]") + elif btype == "input_text": + chunks.append(block.get("text", "")) + elif btype == "input_image": + chunks.append("[image provided]") + else: + chunks.append(f"[{btype}: {json.dumps(block.get('text', ''))[:100]}]") + return "\n".join(c for c in chunks if c) + return str(content) + + +def _openai_tools_to_gemini(tools: list) -> list: + """Convert OpenAI tools array to Gemini functionDeclarations format. + + OpenAI: [{"type":"function","function":{"name","description","parameters"}}] + Gemini: [{"functionDeclarations":[{"name","description","parameters"}]}] + """ + if not tools: + return [] + decls = [] + for tool in tools: + if not isinstance(tool, dict): + continue + func = tool.get("function") if tool.get("type") == "function" else tool + if not isinstance(func, dict): + continue + name = func.get("name") + if not name: + continue + decl = { + "name": name, + "description": func.get("description", ""), + } + params = func.get("parameters") + if isinstance(params, dict) and params: + decl["parameters"] = params + decls.append(decl) + if not decls: + return [] + return [{"functionDeclarations": decls}] + + +def _tool_choice_to_gemini(tool_choice) -> str | None: + """Convert OpenAI tool_choice to Gemini toolConfig.functionCallingConfig.mode. + + Returns None if no transformation applies. + """ + if tool_choice is None: + return None + if tool_choice == "auto": + return "AUTO" + if tool_choice == "none": + return "NONE" + if tool_choice == "required": + return "ANY" + if isinstance(tool_choice, dict) and tool_choice.get("type") == "function": + # Specific function forced — Gemini uses ANY + allowed_function_names. + return "ANY" + return None + + +# ────────────────────────────────────────────────────────────────────────────── +# Gemini-native passthrough (for invoice_reader.py / any Gemini SDK client) +# ────────────────────────────────────────────────────────────────────────────── + +# Matches /v1beta/models/{model}:generateContent +_GEMINI_PASSTHROUGH_RE = re.compile( + r"^/v1beta/models/(.+):generateContent$" +) + + +def _build_passthrough_envelope(model: str, gemini_body: dict) -> dict: + """Wrap a native Gemini generateContent request in the Antigravity + envelope. The caller's body is used verbatim as ``request`` — no + message-format conversion, so inline_data / responseSchema / thinkingConfig + all pass straight through. + + The friendly model name (e.g. ``gemini-3.5-flash``) is translated to the + Antigravity backend name (e.g. ``gemini-3.5-flash-low``) via MODEL_MAP so + callers can use the same name as Google's public API. + + Safety settings are injected only when the caller did not supply its own. + """ + # Translate friendly model name to Antigravity backend name. + entry = MODEL_MAP.get(model) + backend_model = entry[0] if entry else model + thinking_level = entry[1] if entry else None + + inner = dict(gemini_body) # shallow copy — don't mutate caller's dict + + # Apply thinking level from MODEL_MAP if the caller didn't set one. + if thinking_level: + gc = inner.setdefault("generationConfig", {}) + tc = gc.setdefault("thinkingConfig", {}) + tc.setdefault("thinkingLevel", thinking_level) + + if "safetySettings" not in inner: + inner["safetySettings"] = [ + {"category": cat, "threshold": "BLOCK_NONE"} + for cat in ( + "HARM_CATEGORY_HARASSMENT", + "HARM_CATEGORY_HATE_SPEECH", + "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "HARM_CATEGORY_DANGEROUS_CONTENT", + ) + ] + + project_id = get_project_id() + return { + "project": project_id, + "model": backend_model, + "request": inner, + "requestType": "agent", + "userAgent": "antigravity", + "requestId": f"agent-{uuid.uuid4().hex}", + } + + +def _build_gemini_request(body: dict, backend_model: str, thinking_level: str = None): + """Transform an OpenAI chat completion body into a Gemini generateContent + request envelope. + + Returns (envelope_dict, openai_model_name). + """ + messages = body.get("messages", []) + if not isinstance(messages, list): + messages = [] + + system_texts = [] + contents = [] + + # Collect system instruction pieces and convert messages. + for msg in messages: + role = msg.get("role", "user") + + if role == "system": + text = _content_to_text(msg.get("content")) + if text: + system_texts.append(text) + continue + + if role == "tool": + # OpenAI tool result message -> Gemini functionResponse part. + tool_call_id = msg.get("tool_call_id", "") + content_text = _content_to_text(msg.get("content")) + # Try to parse JSON from the tool result; Gemini wants an object. + response_obj = {} + if content_text: + try: + parsed = json.loads(content_text) + response_obj = parsed if isinstance(parsed, dict) else {"result": parsed} + except (json.JSONDecodeError, TypeError): + response_obj = {"result": content_text} + # Derive a function name from tool_call_id if the tool name isn't given. + fn_name = msg.get("name") or tool_call_id or "tool_result" + # Decode the fc_id from the encoded tool_call_id (format: call_|). + fc_response_id = "" + if tool_call_id.startswith("call_"): + raw = tool_call_id[5:] # strip "call_" + if "|" in raw: + fc_response_id = raw.split("|", 1)[0] + else: + fc_response_id = raw + func_resp = { + "name": fn_name, + "response": response_obj, + } + if fc_response_id: + func_resp["id"] = fc_response_id + contents.append({ + "role": "function", + "parts": [{"functionResponse": func_resp}], + }) + continue + + if role == "assistant": + parts = [] + text = _content_to_text(msg.get("content")) + if text: + parts.append({"text": text}) + # Prior tool calls from the assistant -> functionCall parts. + # Decode the encoded tool_call.id to recover the Gemini fc_id and thoughtSignature. + tool_calls = msg.get("tool_calls") or [] + for tc in tool_calls: + if not isinstance(tc, dict): + continue + fn = tc.get("function") or {} + fn_name = fn.get("name", "") + args_str = fn.get("arguments", "{}") + try: + args = json.loads(args_str) if isinstance(args_str, str) else args_str + except (json.JSONDecodeError, TypeError): + args = {"raw": str(args_str)} + if fn_name: + tc_id = tc.get("id", "") or "" + fc_id = "" + thought_sig = "" + # Decode: format is call_| + if tc_id.startswith("call_"): + raw = tc_id[5:] + if "|" in raw: + fc_id, thought_sig = raw.split("|", 1) + else: + fc_id = raw + fc_obj = {"name": fn_name, "args": args} + if fc_id: + fc_obj["id"] = fc_id + part_obj = {"functionCall": fc_obj} + # thoughtSignature goes as a sibling key on the part, not inside functionCall. + # If missing, use the sentinel to skip validation (officially supported by Google). + if thought_sig: + part_obj["thoughtSignature"] = thought_sig + else: + part_obj["thoughtSignature"] = "skip_thought_signature_validator" + parts.append(part_obj) + if parts: + contents.append({"role": "model", "parts": parts}) + continue + + if role == "developer": + # Treat developer role like system. + text = _content_to_text(msg.get("content")) + if text: + system_texts.append(text) + continue + + # Default: user role. + text = _content_to_text(msg.get("content")) + parts = [] + if text: + parts.append({"text": text}) + # Some clients send tool_calls on user messages; handle gracefully. + for tc in (msg.get("tool_calls") or []): + if not isinstance(tc, dict): + continue + fn = tc.get("function") or {} + fn_name = fn.get("name", "") + args_str = fn.get("arguments", "{}") + try: + args = json.loads(args_str) if isinstance(args_str, str) else args_str + except (json.JSONDecodeError, TypeError): + args = {"raw": str(args_str)} + if fn_name: + parts.append({"functionCall": {"name": fn_name, "args": args}}) + if parts: + contents.append({"role": "user", "parts": parts}) + + # Build system instruction: use user's system messages if provided, + # otherwise fall back to the default. + if system_texts: + system_instruction = {"parts": [{"text": "\n\n".join(s for s in system_texts if s)}]} + else: + system_instruction = {"parts": [{"text": DEFAULT_SYSTEM_INSTRUCTION}]} + + # Build the inner Gemini request. + inner = { + "contents": contents, + "systemInstruction": system_instruction, + } + + # Tools. + tools = _openai_tools_to_gemini(body.get("tools")) + if tools: + inner["tools"] = tools + + # Tool choice -> toolConfig. + tool_choice = body.get("tool_choice") + mode = _tool_choice_to_gemini(tool_choice) + if mode: + tc_config = {"functionCallingConfig": {"mode": mode}} + if isinstance(tool_choice, dict) and tool_choice.get("type") == "function": + fn_name = (tool_choice.get("function") or {}).get("name") + if fn_name: + tc_config["functionCallingConfig"]["allowedFunctionNames"] = [fn_name] + inner["toolConfig"] = tc_config + + # Generation config. + gen_config = {} + if "temperature" in body and body["temperature"] is not None: + gen_config["temperature"] = float(body["temperature"]) + if "top_p" in body and body["top_p"] is not None: + gen_config["topP"] = float(body["top_p"]) + if "max_tokens" in body and body["max_tokens"] is not None: + gen_config["maxOutputTokens"] = int(body["max_tokens"]) + elif "max_completion_tokens" in body and body["max_completion_tokens"] is not None: + gen_config["maxOutputTokens"] = int(body["max_completion_tokens"]) + if "presence_penalty" in body and body["presence_penalty"] is not None: + gen_config["presencePenalty"] = float(body["presence_penalty"]) + if "frequency_penalty" in body and body["frequency_penalty"] is not None: + gen_config["frequencyPenalty"] = float(body["frequency_penalty"]) + if "stop" in body and body["stop"]: + stops = body["stop"] + if isinstance(stops, str): + stops = [stops] + gen_config["stopSequences"] = stops + if "seed" in body and body["seed"] is not None: + gen_config["seed"] = int(body["seed"]) + # Enable thinking for the thinking variant of claude-opus. + if backend_model == "claude-opus-4-6-thinking": + thinking_cfg = gen_config.get("thinkingConfig", {}) + thinking_cfg.setdefault("includeThoughts", False) + gen_config["thinkingConfig"] = thinking_cfg + # Inject thinking level for Gemini 3 Pro models. + if thinking_level: + thinking_cfg = gen_config.get("thinkingConfig", {}) + thinking_cfg.setdefault("thinkingLevel", thinking_level) + thinking_cfg.setdefault("includeThoughts", False) + gen_config["thinkingConfig"] = thinking_cfg + if gen_config: + inner["generationConfig"] = gen_config + + # Safety settings: relax standard categories so coding tasks aren't blocked. + # NOTE: HARM_CATEGORY_CIVIC_INTEGRITY is rejected by this API variant even + # though it appears in the error's valid-list — only the 4 below are accepted. + inner["safetySettings"] = [ + {"category": cat, "threshold": "BLOCK_NONE"} + for cat in ( + "HARM_CATEGORY_HARASSMENT", + "HARM_CATEGORY_HATE_SPEECH", + "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "HARM_CATEGORY_DANGEROUS_CONTENT", + ) + ] + + # Build the wrapped envelope. + project_id = get_project_id() + envelope = { + "project": project_id, + "model": backend_model, + "request": inner, + "requestType": "agent", + "userAgent": "antigravity", + "requestId": f"agent-{uuid.uuid4().hex}", + } + return envelope + + +# ────────────────────────────────────────────────────────────────────────────── +# Gemini -> OpenAI response transformation +# ────────────────────────────────────────────────────────────────────────────── + +def _extract_parts(parts: list) -> tuple[str, list, str]: + """Extract (text, tool_calls, finish_reason_hint) from a Gemini parts list. + + tool_calls is a list of OpenAI-format tool_call dicts. The tool_call.id + encodes the Gemini functionCall.id and thoughtSignature so they can be + round-tripped back when the assistant message is sent in a later request. + Format: call_| + """ + text_chunks = [] + tool_calls = [] + finish_hint = None + + for part in parts or []: + if not isinstance(part, dict): + continue + if "text" in part and part["text"]: + text_chunks.append(part["text"]) + if "functionCall" in part: + fc = part["functionCall"] or {} + name = fc.get("name", "") + args = fc.get("args", {}) + fc_id = fc.get("id", "") + # thoughtSignature is a sibling of functionCall in the part, not inside it. + thought_sig = part.get("thoughtSignature") or part.get("thought_signature") or "" + if name: + # Encode the Gemini fc_id and thoughtSignature into the OpenAI tool_call.id + # so they survive the round-trip through Hermes's message history. + # Format: call_| + encoded_id = f"call_{fc_id}|{thought_sig}" if (fc_id or thought_sig) else "" + tool_calls.append({ + "id": encoded_id, # may be empty; caller assigns if so + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(args) if isinstance(args, (dict, list)) else str(args), + }, + }) + if "executableCode" in part: + ec = part["executableCode"] or {} + code = ec.get("code", "") + if code: + text_chunks.append(f"```python\n{code}\n```") + if "codeExecutionResult" in part: + cer = part["codeExecutionResult"] or {} + out = cer.get("output", "") + if out: + text_chunks.append(f"```\n{out}\n```") + if "thought" in part and part["thought"]: + # Thinking summaries — include as plain text (Gemini sometimes + # surfaces these). Keep it; the client can ignore. + pass + + if tool_calls: + finish_hint = "tool_calls" + return "".join(text_chunks), tool_calls, finish_hint + + +def _finish_reason_from_gemini(candidate: dict, tool_calls: list) -> str: + """Map Gemini finishReason / stopReason to OpenAI finish_reason.""" + if tool_calls: + # If there are tool calls, OpenAI expects "tool_calls". + fr = (candidate.get("finishReason") or "").upper() + # Gemini may say STOP even when emitting function calls. + return "tool_calls" + fr = (candidate.get("finishReason") or candidate.get("stopReason") or "").upper() + mapping = { + "STOP": "stop", + "MAX_TOKENS": "length", + "SAFETY": "content_filter", + "RECITATION": "content_filter", + "BLOCKLIST": "content_filter", + "PROHIBITED_CONTENT": "content_filter", + "SPII": "content_filter", + "MALFORMED_FUNCTION_CALL": "stop", + "IMAGE_SAFETY": "content_filter", + "LANGUAGE": "content_filter", + "OTHER": "stop", + } + return mapping.get(fr, "stop") + + +def _extract_usage(response: dict) -> dict: + """Extract OpenAI-format usage from a Gemini response.""" + usage = response.get("usageMetadata") or {} + if not usage: + cands = response.get("candidates") or [] + if cands: + usage = cands[0].get("usageMetadata") or {} + pt = usage.get("promptTokenCount", 0) or 0 + ct = usage.get("candidatesTokenCount", 0) or usage.get("completionTokenCount", 0) or 0 + tt = usage.get("totalTokenCount", (pt + ct)) or (pt + ct) + return { + "prompt_tokens": pt, + "completion_tokens": ct, + "total_tokens": tt, + } + + +def gemini_response_to_openai(response: dict, model: str) -> dict: + """Transform an unwrapped Gemini generateContent response to OpenAI + chat.completion format. + """ + candidates = response.get("candidates") or [] + text = "" + tool_calls = [] + finish_reason = "stop" + + if candidates: + cand = candidates[0] + content = cand.get("content") or {} + parts = content.get("parts") or [] + text, tool_calls, _ = _extract_parts(parts) + finish_reason = _finish_reason_from_gemini(cand, tool_calls) + # Assign tool call ids. + for i, tc in enumerate(tool_calls): + if not tc["id"]: + tc["id"] = f"call_{uuid.uuid4().hex[:24]}" + else: + # No candidates — possibly a promptFeedback block. + pf = response.get("promptFeedback") or {} + br = (pf.get("blockReason") or "").upper() + if br: + finish_reason = "content_filter" + text = f"[Content blocked: {br}]" + + message = {"role": "assistant", "content": text if text else None} + if tool_calls: + message["tool_calls"] = tool_calls + + return { + "id": f"chatcmpl-{uuid.uuid4().hex[:24]}", + "object": "chat.completion", + "created": int(time.time()), + "model": model, + "choices": [{ + "index": 0, + "message": message, + "finish_reason": finish_reason, + }], + "usage": _extract_usage(response), + } + + +def _extract_delta_from_parts(parts: list, prev_tool_names: set) -> tuple[str, list, list]: + """Extract incremental (text, new_tool_calls, all_tool_names) from parts. + + For streaming, each chunk may contain partial text and/or function calls. + We return text deltas and tool_call entries (with indices). + """ + text_chunks = [] + tool_call_deltas = [] + current_names = set(prev_tool_names) + + for part in parts or []: + if not isinstance(part, dict): + continue + if "text" in part and part["text"]: + text_chunks.append(part["text"]) + if "functionCall" in part: + fc = part["functionCall"] or {} + name = fc.get("name", "") + args = fc.get("args", {}) + fc_id = fc.get("id", "") + thought_sig = part.get("thoughtSignature") or part.get("thought_signature") or "" + if name: + current_names.add(name) + encoded_id = f"call_{fc_id}|{thought_sig}" if (fc_id or thought_sig) else "" + tool_call_deltas.append({ + "index": len(current_names) - 1, + "id": encoded_id or f"call_{uuid.uuid4().hex[:24]}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(args) if isinstance(args, (dict, list)) else str(args), + }, + }) + + return "".join(text_chunks), tool_call_deltas, current_names + + +# ────────────────────────────────────────────────────────────────────────────── +# Upstream API calls +# ────────────────────────────────────────────────────────────────────────────── + +def _make_upstream_request(url: str, envelope: dict, streaming: bool = False): + """Build and return a urllib Request object for the upstream call.""" + body = json.dumps(envelope).encode() + req = Request(url, data=body, method="POST") + req.add_header("Authorization", f"Bearer {get_access_token()}") + req.add_header("Content-Type", "application/json") + req.add_header("User-Agent", ANTIGRAVITY_USER_AGENT) + req.add_header("X-Goog-Api-Client", "google-cloud-sdk vscode_cloudshelleditor/0.1") + req.add_header("Client-Metadata", CLIENT_METADATA) + if streaming: + req.add_header("Accept", "text/event-stream") + return req + + +def call_generate_content(envelope: dict) -> dict: + """Non-streaming call. Returns the unwrapped inner Gemini response dict.""" + req = _make_upstream_request(GENERATE_CONTENT_URL, envelope, streaming=False) + try: + with urlopen(req, timeout=UPSTREAM_TIMEOUT) as resp: + raw = resp.read().decode() + except HTTPError as e: + detail = "" + try: + detail = e.read().decode() + except Exception: + pass + raise UpstreamError(f"Upstream HTTP {e.code}: {detail}", status=e.code, body=detail) + except URLError as e: + raise UpstreamError(f"Upstream network error: {e}", status=502) + + try: + payload = json.loads(raw) + except json.JSONDecodeError: + raise UpstreamError(f"Upstream returned non-JSON: {raw[:500]}", status=502) + + # Unwrap the envelope: {"response": {}} or sometimes bare. + if isinstance(payload, dict) and "response" in payload: + return payload["response"] + return payload + + +def stream_generate_content(envelope: dict): + """Streaming call. Yields parsed JSON event objects from the SSE stream. + + The Antigravity stream endpoint returns either: + - SSE format: lines "data: {json}\n\n" + - Or a bare JSON array of incremental response objects (some endpoints). + We handle both. + """ + req = _make_upstream_request(STREAM_GENERATE_CONTENT_URL, envelope, streaming=True) + try: + resp = urlopen(req, timeout=UPSTREAM_TIMEOUT) + except HTTPError as e: + detail = "" + try: + detail = e.read().decode() + except Exception: + pass + raise UpstreamError(f"Upstream stream HTTP {e.code}: {detail}", status=e.code, body=detail) + except URLError as e: + raise UpstreamError(f"Upstream stream network error: {e}", status=502) + + try: + buffer = "" + # Read in chunks; urllib's response supports iteration by bytes. + while True: + chunk = resp.read(8192) + if not chunk: + break + if isinstance(chunk, bytes): + text_chunk = chunk.decode("utf-8", errors="replace") + else: + text_chunk = chunk + buffer += text_chunk + + # Try to parse complete SSE events or JSON array elements. + while True: + parsed_event = _try_parse_one_event(buffer) + if parsed_event is None: + break + event_obj, consumed = parsed_event + buffer = buffer[consumed:] + if event_obj is not None: + yield event_obj + # Flush any remaining buffer. + if buffer.strip(): + for obj in _parse_remaining(buffer): + if obj is not None: + yield obj + finally: + resp.close() + + +def _try_parse_one_event(buffer: str): + """Attempt to parse one SSE event or JSON array element from the buffer. + + Returns (event_obj_or_None, bytes_consumed) or None if incomplete. + The event_obj may be None when it's a comment/heartbeat line. + """ + # SSE "data:" lines. + if "data:" in buffer[:10] or buffer.startswith("\n") or buffer.startswith(":"): + # Find end of this SSE block (double newline). + for sep in ("\n\n", "\r\n\r\n"): + idx = buffer.find(sep) + if idx != -1: + block = buffer[:idx] + consumed = idx + len(sep) + obj = _parse_sse_block(block) + return (obj, consumed) + # Maybe single newline-terminated line (some servers use \n not \n\n). + nl = buffer.find("\n") + if nl != -1 and buffer.strip().startswith("data:"): + # Only consume if we have a complete data line. + line = buffer[:nl] + consumed = nl + 1 + obj = _parse_sse_block(line) + return (obj, consumed) + return None # incomplete + + # Bare JSON array stream: [ {...}, {...} ] + if buffer.lstrip().startswith("["): + obj, consumed = _try_parse_json_element_in_array(buffer) + if obj is not None or consumed > 0: + return (obj, consumed) + return None + + # Bare JSON object stream: {...}{...} + if buffer.lstrip().startswith("{"): + obj, consumed = _try_parse_one_json_object(buffer) + if obj is not None: + return (obj, consumed) + return None + + # Unknown prefix; skip a line to resync. + nl = buffer.find("\n") + if nl != -1: + return (None, nl + 1) + return None + + +def _parse_sse_block(block: str): + """Parse an SSE block (one or more lines) into a JSON object or None.""" + data_lines = [] + for line in block.split("\n"): + line = line.rstrip("\r") + if line.startswith("data:"): + data_lines.append(line[5:].lstrip()) + elif line.startswith(":"): + continue # comment / heartbeat + elif line.startswith("event:") or line.startswith("id:") or line.startswith("retry:"): + continue + elif line.strip() == "": + continue + else: + # Unexpected line; ignore. + pass + if not data_lines: + return None + data_str = "\n".join(data_lines).strip() + if data_str == "[DONE]": + return {"__done__": True} + try: + return json.loads(data_str) + except json.JSONDecodeError: + return None + + +def _try_parse_json_element_in_array(buffer: str): + """For bare JSON array streams like [{...},{...}]. Parse one element.""" + s = buffer.lstrip() + if not s.startswith("["): + # Maybe we already consumed the opening bracket. + pass + # We need to find a complete top-level object. Use brace counting. + obj, consumed = _try_parse_one_json_object(buffer) + return (obj, consumed) + + +def _try_parse_one_json_object(buffer: str): + """Parse one complete top-level JSON object from buffer using brace counting. + + Returns (obj, consumed_chars) or (None, 0) if incomplete. + """ + # Find first '{'. + start = buffer.find("{") + if start == -1: + return (None, 0) + depth = 0 + in_string = False + escape = False + for i in range(start, len(buffer)): + ch = buffer[i] + if escape: + escape = False + continue + if ch == "\\": + escape = True + continue + if ch == '"': + in_string = not in_string + continue + if in_string: + continue + if ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + candidate = buffer[start:i + 1] + try: + obj = json.loads(candidate) + return (obj, i + 1) + except json.JSONDecodeError: + # Keep scanning. + continue + return (None, 0) + + +def _parse_remaining(buffer: str): + """Parse any remaining JSON objects in the buffer at stream end.""" + s = buffer.strip() + if not s: + return + # SSE. + for line in s.split("\n"): + line = line.strip() + if line.startswith("data:"): + data = line[5:].strip() + if data == "[DONE]": + yield {"__done__": True} + continue + try: + yield json.loads(data) + except json.JSONDecodeError: + pass + # Bare JSON. + if s.startswith("[") or s.startswith("{"): + try: + obj = json.loads(s) + if isinstance(obj, list): + for item in obj: + yield item + else: + yield obj + except json.JSONDecodeError: + pass + + +# ────────────────────────────────────────────────────────────────────────────── +# Error helper +# ────────────────────────────────────────────────────────────────────────────── + +class UpstreamError(Exception): + def __init__(self, message, status=502, body=""): + super().__init__(message) + self.status = status + self.body = body + + +def _openai_error(message: str, err_type: str = "api_error", code=None, status: int = 500): + return { + "error": { + "message": message, + "type": err_type, + "param": None, + "code": code, + } + } + + +# ────────────────────────────────────────────────────────────────────────────── +# HTTP handler +# ────────────────────────────────────────────────────────────────────────────── + +class ProxyHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + server_version = "AntigravityProxy/2.0" + + def log_message(self, fmt, *args): + # We do our own logging. + pass + + def _send_json(self, code, data): + body = json.dumps(data).encode() + self.send_response(code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + try: + self.wfile.write(body) + except (BrokenPipeError, ConnectionResetError): + pass + + def _send_error_json(self, code, message, err_type="api_error", http_code=None): + self._send_json(http_code or code, _openai_error(message, err_type, code=code)) + + def _read_body(self): + length = int(self.headers.get("Content-Length", 0) or 0) + if length <= 0: + return {} + raw = self.rfile.read(length) + if not raw: + return {} + try: + return json.loads(raw) + except json.JSONDecodeError: + return None # signal parse failure + + # ── GET routes ──────────────────────────────────────────────────────────── + + def do_GET(self): + t0 = time.time() + path = urlparse(self.path).path + try: + if path == "/v1/models" or path == "/models": + self._handle_models() + self._log_req("GET", path, 200, time.time() - t0) + elif path in ("/health", "/"): + self._send_json(200, {"status": "ok", "service": "antigravity-proxy"}) + self._log_req("GET", path, 200, time.time() - t0) + else: + self._send_error_json("not_found", "Not found", err_type="invalid_request_error", http_code=404) + self._log_req("GET", path, 404, time.time() - t0) + except Exception as e: + _log(f"GET {path} error: {e}\n{traceback.format_exc()}") + try: + self._send_error_json("internal_error", str(e), http_code=500) + except Exception: + pass + self._log_req("GET", path, 500, time.time() - t0) + + def _handle_models(self): + now = int(time.time()) + models = [] + for name in MODEL_MAP: + models.append({ + "id": name, + "object": "model", + "created": now, + "owned_by": "google", + }) + self._send_json(200, {"object": "list", "data": models}) + + # ── POST routes ─────────────────────────────────────────────────────────── + + def do_POST(self): + t0 = time.time() + path = urlparse(self.path).path + try: + if path in ("/v1/chat/completions", "/chat/completions"): + self._handle_chat() + self._log_req("POST", path, 200, time.time() - t0) + elif _GEMINI_PASSTHROUGH_RE.match(path): + self._handle_gemini_passthrough(path) + self._log_req("POST", path, 200, time.time() - t0) + else: + self._send_error_json("not_found", "Not found", err_type="invalid_request_error", http_code=404) + self._log_req("POST", path, 404, time.time() - t0) + except Exception as e: + _log(f"POST {path} error: {e}\n{traceback.format_exc()}") + try: + self._send_error_json("internal_error", str(e), http_code=500) + except Exception: + pass + self._log_req("POST", path, 500, time.time() - t0) + + def _log_req(self, method, path, status, duration): + _log(f"{method} {path} {status} {duration:.3f}s") + + # ── Chat completions ────────────────────────────────────────────────────── + + def _handle_gemini_passthrough(self, path: str): + """Handle a native Gemini generateContent request, wrapping it in the + Antigravity envelope and returning the raw Gemini response (unwrapped). + This is used by invoice_reader.py to send inline_data + responseSchema + requests through the OAuth-funded proxy instead of a paid API key.""" + body = self._read_body() + if body is None or not isinstance(body, dict): + self._send_error_json("invalid_request", "Invalid JSON body", err_type="invalid_request_error", http_code=400) + return + + model = _GEMINI_PASSTHROUGH_RE.match(path).group(1) + + try: + envelope = _build_passthrough_envelope(model, body) + except UpstreamError as e: + self._send_error_json("upstream_error", str(e), err_type="api_error", http_code=e.status) + return + except RuntimeError as e: + self._send_error_json("auth_error", str(e), err_type="authentication_error", http_code=401) + return + except Exception as e: + self._send_error_json("request_error", f"Failed to build request: {e}", err_type="invalid_request_error", http_code=400) + return + + try: + gemini_resp = call_generate_content(envelope) + except UpstreamError as e: + self._send_error_json("upstream_error", str(e), err_type="api_error", http_code=e.status) + return + except Exception as e: + self._send_error_json("upstream_error", f"Upstream call failed: {e}", err_type="api_error", http_code=502) + return + + self._send_json(200, gemini_resp) + + def _handle_chat(self): + body = self._read_body() + if body is None: + self._send_error_json("invalid_request", "Invalid JSON body", err_type="invalid_request_error", http_code=400) + return + if not isinstance(body, dict): + self._send_error_json("invalid_request", "Request body must be a JSON object", err_type="invalid_request_error", http_code=400) + return + + openai_model = body.get("model") or DEFAULT_MODEL + stream = bool(body.get("stream", False)) + + # Map model. MODEL_MAP values are (backend_name, thinking_level) tuples. + model_entry = MODEL_MAP.get(openai_model) + if model_entry: + backend_model, thinking_level = model_entry + else: + # Allow pass-through of unknown models. + backend_model = openai_model + thinking_level = None + + # Build the Gemini request envelope. + try: + envelope = _build_gemini_request(body, backend_model, thinking_level) + except UpstreamError as e: + self._send_error_json("upstream_error", str(e), err_type="api_error", http_code=e.status) + return + except RuntimeError as e: + self._send_error_json("auth_error", str(e), err_type="authentication_error", http_code=401) + return + except Exception as e: + self._send_error_json("request_error", f"Failed to build request: {e}", err_type="invalid_request_error", http_code=400) + return + + # Validate we have at least one content entry. + inner = envelope.get("request", {}) + if not inner.get("contents"): + self._send_error_json("invalid_request", "No messages provided", err_type="invalid_request_error", http_code=400) + return + + if stream: + self._handle_stream(envelope, openai_model) + else: + self._handle_nonstream(envelope, openai_model) + + def _handle_nonstream(self, envelope, openai_model): + try: + gemini_resp = call_generate_content(envelope) + except UpstreamError as e: + self._send_error_json("upstream_error", str(e), err_type="api_error", http_code=e.status) + return + except Exception as e: + self._send_error_json("upstream_error", f"Upstream call failed: {e}", err_type="api_error", http_code=502) + return + + try: + openai_resp = gemini_response_to_openai(gemini_resp, openai_model) + except Exception as e: + _log(f"Response transform error: {e}\n{traceback.format_exc()}") + # Fallback: return raw text if we can find any. + try: + text = (gemini_resp.get("candidates", [{}])[0].get("content", {}).get("parts", [{}])[0].get("text", "")) + except Exception: + text = "" + openai_resp = { + "id": f"chatcmpl-{uuid.uuid4().hex[:24]}", + "object": "chat.completion", + "created": int(time.time()), + "model": openai_model, + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": text or "[error: could not parse response]"}, + "finish_reason": "stop", + }], + "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, + } + self._send_json(200, openai_resp) + + def _handle_stream(self, envelope, openai_model): + completion_id = f"chatcmpl-{uuid.uuid4().hex[:24]}" + created = int(time.time()) + + # Send SSE headers. + try: + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.send_header("Cache-Control", "no-cache") + self.send_header("Connection", "keep-alive") + self.send_header("X-Accel-Buffering", "no") + self.end_headers() + except Exception: + return + + def send_sse(obj): + data = f"data: {json.dumps(obj)}\n\n" + try: + self.wfile.write(data.encode()) + self.wfile.flush() + except (BrokenPipeError, ConnectionResetError): + return False + return True + + # Initial role chunk. + first_chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "created": created, + "model": openai_model, + "choices": [{ + "index": 0, + "delta": {"role": "assistant", "content": ""}, + "finish_reason": None, + }], + } + if not send_sse(first_chunk): + return + + try: + gen = stream_generate_content(envelope) + except UpstreamError as e: + # Streaming failed — fall back to non-streaming and return as a single chunk. + # This handles the case where the streamGenerateContent endpoint rejects + # tool-result round-trips with 400 "invalid argument". + if e.status == 400: + try: + resp = call_generate_content(envelope) + openai_resp = gemini_response_to_openai(resp, openai_model) + msg = openai_resp["choices"][0]["message"] + delta = {} + if msg.get("content"): + delta["content"] = msg["content"] + if msg.get("tool_calls"): + delta["tool_calls"] = msg["tool_calls"] + chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "created": created, + "model": openai_model, + "choices": [{ + "index": 0, + "delta": delta, + "finish_reason": openai_resp["choices"][0].get("finish_reason", "stop"), + }], + } + if openai_resp.get("usage"): + chunk["usage"] = openai_resp["usage"] + send_sse(chunk) + try: + self.wfile.write(b"data: [DONE]\n\n") + self.wfile.flush() + except Exception: + pass + return + except Exception: + pass # fall through to error chunk + err_chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "created": created, + "model": openai_model, + "choices": [{ + "index": 0, + "delta": {"content": f"[Error: {e}]"}, + "finish_reason": "stop", + }], + } + send_sse(err_chunk) + send_sse({"__done__": True}) # sentinel handled below + try: + self.wfile.write(b"data: [DONE]\n\n") + self.wfile.flush() + except Exception: + pass + return + except Exception as e: + err_chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "created": created, + "model": openai_model, + "choices": [{ + "index": 0, + "delta": {"content": f"[Error: {e}]"}, + "finish_reason": "stop", + }], + } + send_sse(err_chunk) + try: + self.wfile.write(b"data: [DONE]\n\n") + self.wfile.flush() + except Exception: + pass + return + + finish_reason = "stop" + seen_tool_names = set() + total_usage = None + events_processed = 0 + + try: + for event in gen: + if not isinstance(event, dict): + continue + if event.get("__done__"): + break + + events_processed += 1 + + # Unwrap envelope if present. + resp = event.get("response", event) + + # Check for usage metadata (often on the last chunk). + usage_meta = resp.get("usageMetadata") + if usage_meta: + total_usage = usage_meta + + candidates = resp.get("candidates") or [] + if not candidates: + continue + + cand = candidates[0] + content = cand.get("content") or {} + parts = content.get("parts") or [] + + text_delta, tool_call_deltas, seen_tool_names = _extract_delta_from_parts( + parts, seen_tool_names + ) + + # Determine finish reason. + fr = cand.get("finishReason") + if fr: + if tool_call_deltas: + finish_reason = "tool_calls" + else: + fr_up = str(fr).upper() + mapping = { + "STOP": "stop", + "MAX_TOKENS": "length", + "SAFETY": "content_filter", + "RECITATION": "content_filter", + "BLOCKLIST": "content_filter", + "PROHIBITED_CONTENT": "content_filter", + "SPII": "content_filter", + "IMAGE_SAFETY": "content_filter", + "LANGUAGE": "content_filter", + "OTHER": "stop", + } + finish_reason = mapping.get(fr_up, "stop") + + # Build delta. + delta = {} + if text_delta: + delta["content"] = text_delta + if tool_call_deltas: + delta["tool_calls"] = tool_call_deltas + + if not delta and not fr: + # Empty chunk with nothing — skip. + continue + + chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "created": created, + "model": openai_model, + "choices": [{ + "index": 0, + "delta": delta, + "finish_reason": None, + }], + } + if not send_sse(chunk): + return + + # Final chunk with finish_reason. + final_chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "created": created, + "model": openai_model, + "choices": [{ + "index": 0, + "delta": {}, + "finish_reason": finish_reason, + }], + } + # Include usage in the final chunk if the client requested + # stream_options.include_usage (OpenAI convention). + if total_usage: + pt = total_usage.get("promptTokenCount", 0) or 0 + ct = total_usage.get("candidatesTokenCount", 0) or total_usage.get("completionTokenCount", 0) or 0 + tt = total_usage.get("totalTokenCount", (pt + ct)) or (pt + ct) + final_chunk["usage"] = { + "prompt_tokens": pt, + "completion_tokens": ct, + "total_tokens": tt, + } + send_sse(final_chunk) + + # Termination sentinel. + try: + self.wfile.write(b"data: [DONE]\n\n") + self.wfile.flush() + except Exception: + pass + + except Exception as e: + _log(f"Streaming error: {e}") + # If the stream endpoint returned 400, try non-streaming as fallback. + # The streamGenerateContent endpoint sometimes rejects tool-result round-trips. + if "400" in str(e): + try: + resp = call_generate_content(envelope) + openai_resp = gemini_response_to_openai(resp, openai_model) + msg = openai_resp["choices"][0]["message"] + delta = {} + if msg.get("content"): + delta["content"] = msg["content"] + if msg.get("tool_calls"): + delta["tool_calls"] = msg["tool_calls"] + chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "created": created, + "model": openai_model, + "choices": [{ + "index": 0, + "delta": delta, + "finish_reason": openai_resp["choices"][0].get("finish_reason", "stop"), + }], + } + if openai_resp.get("usage"): + chunk["usage"] = openai_resp["usage"] + send_sse(chunk) + try: + self.wfile.write(b"data: [DONE]\n\n") + self.wfile.flush() + except Exception: + pass + return + except Exception: + pass # fall through to error chunk + # Try to send an error chunk then close. + try: + err_chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "created": created, + "model": openai_model, + "choices": [{ + "index": 0, + "delta": {"content": f"\n[stream error: {e}]"}, + "finish_reason": "stop", + }], + } + send_sse(err_chunk) + self.wfile.write(b"data: [DONE]\n\n") + self.wfile.flush() + except Exception: + pass + + +# ────────────────────────────────────────────────────────────────────────────── +# Server entry point +# ────────────────────────────────────────────────────────────────────────────── + +def _preflight_check(): + """Validate the token file exists and is readable before starting.""" + if not os.path.isfile(TOKEN_FILE): + _log(f"ERROR: Token file not found at {TOKEN_FILE}") + _log("Run the Antigravity CLI (`agy`) to authenticate first.") + return False + data = _read_token_from_disk() + if not data: + _log("ERROR: Token file exists but could not be parsed.") + return False + token_obj = data.get("token") or {} + if not token_obj.get("refresh_token"): + _log("ERROR: No refresh_token in token file. Re-run `agy` to authenticate.") + return False + _log(f"Token file OK (auth_method={data.get('auth_method', 'unknown')})") + return True + + +def main(): + parser = argparse.ArgumentParser( + description="Antigravity OpenAI-compatible proxy (Cloud Code Assist API)" + ) + parser.add_argument("--port", type=int, default=8877, help="Port to listen on (default 8877)") + parser.add_argument("--host", default="127.0.0.1", help="Host to bind to (default 127.0.0.1)") + args = parser.parse_args() + + if not _preflight_check(): + sys.exit(1) + + # Pre-discover the project ID so the first request is fast. + try: + get_project_id() + except Exception as e: + _log(f"WARNING: Could not pre-discover project ID: {e}") + _log("It will be discovered on the first request.") + + server = ThreadingHTTPServer((args.host, args.port), ProxyHandler) + server.daemon_threads = True + _log(f"Antigravity proxy listening on http://{args.host}:{args.port}/v1") + _log(f" Models: {', '.join(MODEL_MAP.keys())}") + _log(f" Token: {TOKEN_FILE}") + _log(f" Press Ctrl+C to stop") + + try: + server.serve_forever() + except KeyboardInterrupt: + _log("\nShutting down...") + server.shutdown() + + +if __name__ == "__main__": + main()