feat: Add mTLS support to Google Cloud Telemetry exporter

This change enables the Google Cloud Telemetry exporter to use mTLS endpoints. It checks for the availability of client certificates and respects the GOOGLE_API_USE_CLIENT_CERTIFICATE environment variables to determine whether to use the mTLS-specific endpoint and configure the session accordingly.

PiperOrigin-RevId: 911581237

Author: google-genai-bot

src/google/adk/telemetry/google_cloud.py

@@ -14,13 +14,17 @@
 
 from __future__ import annotations
 
+import enum
 import logging
 import os
+from typing import Any
+from typing import Callable
 from typing import cast
 from typing import Optional
 from typing import TYPE_CHECKING
 
 import google.auth
+from google.auth.transport import mtls
 from opentelemetry.sdk._logs import LogRecordProcessor
 from opentelemetry.sdk._logs.export import BatchLogRecordProcessor
 from opentelemetry.sdk.metrics.export import MetricReader
@@ -40,6 +44,19 @@
 _GCP_LOG_NAME_ENV_VARIABLE_NAME = 'GOOGLE_CLOUD_DEFAULT_LOG_NAME'
 _DEFAULT_LOG_NAME = 'adk-otel'
 
+_DEFAULT_TELEMETRY_TRACES_ENPOINT = 'https://telemetry.googleapis.com/v1/traces'
+_DEFAULT_MTLS_TELEMETRY_TRACES_ENPOINT = (
+    'https://telemetry.mtls.googleapis.com/v1/traces'
+)
+
+
+class _MtlsEndpoint(enum.Enum):
+  """The mTLS endpoint setting."""
+
+  AUTO = 'auto'
+  ALWAYS = 'always'
+  NEVER = 'never'
+
 
 def get_gcp_exporters(
     enable_cloud_tracing: bool = False,
@@ -100,10 +117,24 @@ def _get_gcp_span_exporter(credentials: Credentials) -> SpanProcessor:
   from google.auth.transport.requests import AuthorizedSession
   from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 
+  session = 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_TRACES_ENPOINT
+
   return BatchSpanProcessor(
       OTLPSpanExporter(
-          session=AuthorizedSession(credentials=credentials),
-          endpoint='https://telemetry.googleapis.com/v1/traces',
+          session=session,
+          endpoint=endpoint,
       )
   )
 
@@ -158,3 +189,62 @@ def get_gcp_resource(project_id: Optional[str] = None) -> Resource:
         ' GCE, GKE or CloudRun related resource attributes may be missing'
     )
   return resource
+
+
+def _get_api_endpoint(
+    client_cert_source: Callable[[], tuple[bytes, bytes]] | None = None,
+) -> str:
+  """Returns API endpoint based on mTLS configuration and cert availability.
+
+  Args:
+      client_cert_source: A callable that returns the client certificate and
+        key, or None.
+
+  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:
+    logger.warning(
+        'Environment variable `GOOGLE_API_USE_MTLS_ENDPOINT` must be one of '
+        '%s. Defaulting to %s.',
+        [e.value for e in _MtlsEndpoint],
+        _MtlsEndpoint.AUTO.value,
+    )
+    use_mtls_endpoint = _MtlsEndpoint.AUTO
+
+  if (use_mtls_endpoint is _MtlsEndpoint.ALWAYS) or (
+      use_mtls_endpoint is _MtlsEndpoint.AUTO and client_cert_source
+  ):
+    return _DEFAULT_MTLS_TELEMETRY_TRACES_ENPOINT
+
+  return _DEFAULT_TELEMETRY_TRACES_ENPOINT
+
+
+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.
+  """
+  try:
+    return bool(mtls.should_use_client_cert())
+  except (ImportError, AttributeError):
+    use_client_cert_str = os.getenv(
+        'GOOGLE_API_USE_CLIENT_CERTIFICATE', 'false'
+    ).lower()
+    if use_client_cert_str not in ('true', 'false'):
+      logger.warning(
+          'Environment variable `GOOGLE_API_USE_CLIENT_CERTIFICATE` must be'
+          ' either `true` or `false`'
+      )
+    return use_client_cert_str == 'true'

tests/unittests/telemetry/test_google_cloud.py

@@ -16,8 +16,18 @@
 from typing import Optional
 from unittest import mock
 
+from google.adk.telemetry import google_cloud
+from google.adk.telemetry.google_cloud import _DEFAULT_MTLS_TELEMETRY_TRACES_ENPOINT
+from google.adk.telemetry.google_cloud import _DEFAULT_TELEMETRY_TRACES_ENPOINT
+from google.adk.telemetry.google_cloud import _get_api_endpoint
+from google.adk.telemetry.google_cloud import _get_gcp_span_exporter
+from google.adk.telemetry.google_cloud import _use_client_cert_effective
 from google.adk.telemetry.google_cloud import get_gcp_exporters
 from google.adk.telemetry.google_cloud import get_gcp_resource
