99 lines
3.4 KiB
Python
99 lines
3.4 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import platform
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
from typing import Any, Callable
|
|
|
|
from .errors import ProxyError
|
|
|
|
CLOUD_CODE_ENDPOINT = "https://cloudcode-pa.googleapis.com"
|
|
ANTIGRAVITY_LOAD_CODE_ASSIST_METADATA = {
|
|
"ideType": "ANTIGRAVITY",
|
|
"platform": "PLATFORM_UNSPECIFIED",
|
|
"pluginType": "GEMINI",
|
|
}
|
|
TIER_LEGACY = "legacy-tier"
|
|
PROJECT_ONBOARD_MAX_ATTEMPTS = 5
|
|
PROJECT_ONBOARD_INTERVAL_SECONDS = 2
|
|
|
|
|
|
def antigravity_user_agent() -> str:
|
|
version = os.getenv("PI_AI_ANTIGRAVITY_VERSION") or "2.1.4"
|
|
system = platform.system().lower()
|
|
os_name = "windows" if system.startswith("win") else ("darwin" if system == "darwin" else system or "linux")
|
|
machine = platform.machine().lower()
|
|
arch = "amd64" if machine in {"x86_64", "x64"} else ("386" if machine in {"i386", "i686"} else machine or "arm64")
|
|
return f"antigravity/hub/{version} {os_name}/{arch}"
|
|
|
|
|
|
def _post_json(url: str, body: dict[str, Any], headers: dict[str, str]) -> dict[str, Any]:
|
|
req = urllib.request.Request(
|
|
url,
|
|
data=json.dumps(body).encode("utf-8"),
|
|
headers=headers,
|
|
method="POST",
|
|
)
|
|
try:
|
|
with urllib.request.urlopen(req, timeout=120) as resp:
|
|
return json.loads(resp.read().decode("utf-8") or "{}")
|
|
except urllib.error.HTTPError as e:
|
|
detail = e.read().decode("utf-8", "replace")
|
|
raise ProxyError(f"Cloud Code Assist API error ({e.code}): {detail}", status=e.code) from e
|
|
|
|
|
|
def read_project_id(value: object) -> str | None:
|
|
if isinstance(value, str) and value:
|
|
return value
|
|
if isinstance(value, dict) and isinstance(value.get("id"), str) and value["id"]:
|
|
return value["id"]
|
|
return None
|
|
|
|
|
|
def read_default_tier(allowed_tiers: object) -> str:
|
|
if not isinstance(allowed_tiers, list):
|
|
return TIER_LEGACY
|
|
for tier in allowed_tiers:
|
|
if isinstance(tier, dict) and tier.get("isDefault") and isinstance(tier.get("id"), str) and tier["id"]:
|
|
return tier["id"]
|
|
return TIER_LEGACY
|
|
|
|
|
|
def load_or_onboard_project(
|
|
access_token: str,
|
|
*,
|
|
post_json: Callable[[str, dict[str, Any], dict[str, str]], dict[str, Any]] | None = None,
|
|
sleep: Callable[[float], None] = time.sleep,
|
|
) -> str:
|
|
post_json = post_json or _post_json
|
|
headers = {
|
|
"Authorization": f"Bearer {access_token}",
|
|
"Content-Type": "application/json",
|
|
"User-Agent": antigravity_user_agent(),
|
|
}
|
|
load_payload = post_json(
|
|
f"{CLOUD_CODE_ENDPOINT}/v1internal:loadCodeAssist",
|
|
{"metadata": ANTIGRAVITY_LOAD_CODE_ASSIST_METADATA},
|
|
headers,
|
|
)
|
|
existing = read_project_id(load_payload.get("cloudaicompanionProject"))
|
|
if existing:
|
|
return existing
|
|
|
|
onboard_body = {
|
|
"tierId": read_default_tier(load_payload.get("allowedTiers")),
|
|
"metadata": ANTIGRAVITY_LOAD_CODE_ASSIST_METADATA,
|
|
}
|
|
for attempt in range(1, PROJECT_ONBOARD_MAX_ATTEMPTS + 1):
|
|
if attempt > 1:
|
|
sleep(PROJECT_ONBOARD_INTERVAL_SECONDS)
|
|
op = post_json(f"{CLOUD_CODE_ENDPOINT}/v1internal:onboardUser", onboard_body, headers)
|
|
if not op.get("done"):
|
|
continue
|
|
project_id = read_project_id((op.get("response") or {}).get("cloudaicompanionProject"))
|
|
if project_id:
|
|
return project_id
|
|
raise ProxyError("onboardUser did not return a project id", status=502)
|