fix: Add mTLS and telemetry endpoint configurations to preview AdkApp

Telemetry exports from preview Reasoning Engine ADK applications were failing with 401 Unauthorized because the preview template lacked mTLS support for telemetry requests.

This fix ports the mTLS endpoint configuration logic (`_use_client_cert_effective()`, `_get_api_endpoint()`, and `configure_mtls_channel()`) from the production template, resolving telemetry issues for preview applications.

PiperOrigin-RevId: 944004265

Author: vertex-sdk-bot

tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py

@@ -26,6 +26,7 @@
 from google.cloud.aiplatform import initializer
 from vertexai.agent_engines import _utils
 from vertexai.preview import reasoning_engines
+from vertexai.preview.reasoning_engines.templates import adk as adk_template
 from google.genai import types
 import pytest
 import uuid
@@ -894,6 +895,7 @@ def test_tracing_setup(
         ):
             app.set_up()
 
+        import opentelemetry.sdk.version
         expected_attributes = {
             "cloud.account.id": _TEST_PROJECT_ID,
             "cloud.platform": "gcp.agent_engine",
@@ -906,8 +908,7 @@ def test_tracing_setup(
             "some-attribute": "some-value",
             "telemetry.sdk.language": "python",
             "telemetry.sdk.name": "opentelemetry",
-            "telemetry.sdk.version": "1.39.0",
-            "some-attribute": "some-value",
+            "telemetry.sdk.version": opentelemetry.sdk.version.__version__,
         }
 
         otlp_span_exporter_mock.assert_called_once_with(
@@ -1104,3 +1105,60 @@ async def test_bidi_stream_query_invalid_first_request(self):
         ):
             async for _ in app.bidi_stream_query(request_queue):
                 pass
+
+
+class TestAdkAppMtls:
+    """Test cases for mTLS functionality in preview AdkApp."""
+
+    def test_use_client_cert_effective_with_should_use_client_cert(self):
+        """Verifies that it respects the google-auth mTLS enablement check."""
+        with mock.patch.object(
+            adk_template.mtls,
+            "should_use_client_cert",
+            return_value=True,
+            create=True,
+        ):
+            assert adk_template._use_client_cert_effective() is True
+
+    @mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"})
+    def test_use_client_cert_effective_with_env_var_true(self):
+        """Verifies that it falls back to the environment variable if google-auth check fails."""
+        with mock.patch.object(
+            adk_template.mtls,
+            "should_use_client_cert",
+            side_effect=AttributeError,
+            create=True,
+        ):
+            assert adk_template._use_client_cert_effective() is True
+
+    @mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "false"})
+    def test_use_client_cert_effective_with_env_var_false(self):
+        """Verifies that it respects the environment variable being set to false."""
+        with mock.patch.object(
+            adk_template.mtls,
+            "should_use_client_cert",
+            side_effect=AttributeError,
+            create=True,
+        ):
+            assert adk_template._use_client_cert_effective() is False
+
+    def test_get_api_endpoint_default(self):
+        """Verifies the default telemetry endpoint is returned when no mTLS is configured."""
+        assert (
+            adk_template._get_api_endpoint() == adk_template._DEFAULT_TELEMETRY_ENDPOINT
+        )
+
+    @mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "always"})
+    def test_get_api_endpoint_always_with_cert(self):
+        """Verifies the mTLS endpoint is used when forced and a certificate is available."""
+        assert (
+            adk_template._get_api_endpoint(client_cert_source=b"cert")
+            == adk_template._DEFAULT_MTLS_TELEMETRY_ENDPOINT
+        )
+
+    @mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "auto"})
+    def test_get_api_endpoint_auto_no_cert(self):
+        """Verifies it falls back to regular endpoint even if forced if no certificate is provided."""
+        assert (
+            adk_template._get_api_endpoint() == adk_template._DEFAULT_TELEMETRY_ENDPOINT
+        )

vertexai/preview/reasoning_engines/templates/adk.py

@@ -27,8 +27,12 @@
 import asyncio
 from collections.abc import Awaitable
 import queue
+import os
 import sys
 import threading
+import enum
+from google.auth.transport import mtls
+from google.auth.transport import requests as requests_auth
 
 
 if TYPE_CHECKING:
@@ -97,6 +101,18 @@
 
 _DEFAULT_APP_NAME = "default-app-name"
 _DEFAULT_USER_ID = "default-user-id"
