Skip to content

fix: Make the Telemetry API enablement check non-fatal in AdkApp.set_up(). #7073

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: Make the Telemetry API enablement check non-fatal in AdkApp.set_up(). #7073

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 31 additions & 15 deletions agentplatform/agent_engines/templates/adk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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))

Expand Down
65 changes: 63 additions & 2 deletions tests/unit/agentplatform/frameworks/test_frameworks_adk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@
)
from google.genai import types
import pytest
import requests
from google.adk.sessions.base_session_service import BaseSessionService


Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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"})
Expand Down Expand Up @@ -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))
Expand Down
65 changes: 63 additions & 2 deletions tests/unit/vertex_adk/test_agent_engine_templates_adk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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"})
Expand Down Expand Up @@ -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))
Expand Down
46 changes: 31 additions & 15 deletions vertexai/agent_engines/templates/adk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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))

Expand Down
46 changes: 31 additions & 15 deletions vertexai/preview/reasoning_engines/templates/adk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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))

Expand Down
Loading