feat(integrations): add google antigravity sdk tracing

py/noxfile.py

@@ -473,6 +473,17 @@ def test_claude_agent_sdk(session, version):
     _run_tests(session, f"{INTEGRATION_DIR}/claude_agent_sdk/test_claude_agent_sdk.py", version=version)
 
 
+ANTIGRAVITY_VERSIONS = _get_matrix_versions("google-antigravity")
+
+
+@nox.session()
+@nox.parametrize("version", ANTIGRAVITY_VERSIONS, ids=ANTIGRAVITY_VERSIONS)
+def test_antigravity(session, version):
+    _install_test_deps(session)
+    _install_matrix_dep(session, "google-antigravity", version)
+    _run_tests(session, f"{INTEGRATION_DIR}/antigravity/test_antigravity.py", version=version)
+
+
 CURSOR_SDK_VERSIONS = _get_matrix_versions("cursor-sdk")
 
 

py/pyproject.toml

@@ -536,6 +536,10 @@ latest = "claude-agent-sdk==0.2.179"
 latest = "cursor-sdk==1.3.1"
 "1.0.25" = "cursor-sdk==1.0.25"
 
+[tool.braintrust.matrix.google-antigravity]
+latest = "google-antigravity==0.1.23"
+"0.1.23" = "google-antigravity==0.1.23"
+
 [tool.braintrust.matrix.agno]
 latest = "agno==3.1.12"
 "2.4.0" = "agno==2.4.0"
@@ -699,6 +703,7 @@ agentscope = ["agentscope"]
 agno = ["agno"]
 autogen = ["autogen-agentchat"]
 anthropic = ["anthropic"]
+antigravity = ["google-antigravity"]
 bedrock_runtime = ["boto3", "botocore"]
 cohere = ["cohere"]
 claude_agent_sdk = ["claude-agent-sdk"]
@@ -741,6 +746,7 @@ crewai = "crewai"
 cursor-sdk = "cursor_sdk"
 dspy = "dspy"
 google-adk = "google.adk"
+google-antigravity = "google.antigravity"
 google-genai = "google.genai"
 google-cloud-discoveryengine = "google.cloud.discoveryengine_v1"
 instructor = "instructor"

py/src/braintrust/auto.py