+_DEFAULT_TELEMETRY_ENDPOINT = "https://telemetry.googleapis.com/v1/traces"
+_DEFAULT_MTLS_TELEMETRY_ENDPOINT = "https://telemetry.mtls.googleapis.com/v1/traces"
+
+
+class _MtlsEndpoint(enum.Enum):
+    """The mTLS endpoint setting."""
+
+    AUTO = "auto"
+    ALWAYS = "always"
+    NEVER = "never"
+
+
 _TELEMETRY_API_DISABLED_WARNING = (
     "Tracing integration for Agent Engine has migrated to a new API.\n"
     "The 'telemetry.googleapis.com' has not been enabled in project %s. \n"
@@ -250,6 +266,64 @@ def _warn(msg: str):
     _warn._LOGGER.warning(msg)  # pyright: ignore[reportFunctionMemberAccess]
 
 
+def _get_api_endpoint(client_cert_source: bytes | None = None) -> str:
+    """Returns API endpoint based on mTLS configuration and cert availability.
+
+    Args:
+        client_cert_source (bytes | None): The client certificate source.
+
+    Returns:
+        str: The API endpoint to be used.
+    """
+    use_mtls_endpoint_str = os.getenv(
+        "GOOGLE_API_USE_MTLS_ENDPOINT", _MtlsEndpoint.AUTO.value
+    ).lower()
+
+    try:
+        use_mtls_endpoint = _MtlsEndpoint(use_mtls_endpoint_str)
+    except ValueError:
+        _warn(
+            f"Environment variable `GOOGLE_API_USE_MTLS_ENDPOINT` must be one of "
+            f"{[e.value for e in _MtlsEndpoint]}. Defaulting to"
+            f" {_MtlsEndpoint.AUTO.value}."
+        )
+        use_mtls_endpoint = _MtlsEndpoint.AUTO
+
+    if (use_mtls_endpoint == _MtlsEndpoint.ALWAYS) or (
+        use_mtls_endpoint == _MtlsEndpoint.AUTO and client_cert_source
+    ):
+        return _DEFAULT_MTLS_TELEMETRY_ENDPOINT
+
+    return _DEFAULT_TELEMETRY_ENDPOINT
+
+
+def _use_client_cert_effective() -> bool:
+    """Returns whether client certificate should be used for mTLS.
+
+    This checks if the google-auth version supports should_use_client_cert
+    automatic mTLS enablement. Alternatively, it reads from the
+    GOOGLE_API_USE_CLIENT_CERTIFICATE env var.
+
+    Returns:
+        bool: whether client certificate should be used for mTLS.
+    """
+    # check if google-auth version supports should_use_client_cert for automatic
+    # mTLS enablement
+    try:
+        return mtls.should_use_client_cert()
+    except (ImportError, AttributeError):
+        # if unsupported, fallback to reading from env var
+        use_client_cert_str = os.getenv(
+            "GOOGLE_API_USE_CLIENT_CERTIFICATE", "false"
+        ).lower()
+        if use_client_cert_str not in ("true", "false"):
+            _warn(
+                "Environment variable `GOOGLE_API_USE_CLIENT_CERTIFICATE` must be"
+                " either `true` or `false`"
+            )
+        return use_client_cert_str == "true"
+
+
 async def _force_flush_otel(tracing_enabled: bool, logging_enabled: bool):
     try:
         import opentelemetry.trace
@@ -393,12 +467,24 @@ def _detect_cloud_resource_id(project_id: str) -> Optional[str]:
         otlp_http_version = opentelemetry.exporter.otlp.proto.http.version.__version__
         user_agent = f"Vertex-Agent-Engine/{vertex_sdk_version} OTel-OTLP-Exporter-Python/{otlp_http_version}"
 
+        session = requests_auth.AuthorizedSession(credentials=credentials)
+
+        use_client_cert = _use_client_cert_effective()
+        if use_client_cert:
+            client_cert_source = (
+                mtls.default_client_cert_source()
+                if mtls.has_default_client_cert_source()
+                else None
+            )
+            session.configure_mtls_channel()
+            endpoint = _get_api_endpoint(client_cert_source)
+        else:
+            endpoint = _DEFAULT_TELEMETRY_ENDPOINT
+
         span_exporter = (
             opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter(
-                session=google.auth.transport.requests.AuthorizedSession(
-                    credentials=credentials
-                ),
-                endpoint="https://telemetry.googleapis.com/v1/traces",
+                session=session,
+                endpoint=endpoint,
                 headers={"User-Agent": user_agent},
             )
         )
@@ -1620,10 +1706,20 @@ def _warn_if_telemetry_api_disabled(self):
         except (ImportError, AttributeError):
             return
         credentials, project = google.auth.default()
-        session = google.auth.transport.requests.AuthorizedSession(
-            credentials=credentials
-        )
-        r = session.post("https://telemetry.googleapis.com/v1/traces", data=None)
+        session = requests_auth.AuthorizedSession(credentials=credentials)
+
+        use_client_cert = _use_client_cert_effective()
+        if use_client_cert:
+            client_cert_source = (
+                mtls.default_client_cert_source()
+                if mtls.has_default_client_cert_source()
+                else None
+            )
+            session.configure_mtls_channel()
+            endpoint = _get_api_endpoint(client_cert_source)
+        else:
+            endpoint = _DEFAULT_TELEMETRY_ENDPOINT
+        r = session.post(endpoint, data=None)
         if "Telemetry API has not been used in project" in r.text:
             _warn(_TELEMETRY_API_DISABLED_WARNING % (project, project))