feat(integrations): add discovery engine v1 instrumentation

.agents/skills/sdk-integrations/SKILL.md

@@ -263,6 +263,28 @@ Assert on emitted spans (not just provider return values):
 
 For streaming, assert both the provider iterator/async-iterator still works AND the final span has aggregated `output` + stream-specific `metrics`.
 
+### Streaming lifecycle review
+
+Review every exit path, not just full consumption. Check the real provider's
+iterator, context-manager, cancellation, and garbage-collection behavior before
+choosing a wrapper. A proxy can preserve transport cleanup while still losing
+the span's final output and end time.
+
+- Cover exhaustion, provider errors, explicit close/cancel, and context-manager
+  exit where the provider supports it. Preserve exception and return semantics.
+- Cover `break` followed by dropping the last stream reference, and dropping a
+  stream before consuming any chunks. `break` alone does not close a retained
+  iterator; do not promise immediate finalization while callers still hold it.
+- Use existing recordings for partial-consumption tests. Drop the proxy, force
+  collection, and assert partial (or empty) output, `metrics.end`, correct
+  parentage, and no duplicate finalization after explicit close or exhaustion.
+- A GC fallback must not retain the stream through its callback or closure.
+  For async streams, do not run or schedule event-loop work from a finalizer;
+  finalize trace state and preserve the provider's own cleanup behavior.
+- Check that neither iteration nor cleanup leaves the stream span current in
+  the caller's context. GC is a best-effort fallback, not a substitute for
+  deterministic cleanup when the caller explicitly closes the stream.
+
 Cassettes live in `integrations/<provider>/cassettes/<version>/` (e.g. `cassettes/latest/`, `cassettes/0.48.0/`). Nox sets `BRAINTRUST_TEST_PACKAGE_VERSION` so cassettes land correctly. Do not add per-test `vcr_cassette_dir` / `cassette_library_dir` fixtures — `integrations/conftest.py` handles it. Re-record only when behavior intentionally changed. Sanitize binary media in both request and response bodies so checked-in cassettes do not retain large base64 payloads.
 
 Confirm the exact session name from `noxfile.py` — don't assume it matches the folder.

.agents/skills/sdk-vcr-workflows/SKILL.md

@@ -266,6 +266,82 @@ Important differences:
 
 Do not try to force ordinary HTTP VCR patterns onto Claude Agent SDK subprocess tests.
 
+## Discovery Engine Recording
+
+Discovery Engine tests use HTTP VCR for sync REST calls and the test-only
+`integrations/google_discoveryengine/_test_grpc.py` helper for async gRPC calls. Both
+recording formats live under `py/src/braintrust/integrations/google_discoveryengine/cassettes/<version>/`.
+
+### Prerequisites
+
+- Install `gcloud` and use a Google Cloud project with billing and the
+  `discoveryengine.googleapis.com` API enabled.
+- Authenticate with Application Default Credentials (ADC). The account needs
+  permission to invoke Discovery Engine in the target project. A `GEMINI_API_KEY`
+  alone does not authenticate these tests.
+- For answer and conversation tests, create a search app with generative responses
+  enabled and attach a populated datastore. Wait for indexing to finish before
+  recording. Ranking and grounding-check tests do not require indexed documents.
+
+```sh
+gcloud config set project YOUR_PROJECT_ID
+gcloud auth application-default login
+gcloud auth application-default set-quota-project YOUR_PROJECT_ID
+```
+
+The recording fixture obtains an access token using
+`gcloud auth application-default print-access-token`, outside the recorded call.
+Playback uses anonymous credentials and needs no Google account.
+
+Live recording requires explicit resource IDs through environment variables.
+There are no private-project defaults in the test code. Tests use the `global`
+location. Ranking and grounding checks require the project; answer and
+conversation coverage also requires the app and datastore.
+
+A datastore can be populated from Google's public Alphabet earnings-report
+sample PDFs at `gs://cloud-samples-data/gen-app-builder/search/alphabet-investor-pdfs`.
+Wait for indexing to finish before recording.
+
+```sh
+export BRAINTRUST_GOOGLE_DISCOVERYENGINE_PROJECT="your-project"
+export BRAINTRUST_GOOGLE_DISCOVERYENGINE_APP="your-app-id"
+export BRAINTRUST_GOOGLE_DISCOVERYENGINE_DATASTORE="your-datastore-id"
+```
+
+### Record and replay
+
+From `py/`, select a focused scenario. `--vcr-record=all` enables recording for
+both REST and gRPC; the gRPC helper otherwise requires an existing cassette.
+
+```sh
+# REST ranking, including manual/setup entry-point coverage.
+mise exec -- nox -s 'test_google_discoveryengine(latest)' -- --vcr-record=all -k 'test_rank and not test_rank_output_limit'
+
+# Async gRPC ranking.
+mise exec -- nox -s 'test_google_discoveryengine(latest)' -- --vcr-record=all -k 'test_async_grpc and rank'
+
+# Verify all recordings without network access to Google.
+mise exec -- nox -R -s 'test_google_discoveryengine(latest)' -- --vcr-record=none
+```
+
+Playback derives resource paths from the checked-in REST cassettes and ignores
+resource environment overrides. No access to the recorded project's resources is
+needed. The auto-instrument subprocess test also reads its ranking resource from
+the cassette. Keep the project's recordings consistent when recording against a
+different project.
+
+The gRPC helper temporarily replaces one callable on one client transport,
+records real protobuf requests/responses and errors as JSON, then restores the
+callable. Playback reconstructs provider response objects and checks requests;
+it does not exercise networking, retries, or real cancellation behavior.
+HTTP authorization headers are filtered by shared VCR configuration; gRPC
+recordings omit credentials and transport metadata. Inspect new recordings before
+checking them in.
+
+`generate_grounded_content` and `stream_generate_grounded_content` are deferred:
+the v1 endpoints returned method-not-found errors during initial recording.
+Do not fabricate success recordings for those methods.
+
 ## Relationship To Other Skills
 
 - Use `sdk-integrations` when the main task is integration implementation, patchers, tracing, or provider package structure.

