diff --git a/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py b/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py index e943ceee96..dad1ee1fc7 100644 --- a/tests/unit/vertex_adk/test_reasoning_engine_templates_adk.py +++ b/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 + ) diff --git a/vertexai/preview/reasoning_engines/templates/adk.py b/vertexai/preview/reasoning_engines/templates/adk.py index 469ecf71cc..41a4c164e3 100644 --- a/vertexai/preview/reasoning_engines/templates/adk.py +++ b/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))