diff --git a/agentplatform/agent_engines/templates/adk.py b/agentplatform/agent_engines/templates/adk.py index ff3aecbb90..ce9c9bb967 100644 --- a/agentplatform/agent_engines/templates/adk.py +++ b/agentplatform/agent_engines/templates/adk.py @@ -128,6 +128,10 @@ _DEFAULT_TELEMETRY_ENDPOINT = "https://telemetry.googleapis.com/v1/traces" _DEFAULT_MTLS_TELEMETRY_ENDPOINT = "https://telemetry.mtls.googleapis.com/v1/traces" +# Timeout (in seconds) for the best-effort Telemetry API enablement check. This +# check runs on the agent server's critical startup path, so it must fail fast +# rather than inherit AuthorizedSession's 120s default. +_TELEMETRY_API_CHECK_TIMEOUT_SECONDS = 5.0 class _MtlsEndpoint(enum.Enum): @@ -601,22 +605,34 @@ def _validate_run_config(run_config: Optional[Dict[str, Any]]): def _warn_if_telemetry_api_disabled(): - """Warn if telemetry API is disabled.""" - credentials, project = google.auth.default() - 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 + """Warns if the Telemetry API is disabled. + + This is a best-effort diagnostic that runs from `set_up()`, on the agent + server's critical startup path. No failure here (blocked egress, an + unreachable endpoint, unavailable credentials) may prevent the container + from starting, so the whole body is non-fatal. + """ + try: + credentials, project = google.auth.default() + 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, timeout=_TELEMETRY_API_CHECK_TIMEOUT_SECONDS ) - session.configure_mtls_channel() - endpoint = _get_api_endpoint(client_cert_source) - else: - endpoint = _DEFAULT_TELEMETRY_ENDPOINT - r = session.post(endpoint, data=None) + except Exception as e: # pylint: disable=broad-exception-caught + _warn(f"Could not verify whether the Telemetry API is enabled: {e}") + return if "Telemetry API has not been used in project" in r.text: _warn(_TELEMETRY_API_DISABLED_WARNING % (project, project)) diff --git a/tests/unit/agentplatform/frameworks/test_frameworks_adk.py b/tests/unit/agentplatform/frameworks/test_frameworks_adk.py index d88cd56074..99e9881cc4 100644 --- a/tests/unit/agentplatform/frameworks/test_frameworks_adk.py +++ b/tests/unit/agentplatform/frameworks/test_frameworks_adk.py @@ -41,6 +41,7 @@ ) from google.genai import types import pytest +import requests from google.adk.sessions.base_session_service import BaseSessionService @@ -1733,6 +1734,32 @@ def test_update_default_telemetry_enablement( class TestAdkAppMtls: """Test cases for mTLS functionality in AdkApp.""" + @pytest.fixture(autouse=True) + def _isolate_global_tracer_provider(self): + """Isolates these tests from span processors left by earlier tests. + + `_default_instrumentor_builder` shuts down whatever span processor is + already installed on the global tracer provider. A processor left there + by another test in the same process holds a real `AuthorizedSession`, + and closing it raises `TypeError` while these tests have + `AuthorizedSession` patched out, so the outcome depends on which other + tests happen to share the shard. + """ + import opentelemetry.sdk.trace + import opentelemetry.trace + + tracer_provider = opentelemetry.trace.get_tracer_provider() + original = getattr(tracer_provider, "_active_span_processor", None) + if original is not None: + tracer_provider._active_span_processor = ( + opentelemetry.sdk.trace.SynchronousMultiSpanProcessor() + ) + try: + yield + finally: + if original is not None: + tracer_provider._active_span_processor = original + def setup_method(self): import opentelemetry.trace @@ -1865,7 +1892,9 @@ def test_warn_if_telemetry_api_disabled_with_mtls( mock_session.configure_mtls_channel.assert_called_once() # Verify the check was performed against the mTLS endpoint mock_session.post.assert_called_once_with( - adk_template._DEFAULT_MTLS_TELEMETRY_ENDPOINT, data=None + adk_template._DEFAULT_MTLS_TELEMETRY_ENDPOINT, + data=None, + timeout=adk_template._TELEMETRY_API_CHECK_TIMEOUT_SECONDS, ) @mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "invalid_value"}) @@ -1957,7 +1986,39 @@ def test_warn_if_telemetry_api_disabled_no_mtls( mock_session.configure_mtls_channel.assert_not_called() # Verify the check was performed against the regular endpoint mock_session.post.assert_called_once_with( - adk_template._DEFAULT_TELEMETRY_ENDPOINT, data=None + adk_template._DEFAULT_TELEMETRY_ENDPOINT, + data=None, + timeout=adk_template._TELEMETRY_API_CHECK_TIMEOUT_SECONDS, + ) + + @mock.patch("google.auth.default", return_value=(mock.Mock(), _TEST_PROJECT)) + @mock.patch.object(adk_template.requests_auth, "AuthorizedSession") + def test_warn_if_telemetry_api_disabled_survives_connection_error( + self, + mock_session_cls, + mock_auth_default, + ): + """The telemetry check must never propagate a failure to its caller. + + Regression test for b/546241881: this check runs from `set_up()` on the + agent server's startup path, where an unhandled `ConnectionError` took + down the container instead of degrading to a warning. + """ + mock_session = mock_session_cls.return_value + mock_session.post.side_effect = requests.exceptions.ConnectionError( + "('Connection aborted.', RemoteDisconnected('Remote end closed" + " connection without response'))" + ) + + with mock.patch.object( + adk_template, "_use_client_cert_effective", return_value=False + ): + with mock.patch.object(adk_template, "_warn") as mock_warn: + adk_template._warn_if_telemetry_api_disabled() + + mock_warn.assert_called_once() + assert "Could not verify whether the Telemetry API is enabled" in ( + mock_warn.call_args.args[0] ) @mock.patch("google.auth.default", return_value=(mock.Mock(), _TEST_PROJECT)) diff --git a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py index c80e96a5b5..a359289406 100644 --- a/tests/unit/vertex_adk/test_agent_engine_templates_adk.py +++ b/tests/unit/vertex_adk/test_agent_engine_templates_adk.py @@ -40,6 +40,7 @@ from vertexai.agent_engines.templates import adk as adk_template from google.genai import types import pytest +import requests from google.adk.sessions.base_session_service import BaseSessionService @@ -1409,6 +1410,32 @@ def test_update_default_telemetry_enablement( class TestAdkAppMtls: """Test cases for mTLS functionality in AdkApp.""" + @pytest.fixture(autouse=True) + def _isolate_global_tracer_provider(self): + """Isolates these tests from span processors left by earlier tests. + + `_default_instrumentor_builder` shuts down whatever span processor is + already installed on the global tracer provider. A processor left there + by another test in the same process holds a real `AuthorizedSession`, + and closing it raises `TypeError` while these tests have + `AuthorizedSession` patched out, so the outcome depends on which other + tests happen to share the shard. + """ + import opentelemetry.sdk.trace + import opentelemetry.trace + + tracer_provider = opentelemetry.trace.get_tracer_provider() + original = getattr(tracer_provider, "_active_span_processor", None) + if original is not None: + tracer_provider._active_span_processor = ( + opentelemetry.sdk.trace.SynchronousMultiSpanProcessor() + ) + try: + yield + finally: + if original is not None: + tracer_provider._active_span_processor = original + 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( @@ -1536,7 +1563,9 @@ def test_warn_if_telemetry_api_disabled_with_mtls( mock_session.configure_mtls_channel.assert_called_once() # Verify the check was performed against the mTLS endpoint mock_session.post.assert_called_once_with( - adk_template._DEFAULT_MTLS_TELEMETRY_ENDPOINT, data=None + adk_template._DEFAULT_MTLS_TELEMETRY_ENDPOINT, + data=None, + timeout=adk_template._TELEMETRY_API_CHECK_TIMEOUT_SECONDS, ) @mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "invalid_value"}) @@ -1628,7 +1657,39 @@ def test_warn_if_telemetry_api_disabled_no_mtls( mock_session.configure_mtls_channel.assert_not_called() # Verify the check was performed against the regular endpoint mock_session.post.assert_called_once_with( - adk_template._DEFAULT_TELEMETRY_ENDPOINT, data=None + adk_template._DEFAULT_TELEMETRY_ENDPOINT, + data=None, + timeout=adk_template._TELEMETRY_API_CHECK_TIMEOUT_SECONDS, + ) + + @mock.patch("google.auth.default", return_value=(mock.Mock(), _TEST_PROJECT)) + @mock.patch.object(adk_template.requests_auth, "AuthorizedSession") + def test_warn_if_telemetry_api_disabled_survives_connection_error( + self, + mock_session_cls, + mock_auth_default, + ): + """The telemetry check must never propagate a failure to its caller. + + Regression test for b/546241881: this check runs from `set_up()` on the + agent server's startup path, where an unhandled `ConnectionError` took + down the container instead of degrading to a warning. + """ + mock_session = mock_session_cls.return_value + mock_session.post.side_effect = requests.exceptions.ConnectionError( + "('Connection aborted.', RemoteDisconnected('Remote end closed" + " connection without response'))" + ) + + with mock.patch.object( + adk_template, "_use_client_cert_effective", return_value=False + ): + with mock.patch.object(adk_template, "_warn") as mock_warn: + adk_template._warn_if_telemetry_api_disabled() + + mock_warn.assert_called_once() + assert "Could not verify whether the Telemetry API is enabled" in ( + mock_warn.call_args.args[0] ) @mock.patch("google.auth.default", return_value=(mock.Mock(), _TEST_PROJECT)) diff --git a/vertexai/agent_engines/templates/adk.py b/vertexai/agent_engines/templates/adk.py index 25af45c2dd..76398ba03c 100644 --- a/vertexai/agent_engines/templates/adk.py +++ b/vertexai/agent_engines/templates/adk.py @@ -112,6 +112,10 @@ _DEFAULT_TELEMETRY_ENDPOINT = "https://telemetry.googleapis.com/v1/traces" _DEFAULT_MTLS_TELEMETRY_ENDPOINT = "https://telemetry.mtls.googleapis.com/v1/traces" +# Timeout (in seconds) for the best-effort Telemetry API enablement check. This +# check runs on the agent server's critical startup path, so it must fail fast +# rather than inherit AuthorizedSession's 120s default. +_TELEMETRY_API_CHECK_TIMEOUT_SECONDS = 5.0 class _MtlsEndpoint(enum.Enum): @@ -585,22 +589,34 @@ def _validate_run_config(run_config: Optional[Dict[str, Any]]): def _warn_if_telemetry_api_disabled(): - """Warn if telemetry API is disabled.""" - credentials, project = google.auth.default() - 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 + """Warns if the Telemetry API is disabled. + + This is a best-effort diagnostic that runs from `set_up()`, on the agent + server's critical startup path. No failure here (blocked egress, an + unreachable endpoint, unavailable credentials) may prevent the container + from starting, so the whole body is non-fatal. + """ + try: + credentials, project = google.auth.default() + 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, timeout=_TELEMETRY_API_CHECK_TIMEOUT_SECONDS ) - session.configure_mtls_channel() - endpoint = _get_api_endpoint(client_cert_source) - else: - endpoint = _DEFAULT_TELEMETRY_ENDPOINT - r = session.post(endpoint, data=None) + except Exception as e: # pylint: disable=broad-exception-caught + _warn(f"Could not verify whether the Telemetry API is enabled: {e}") + return if "Telemetry API has not been used in project" in r.text: _warn(_TELEMETRY_API_DISABLED_WARNING % (project, project)) diff --git a/vertexai/preview/reasoning_engines/templates/adk.py b/vertexai/preview/reasoning_engines/templates/adk.py index 4138a69ecf..ed8909d804 100644 --- a/vertexai/preview/reasoning_engines/templates/adk.py +++ b/vertexai/preview/reasoning_engines/templates/adk.py @@ -103,6 +103,10 @@ _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" +# Timeout (in seconds) for the best-effort Telemetry API enablement check. This +# check runs on the agent server's critical startup path, so it must fail fast +# rather than inherit AuthorizedSession's 120s default. +_TELEMETRY_API_CHECK_TIMEOUT_SECONDS = 5.0 class _MtlsEndpoint(enum.Enum): @@ -1724,27 +1728,39 @@ def _tracing_enabled(self) -> bool: ) def _warn_if_telemetry_api_disabled(self): - """Warn if telemetry API is disabled.""" + """Warns if the Telemetry API is disabled. + + This is a best-effort diagnostic that runs from `set_up()`, on the + agent server's critical startup path. No failure here (blocked egress, + an unreachable endpoint, unavailable credentials) may prevent the + container from starting, so the whole body is non-fatal. + """ try: import google.auth.transport.requests import google.auth except (ImportError, AttributeError): return - credentials, project = google.auth.default() - 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 + try: + credentials, project = google.auth.default() + 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, timeout=_TELEMETRY_API_CHECK_TIMEOUT_SECONDS ) - session.configure_mtls_channel() - endpoint = _get_api_endpoint(client_cert_source) - else: - endpoint = _DEFAULT_TELEMETRY_ENDPOINT - r = session.post(endpoint, data=None) + except Exception as e: # pylint: disable=broad-exception-caught + _warn(f"Could not verify whether the Telemetry API is enabled: {e}") + return if "Telemetry API has not been used in project" in r.text: _warn(_TELEMETRY_API_DISABLED_WARNING % (project, project))