py/noxfile.py

@@ -569,6 +569,17 @@ def test_google_genai(session, version):
     _run_tests(session, f"{INTEGRATION_DIR}/google_genai/test_google_genai.py", version=version)
 
 
+GOOGLE_DISCOVERYENGINE_VERSIONS = _get_matrix_versions("google-cloud-discoveryengine")
+
+
+@nox.session()
+@nox.parametrize("version", GOOGLE_DISCOVERYENGINE_VERSIONS, ids=GOOGLE_DISCOVERYENGINE_VERSIONS)
+def test_google_discoveryengine(session, version):
+    _install_test_deps(session)
+    _install_matrix_dep(session, "google-cloud-discoveryengine", version)
+    _run_tests(session, f"{INTEGRATION_DIR}/google_discoveryengine", version=version)
+
+
 DSPY_VERSIONS = _get_matrix_versions("dspy")
 
 

py/pyproject.toml

@@ -267,6 +267,7 @@ lint = [
     "dspy",
     "google-adk",
     "google-genai",
+    "google-cloud-discoveryengine",
     "instructor",
     "litellm>=1.83.10",
     "livekit-agents",
@@ -462,6 +463,9 @@ latest = "google-genai==2.23.0"
 "1.75.0" = "google-genai==1.75.0"
 "1.30.0" = "google-genai==1.30.0"
 
+[tool.braintrust.matrix.google-cloud-discoveryengine]
+latest = "google-cloud-discoveryengine==0.20.3"
+
 [tool.braintrust.matrix.dspy]
 latest = "dspy==3.3.1"
 "2.6.0" = "dspy==2.6.0"
@@ -571,6 +575,7 @@ cursor_sdk = ["cursor-sdk"]
 crewai = ["crewai"]
 dspy = ["dspy"]
 google_genai = ["google-genai"]
+google_discoveryengine = ["google-cloud-discoveryengine"]
 huggingface_hub = ["huggingface-hub"]
 harbor = ["harbor"]
 instructor = ["instructor"]
@@ -604,6 +609,7 @@ cursor-sdk = "cursor_sdk"
 dspy = "dspy"
 google-adk = "google.adk"
 google-genai = "google.genai"
+google-cloud-discoveryengine = "google.cloud.discoveryengine_v1"
 litellm = "litellm"
 livekit-agents = "livekit.agents"
 mistralai = "mistralai"

py/src/braintrust/auto.py

@@ -20,6 +20,7 @@
     CrewAIIntegration,
     CursorSDKIntegration,
     DSPyIntegration,
+    GoogleDiscoveryEngineIntegration,
     GoogleGenAIIntegration,
     HuggingFaceHubIntegration,
     InstructorIntegration,
@@ -64,6 +65,7 @@ def auto_instrument(
     ai_sdk: bool = True,
     pydantic_ai: bool = True,
     google_genai: bool = True,
+    google_discoveryengine: bool = True,
     instructor: bool = True,
     openrouter: bool = True,
     mistral: bool = True,
@@ -102,6 +104,7 @@ def auto_instrument(
         litellm: Enable LiteLLM instrumentation (default: True)
         ai_sdk: Enable Vercel AI SDK for Python instrumentation (default: True)
         pydantic_ai: Enable Pydantic AI instrumentation (default: True)
+        google_discoveryengine: Enable Google Discovery Engine v1 instrumentation (default: True)
         google_genai: Enable Google GenAI instrumentation (default: True)
         instructor: Enable Instructor (structured-output) instrumentation (default: True)
         openrouter: Enable OpenRouter instrumentation (default: True)
@@ -184,6 +187,8 @@ def auto_instrument(
         results["pydantic_ai"] = _instrument_integration(PydanticAIIntegration)
     if google_genai:
         results["google_genai"] = _instrument_integration(GoogleGenAIIntegration)
+    if google_discoveryengine:
+        results["google_discoveryengine"] = _instrument_integration(GoogleDiscoveryEngineIntegration)
     if instructor:
         results["instructor"] = _instrument_integration(InstructorIntegration)
     if openrouter:

py/src/braintrust/integrations/__init__.py

@@ -10,6 +10,7 @@
 from .crewai import CrewAIIntegration
 from .cursor_sdk import CursorSDKIntegration
 from .dspy import DSPyIntegration
+from .google_discoveryengine import GoogleDiscoveryEngineIntegration
 from .google_genai import GoogleGenAIIntegration
 from .huggingface_hub import HuggingFaceHubIntegration
 from .instructor import InstructorIntegration
@@ -41,6 +42,7 @@
     "CrewAIIntegration",
     "CursorSDKIntegration",
     "DSPyIntegration",
+    "GoogleDiscoveryEngineIntegration",
     "GoogleGenAIIntegration",
     "HuggingFaceHubIntegration",
     "InstructorIntegration",

py/src/braintrust/integrations/auto_test_scripts/test_auto_google_discoveryengine.py

@@ -0,0 +1,65 @@
+"""Both import orders and opt-out, using the real REST ranking cassette."""
+
+import inspect
+import subprocess
+import sys
+from pathlib import Path
+from urllib.parse import urlsplit
+
+import yaml
+from braintrust.auto import auto_instrument
+from braintrust.integrations.conftest import _versioned_cassette_dir
+from braintrust.integrations.test_utils import autoinstrument_test_context
+
+
+if len(sys.argv) == 1:
+    for order in ("before", "after"):
+        subprocess.run([sys.executable, __file__, order], check=True)
+    print("SUCCESS")
+    sys.exit(0)
+
+options = {name: False for name in inspect.signature(auto_instrument).parameters}
+assert auto_instrument(**options) == {}
+RankServiceClient = None
+if sys.argv[1] == "before":
+    from google.cloud.discoveryengine_v1 import RankServiceClient
+
+options["google_discoveryengine"] = True
+assert auto_instrument(**options) == {"google_discoveryengine": True}
+assert auto_instrument(**options) == {"google_discoveryengine": True}
+from google.auth.credentials import AnonymousCredentials
+
+
+if sys.argv[1] == "after":
+    from google.cloud.discoveryengine_v1 import RankServiceClient
+
+
+cassette_dir = Path(
+    _versioned_cassette_dir(str(Path(__file__).parent.parent / "google_discoveryengine" / "cassettes"))
+)
+cassette = yaml.safe_load((cassette_dir / "test_rank.yaml").read_text())
+ranking_config = urlsplit(cassette["interactions"][0]["request"]["uri"]).path.removeprefix("/v1/").split(":rank")[0]
+
+assert RankServiceClient is not None
+with autoinstrument_test_context(
+    "test_rank", integration="google_discoveryengine", vcr_config={"record_mode": "none"}
+) as memory_logger:
+    client = RankServiceClient(transport="rest", credentials=AnonymousCredentials())
+    result = client.rank(
+        request={
+            "ranking_config": ranking_config,
+            "model": "semantic-ranker-512@latest",
+            "query": "What is Braintrust?",
+            "records": [
+                {"id": "1", "content": "Braintrust is a platform for evaluating and monitoring AI applications."},
+                {"id": "2", "content": "The moon orbits the Earth."},
+            ],
+            "top_n": 1,
+        },
+        retry=None,
+    )
+    assert result.records[0].id == "1"
+    spans = memory_logger.pop()
+    assert len(spans) == 1
+    assert spans[0]["metadata"]["provider"] == "google"
+    assert spans[0]["context"]["span_origin"]["instrumentation"]["name"] == "google-discoveryengine-auto"

py/src/braintrust/integrations/google_discoveryengine/__init__.py

@@ -0,0 +1,12 @@
+"""Braintrust integration for google-cloud-discoveryengine v1."""
+
+from .integration import GoogleDiscoveryEngineIntegration
+from .patchers import wrap_google_discoveryengine
+
+
+__all__ = ["GoogleDiscoveryEngineIntegration", "setup_google_discoveryengine", "wrap_google_discoveryengine"]
+
+
+def setup_google_discoveryengine() -> bool:
+    """Instrument supported v1 clients in this process."""
+    return GoogleDiscoveryEngineIntegration.setup()