+import google.auth.credentials
+from google.auth.transport import mtls
+from google.auth.transport import requests
+from opentelemetry.exporter.otlp.proto.http import trace_exporter
 import pytest
 
 
@@ -89,3 +99,108 @@ def test_get_gcp_resource(
       otel_resource.attributes.get("gcp.project_id", None)
       == expected_project_id
   )
+
+
+@mock.patch.object(mtls, "should_use_client_cert", autospec=True)
+def test_use_client_cert_effective_from_mtls(mock_should_use):
+  mock_should_use.return_value = True
+  assert _use_client_cert_effective()
+
+  mock_should_use.return_value = False
+  assert not _use_client_cert_effective()
+
+
+def test_use_client_cert_effective_from_env(
+    monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+):
+  with mock.patch.object(
+      mtls,
+      "should_use_client_cert",
+      autospec=True,
+      side_effect=AttributeError,
+  ):
+    monkeypatch.setenv("GOOGLE_API_USE_CLIENT_CERTIFICATE", "true")
+    assert _use_client_cert_effective()
+
+    monkeypatch.setenv("GOOGLE_API_USE_CLIENT_CERTIFICATE", "false")
+    assert not _use_client_cert_effective()
+
+    # Test invalid value defaults to False
+    monkeypatch.setenv("GOOGLE_API_USE_CLIENT_CERTIFICATE", "maybe")
+    assert not _use_client_cert_effective()
+    assert (
+        "Environment variable `GOOGLE_API_USE_CLIENT_CERTIFICATE` must be"
+        " either `true` or `false`"
+        in caplog.text
+    )
+
+
+@pytest.mark.parametrize(
+    "env_val, cert_source, expected",
+    [
+        ("auto", lambda: b"cert", _DEFAULT_MTLS_TELEMETRY_TRACES_ENPOINT),
+        ("auto", None, _DEFAULT_TELEMETRY_TRACES_ENPOINT),
+        ("always", None, _DEFAULT_MTLS_TELEMETRY_TRACES_ENPOINT),
+        ("never", lambda: b"cert", _DEFAULT_TELEMETRY_TRACES_ENPOINT),
+        ("invalid", None, _DEFAULT_TELEMETRY_TRACES_ENPOINT),
+    ],
+)
+def test_get_api_endpoint(
+    env_val,
+    cert_source,
+    expected,
+    monkeypatch: pytest.MonkeyPatch,
+    caplog: pytest.LogCaptureFixture,
+):
+  monkeypatch.setenv("GOOGLE_API_USE_MTLS_ENDPOINT", env_val)
+  if env_val == "invalid":
+    assert _get_api_endpoint(cert_source) == expected
+    assert (
+        "Environment variable `GOOGLE_API_USE_MTLS_ENDPOINT` must be one of"
+        in caplog.text
+    )
+  else:
+    assert _get_api_endpoint(cert_source) == expected
+
+
+@mock.patch.object(requests, "AuthorizedSession", autospec=True)
+@mock.patch(
+    "opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter",
+    autospec=True,
+)
+@mock.patch(
+    "google.adk.telemetry.google_cloud.BatchSpanProcessor", autospec=True
+)
+@mock.patch(
+    "google.adk.telemetry.google_cloud._use_client_cert_effective",
+    autospec=True,
+)
+@mock.patch(
+    "google.auth.transport.mtls.has_default_client_cert_source", autospec=True
+)
+@mock.patch(
+    "google.auth.transport.mtls.default_client_cert_source", autospec=True
+)
+def test_get_gcp_span_exporter_mtls(
+    mock_default_cert: mock.MagicMock,
+    mock_has_cert: mock.MagicMock,
+    mock_use_cert: mock.MagicMock,
+    mock_batch: mock.MagicMock,
+    mock_exporter: mock.MagicMock,
+    mock_session: mock.MagicMock,
+):
+  credentials = mock.create_autospec(
+      google.auth.credentials.Credentials, instance=True
+  )
+  mock_use_cert.return_value = True
+  mock_has_cert.return_value = True
+  mock_default_cert.return_value = b"cert"
+
+  _get_gcp_span_exporter(credentials)
+
+  mock_session.assert_called_once_with(credentials=credentials)
+  mock_session.return_value.configure_mtls_channel.assert_called_once()
+  mock_exporter.assert_called_once_with(
+      session=mock_session.return_value,
+      endpoint=_DEFAULT_MTLS_TELEMETRY_TRACES_ENPOINT,
+  )