@@ -47,6 +47,7 @@ def auto_instrument(
     agentscope: bool = True,
     claude_agent_sdk: bool = True,
     cursor_sdk: bool = True,
+    antigravity: bool = True,
     dspy: bool = True,
     adk: bool = True,
     langchain: bool = True,
@@ -92,6 +93,7 @@ def auto_instrument(
         agentscope: Enable AgentScope instrumentation (default: True)
         claude_agent_sdk: Enable Claude Agent SDK instrumentation (default: True)
         cursor_sdk: Enable Cursor SDK instrumentation (default: True)
+        antigravity: Enable Google Antigravity SDK instrumentation (default: True)
         dspy: Enable DSPy instrumentation (default: True)
         adk: Enable Google ADK instrumentation (default: True)
         langchain: Enable LangChain instrumentation (default: True)
@@ -198,6 +200,8 @@ def auto_instrument(
         results["claude_agent_sdk"] = _instrument_integration("claude_agent_sdk", "ClaudeAgentSDKIntegration")
     if cursor_sdk:
         results["cursor_sdk"] = _instrument_integration("cursor_sdk", "CursorSDKIntegration")
+    if antigravity:
+        results["antigravity"] = _instrument_integration("antigravity", "AntigravityIntegration")
     if dspy:
         results["dspy"] = _instrument_integration("dspy", "DSPyIntegration")
     if adk:

py/src/braintrust/integrations/__init__.py

@@ -10,6 +10,7 @@
     "AgnoIntegration": "agno",
     "AISDKIntegration": "ai_sdk",
     "AnthropicIntegration": "anthropic",
+    "AntigravityIntegration": "antigravity",
     "AutoGenIntegration": "autogen",
     "BedrockRuntimeIntegration": "bedrock_runtime",
     "ClaudeAgentSDKIntegration": "claude_agent_sdk",
@@ -45,6 +46,7 @@
     from .agno import AgnoIntegration as AgnoIntegration
     from .ai_sdk import AISDKIntegration as AISDKIntegration
     from .anthropic import AnthropicIntegration as AnthropicIntegration
+    from .antigravity import AntigravityIntegration as AntigravityIntegration
     from .autogen import AutoGenIntegration as AutoGenIntegration
     from .bedrock_runtime import BedrockRuntimeIntegration as BedrockRuntimeIntegration
     from .claude_agent_sdk import ClaudeAgentSDKIntegration as ClaudeAgentSDKIntegration
@@ -79,6 +81,7 @@
     "AgnoIntegration",
     "AISDKIntegration",
     "AnthropicIntegration",
+    "AntigravityIntegration",
     "AutoGenIntegration",
     "BedrockRuntimeIntegration",
     "ClaudeAgentSDKIntegration",

py/src/braintrust/integrations/antigravity/__init__.py

@@ -0,0 +1,34 @@
+"""Braintrust tracing for the Google Antigravity SDK.
+
+Instrument before creating an Agent, then consume its ChatResponse normally.
+"""
+
+from typing import Any
+
+from braintrust.logger import NOOP_SPAN, current_span, init_logger
+
+from .integration import AntigravityIntegration
+from .patchers import AgentPatcher
+
+
+__all__ = ["AntigravityIntegration", "setup_antigravity", "wrap_antigravity_agent"]
+
+
+def setup_antigravity(
+    api_key: str | None = None,
+    project_id: str | None = None,
+    project: str | None = None,
+) -> bool:
+    """Enable tracing for Antigravity chat responses."""
+    if current_span() == NOOP_SPAN:
+        init_logger(project=project, api_key=api_key, project_id=project_id)
+    return AntigravityIntegration.setup()
+
+
+def wrap_antigravity_agent(agent_class: Any) -> Any:
+    """Instrument the Conversation surface used by an Antigravity Agent.
+
+    Patches in place, including Agent classes imported before this call.
+    """
+    AntigravityIntegration.setup()
+    return AgentPatcher.wrap_target(agent_class)

py/src/braintrust/integrations/antigravity/_test_transport.py

@@ -0,0 +1,148 @@
+"""Test-only record/replay for the harness WebSocket, using real SDK parsers.
+
+HTTP VCR cannot intercept the bundled Go runtime's Gemini requests. Record
+actual WebSocket frames instead, then replay through LocalConnection's real
+event processor. No provider responses are synthesized.
+"""
+
+import asyncio
+import json
+import os
+from contextlib import contextmanager
+from pathlib import Path
+
+from braintrust.integrations.conftest import _versioned_cassette_dir
+from google.antigravity.connections.local.event_processor import parse_initialize_response
+from google.antigravity.connections.local.local_connection import LocalConnection, LocalConnectionStrategy
+from google.antigravity.proto import localharness_pb2
+from google.protobuf import json_format
+
+
+def cassette_path(name):
+    base = Path(os.environ.get("BRAINTRUST_INTEGRATIONS_DIR", Path(__file__).resolve().parent.parent))
+    return Path(_versioned_cassette_dir(str(base / "antigravity" / "cassettes"))) / f"{name}.json"
+
+
+def _sanitize(value):
+    if isinstance(value, dict):
+        return {k: "REDACTED" if k in ("apiKey", "api_key") else _sanitize(v) for k, v in value.items()}
+    if isinstance(value, list):
+        return [_sanitize(v) for v in value]
+    if isinstance(value, str):
+        for key in ("GEMINI_API_KEY", "GOOGLE_API_KEY"):
+            secret = os.environ.get(key)
+            if secret:
+                value = value.replace(secret, "REDACTED")
+        value = value.replace(str(Path.cwd()), "WORKSPACE")
+        if value.startswith(("/private/var/folders/", "/tmp/pytest-", "C:\\Users\\")):
+            value = "TEST_WORKSPACE"
+    return value
+
+
+class _RecordingSocket:
+    def __init__(self, ws, frames):
+        self.ws, self.frames = ws, frames
+
+    async def send(self, data):
+        self.frames.append({"send": _sanitize(json.loads(data))})
+        await self.ws.send(data)
+
+    async def recv(self):
+        data = await self.ws.recv()
+        self.frames.append({"recv": _sanitize(json.loads(data))})
+        return data
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        return await self.recv()
+
+    async def close(self):
+        await self.ws.close()
+
+
+class _ReplaySocket:
+    def __init__(self, frames):
+        self.frames = iter(frames)
+        self.next_frame = next(self.frames, None)
+        self.changed = asyncio.Event()
+
+    def _advance(self):
+        self.next_frame = next(self.frames, None)
+        self.changed.set()
+
+    async def send(self, data):
+        assert self.next_frame is not None and "send" in self.next_frame
+        actual = json.loads(data)
+        expected = self.next_frame["send"]
+        # Handshake has machine-specific workspace, credentials and config.
+        # Prompts must match exactly; other events retain their protocol type.
+        assert actual.keys() == expected.keys()
+        if "userInput" in actual:
+            assert actual == expected
+        self._advance()
+
+    async def recv(self):
+        while self.next_frame is None or "recv" not in self.next_frame:
+            self.changed.clear()
+            await self.changed.wait()
+        data = json.dumps(self.next_frame["recv"])
+        self._advance()
+        await asyncio.sleep(0)
+        return data
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        return await self.recv()
+
+    async def close(self):
+        pass
+
+
+@contextmanager
+def cassette_transport(name):
+    path = cassette_path(name)
+    mode = os.environ.get("BRAINTRUST_ANTIGRAVITY_RECORD_MODE", "none" if os.getenv("CI") else "once")
+    recording = mode == "all" or (mode == "once" and not path.exists())
+    original_spawn = LocalConnectionStrategy._spawn_harness_and_connect_ws
+    original_enter = LocalConnectionStrategy.__aenter__
+    frames = []
+
+    async def spawn(strategy, **kwargs):
+        process, ws, url = await original_spawn(strategy, **kwargs)
+        return process, _RecordingSocket(ws, frames), url
+
+    async def enter(strategy):
+        ws = _ReplaySocket(json.loads(path.read_text(encoding="utf-8")))
+        init = localharness_pb2.InitializeConversationEvent(config=strategy._build_harness_config())  # pylint: disable=no-member
+        await ws.send(json_format.MessageToJson(init))
+        event = localharness_pb2.OutputEvent()  # pylint: disable=no-member
+        json_format.Parse(await ws.recv(), event)
+        parsed = parse_initialize_response(event.initialize_conversation_response)
+        strategy._connection = LocalConnection(
+            process=None,
+            ws=ws,
+            tool_runner=strategy._tool_runner,
+            hook_runner=strategy._hook_runner,
+            initial_history=parsed.history,
+            initial_usage=parsed.cumulative_usage,
+            initial_trajectory_usages=parsed.trajectory_usages,
+        )
+
+    if recording:
+        LocalConnectionStrategy._spawn_harness_and_connect_ws = spawn
+    else:
+        if not path.exists():
+            raise FileNotFoundError(f"Missing Antigravity recording: {path}")
+        LocalConnectionStrategy.__aenter__ = enter
+    try:
+        yield
+    finally:
+        LocalConnectionStrategy.__aenter__ = original_enter
+        LocalConnectionStrategy._spawn_harness_and_connect_ws = original_spawn
+    if recording:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(json.dumps(frames, indent=2) + "\n", encoding="utf-8")