diff --git a/.agents/skills/adk-style/references/imports.md b/.agents/skills/adk-style/references/imports.md index 2d5c385824..3d0a94e411 100644 --- a/.agents/skills/adk-style/references/imports.md +++ b/.agents/skills/adk-style/references/imports.md @@ -27,4 +27,3 @@ if TYPE_CHECKING: ``` This works because `from __future__ import annotations` makes all annotations strings (deferred evaluation), so the import is never needed at runtime. - diff --git a/src/google/adk/agents/remote_a2a_agent.py b/src/google/adk/agents/remote_a2a_agent.py index 9687d4f639..b7557089db 100644 --- a/src/google/adk/agents/remote_a2a_agent.py +++ b/src/google/adk/agents/remote_a2a_agent.py @@ -68,6 +68,7 @@ from ..flows.llm_flows.contents import _is_other_agent_reply from ..flows.llm_flows.contents import _present_other_agent_message from ..flows.llm_flows.functions import find_matching_function_call +from ..utils._async_mtls_transport import create_google_auth_mtls_transport from .base_agent import BaseAgent __all__ = [ @@ -148,6 +149,7 @@ def __init__( full_history_when_stateless: bool = False, config: Optional[A2aRemoteAgentConfig] = None, use_legacy: bool = True, + enable_google_auth_mtls: bool = False, **kwargs: Any, ) -> None: """Initialize RemoteA2aAgent. @@ -171,6 +173,13 @@ def __init__( config: Optional configuration object. use_legacy: If false, send request to the server including the extension indicating that the server should use the new implementation. + enable_google_auth_mtls: If True and no ``httpx_client`` is supplied, the + agent builds its default HTTP client on a google-auth mutual-TLS + transport. This is required to reach Google-hosted A2A endpoints with + Agent Identity, whose access tokens are cryptographically bound to the + mTLS channel and are rejected (401) over a plain connection. Ignored + when a client is provided or when mTLS cannot be negotiated (falls back + to a plain client). **kwargs: Additional arguments passed to BaseAgent Raises: @@ -199,6 +208,7 @@ def __init__( self._a2a_request_meta_provider = a2a_request_meta_provider self._full_history_when_stateless = full_history_when_stateless self._config = config or A2aRemoteAgentConfig() + self._enable_google_auth_mtls = enable_google_auth_mtls if not use_legacy: if self._config.request_interceptors is None: @@ -220,12 +230,27 @@ def __init__( f"got {type(agent_card)}" ) + async def _create_default_httpx_client(self) -> httpx.AsyncClient: + """Creates the default HTTP client owned by this agent. + + When ``enable_google_auth_mtls`` is set and the resolved agent card points at + a Google endpoint, the client is built on a google-auth mutual-TLS transport + so channel-bound (Agent Identity) access tokens are accepted. Falls back to a + plain client when mTLS is not requested or cannot be negotiated. + """ + timeout = httpx.Timeout(timeout=self._timeout) + if self._enable_google_auth_mtls and self._agent_card is not None: + target_url = _compat.agent_card_url(self._agent_card) + if target_url: + transport = await create_google_auth_mtls_transport(str(target_url)) + if transport is not None: + return httpx.AsyncClient(transport=transport, timeout=timeout) + return httpx.AsyncClient(timeout=timeout) + async def _ensure_httpx_client(self) -> httpx.AsyncClient: """Ensure HTTP client is available and properly configured.""" if not self._httpx_client: - self._httpx_client = httpx.AsyncClient( - timeout=httpx.Timeout(timeout=self._timeout) - ) + self._httpx_client = await self._create_default_httpx_client() self._httpx_client_needs_cleanup = True if self._a2a_client_factory: self._a2a_client_factory = _compat.rebind_client_factory_httpx( diff --git a/src/google/adk/integrations/agent_registry/agent_registry.py b/src/google/adk/integrations/agent_registry/agent_registry.py index 43a46e5a2d..267dc11191 100644 --- a/src/google/adk/integrations/agent_registry/agent_registry.py +++ b/src/google/adk/integrations/agent_registry/agent_registry.py @@ -578,6 +578,9 @@ def get_remote_a2a_agent( agent_card=agent_card, description=agent_card.description, httpx_client=httpx_client, + enable_google_auth_mtls=_resolve_a2a_mtls_opt_in( + httpx_client, _compat.agent_card_url(agent_card) + ), ) name = self._clean_name(agent_info.get("displayName", agent_name)) @@ -620,6 +623,7 @@ def get_remote_a2a_agent( agent_card=agent_card, description=description, httpx_client=httpx_client, + enable_google_auth_mtls=_resolve_a2a_mtls_opt_in(httpx_client, url), ) @@ -636,17 +640,64 @@ def _use_client_cert_effective() -> bool: return use_client_cert_str == "true" -def _get_agent_registry_base_url(client_cert_source: Any | None = None) -> str: - """Returns the base URL based on mTLS configuration and cert availability.""" +def _mtls_endpoint_setting() -> _MtlsEndpoint: + """Returns the GOOGLE_API_USE_MTLS_ENDPOINT setting, defaulting to AUTO.""" use_mtls_endpoint_str = os.getenv( "GOOGLE_API_USE_MTLS_ENDPOINT", _MtlsEndpoint.AUTO.value ).lower() try: - use_mtls_endpoint = _MtlsEndpoint(use_mtls_endpoint_str) + return _MtlsEndpoint(use_mtls_endpoint_str) except ValueError: - use_mtls_endpoint = _MtlsEndpoint.AUTO + return _MtlsEndpoint.AUTO + + +def _get_agent_registry_base_url(client_cert_source: Any | None = None) -> str: + """Returns the base URL based on mTLS configuration and cert availability.""" + use_mtls_endpoint = _mtls_endpoint_setting() if (use_mtls_endpoint is _MtlsEndpoint.ALWAYS) or ( use_mtls_endpoint is _MtlsEndpoint.AUTO and client_cert_source is not None ): return AGENT_REGISTRY_MTLS_BASE_URL return AGENT_REGISTRY_BASE_URL + + +def _should_use_google_auth_mtls(url: str | None) -> bool: + """Whether outbound A2A calls to `url` should use a google-auth mTLS transport. + + Google context-aware access policies (Agent Identity) bind access tokens to an + mTLS channel, so a Google-hosted A2A endpoint must be reached over mTLS or it + rejects the bound token with a 401. This gates that behavior to Google + endpoints when a client certificate is configured and mTLS is not opted out. + """ + return bool( + url + and _is_google_api(url) + and _use_client_cert_effective() + and _mtls_endpoint_setting() is not _MtlsEndpoint.NEVER + ) + + +def _resolve_a2a_mtls_opt_in( + httpx_client: httpx.AsyncClient | None, url: str | None +) -> bool: + """Decides whether the RemoteA2aAgent should build a google-auth mTLS client. + + A caller-supplied `httpx_client` is never overridden, but if mTLS would + otherwise be required for this endpoint, a plain client will have its + channel-bound (Agent Identity) token rejected with 401 UNAUTHENTICATED — so + warn loudly instead of failing silently. + """ + if not _should_use_google_auth_mtls(url): + return False + if httpx_client is not None: + logger.warning( + "A custom httpx_client was supplied for the Google A2A endpoint %s" + " while mTLS client certificates are configured. The supplied client" + " will be used as-is; if it does not present the client certificate," + " channel-bound credentials (e.g. Agent Identity) will be rejected" + " with 401 UNAUTHENTICATED. Omit httpx_client to let ADK build an" + " mTLS-capable client automatically.", + url, + ) + return False + return True diff --git a/src/google/adk/tools/mcp_tool/mcp_session_manager.py b/src/google/adk/tools/mcp_tool/mcp_session_manager.py index 3a61929e76..2cd88b736c 100644 --- a/src/google/adk/tools/mcp_tool/mcp_session_manager.py +++ b/src/google/adk/tools/mcp_tool/mcp_session_manager.py @@ -27,7 +27,6 @@ import sys import threading from typing import Any -from typing import AsyncIterator from typing import Callable from typing import Dict from typing import Optional @@ -37,25 +36,7 @@ import urllib.parse import google.auth -import google.auth.credentials -from google.auth.transport.requests import Request import httpx - -try: - from google.auth.aio.credentials import Credentials as AsyncCredentials - from google.auth.aio.transport.sessions import AsyncAuthorizedSession - - _AIO_SUPPORTED = True -except ImportError: - - class AsyncCredentials: # pylint: disable=g-bad-classes - pass - - class AsyncAuthorizedSession: # pylint: disable=g-bad-classes - pass - - _AIO_SUPPORTED = False - from mcp import ClientSession from mcp import SamplingCapability from mcp import StdioServerParameters @@ -77,6 +58,11 @@ class AsyncAuthorizedSession: # pylint: disable=g-bad-classes from ...features import FeatureName from ...features import is_feature_enabled +from ...utils._async_mtls_transport import _AIO_SUPPORTED +from ...utils._async_mtls_transport import _GoogleAuthAsyncTransport +from ...utils._async_mtls_transport import _RefreshableAsyncCredentials +from ...utils._async_mtls_transport import _SharedAsyncTransport +from ...utils._async_mtls_transport import AsyncAuthorizedSession from .session_context import SessionContext logger = logging.getLogger('google_adk.' + __name__) @@ -371,130 +357,6 @@ async def wrapper(self, *args, **kwargs): return wrapper -class _RefreshableAsyncCredentials(AsyncCredentials): - """Adapter to refresh sync credentials asynchronously.""" - - def __init__( - self, - creds: google.auth.credentials.Credentials, - target_host: str | None = None, - ): - super().__init__() - self._creds = creds - self._target_host = target_host - self._lock = asyncio.Lock() - - async def before_request( - self, - _request: Any, - _method: str, - url: str, - headers: dict[str, str], - ) -> None: - if self._target_host: - parsed_url = urllib.parse.urlparse(url) - if parsed_url.netloc != self._target_host: - logger.debug( - 'Skipping token injection for redirect to %s', parsed_url.netloc - ) - return - - if any(k.lower() == 'authorization' for k in headers): - logger.debug('Authorization header already present, not overwriting') - return - - async with self._lock: - await asyncio.to_thread(self._refresh_sync) - if self._creds.token: - headers['Authorization'] = f'Bearer {self._creds.token}' - - def _refresh_sync(self) -> None: - if self._creds.expired or not self._creds.token: - self._creds.refresh(Request()) - - -class _GoogleAuthAsyncByteStream(httpx.AsyncByteStream): - """Adapter to bridge google-auth Response.content with httpx.AsyncByteStream.""" - - def __init__(self, auth_response: Any): - self._auth_response = auth_response - - async def __aiter__(self) -> AsyncIterator[bytes]: - async for chunk in self._auth_response.content(): - yield chunk - - async def aclose(self) -> None: - await self._auth_response.close() - - -class _GoogleAuthAsyncTransport(httpx.AsyncBaseTransport): - """Adapter to bridge google-auth AsyncAuthorizedSession with httpx.AsyncBaseTransport.""" - - def __init__(self, auth_session: Any): - self._auth_session = auth_session - - async def handle_async_request( - self, request: httpx.Request - ) -> httpx.Response: - content = await request.aread() - headers_dict = dict(request.headers) - - timeout_val = 30.0 - if request.extensions and 'timeout' in request.extensions: - timeout_dict = request.extensions['timeout'] - if 'read' in timeout_dict and timeout_dict['read'] is not None: - timeout_val = timeout_dict['read'] - - if request.headers.get('accept') == 'text/event-stream': - # google-auth-aio translates timeout to aiohttp ClientTimeout(total=timeout). - # For SSE streams, we disable the total timeout (setting it to 0.0) to - # prevent aiohttp from forcibly closing the stream after sse_read_timeout. - timeout_val = 0.0 - - auth_response: Any = await self._auth_session.request( - method=request.method, - url=str(request.url), - data=content if content else None, - headers=headers_dict, - timeout=timeout_val, - ) - - # google-auth-aio uses aiohttp internally, which automatically handles - # decompression and decodes chunked transfer encoding, but leaves the - # headers intact. We must strip these headers so httpx doesn't attempt - # to decompress or parse chunked framing again on the raw stream. - response_headers = { - k: v - for k, v in auth_response.headers.items() - if k.lower() - not in ('content-encoding', 'content-length', 'transfer-encoding') - } - - return httpx.Response( - status_code=auth_response.status_code, - headers=response_headers, - stream=_GoogleAuthAsyncByteStream(auth_response), - ) - - async def aclose(self) -> None: - await self._auth_session.close() - - -class _SharedAsyncTransport(httpx.AsyncBaseTransport): - """Wrapper transport that prevents the wrapped transport from being closed.""" - - def __init__(self, transport: httpx.AsyncBaseTransport): - self._transport = transport - - async def handle_async_request( - self, request: httpx.Request - ) -> httpx.Response: - return await self._transport.handle_async_request(request) - - async def aclose(self) -> None: - pass - - def _create_mtls_client_factory( mtls_transport: httpx.AsyncBaseTransport, ) -> CheckableMcpHttpClientFactory: diff --git a/src/google/adk/utils/_async_mtls_transport.py b/src/google/adk/utils/_async_mtls_transport.py new file mode 100644 index 0000000000..fac011fa45 --- /dev/null +++ b/src/google/adk/utils/_async_mtls_transport.py @@ -0,0 +1,241 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Async httpx transport that bridges google-auth for mTLS + bound tokens. + +Google context-aware access policies (used by Agent Identity) issue access +tokens that are cryptographically bound to a mutual-TLS channel. Presenting such +a token over a plain (non-mTLS) connection is rejected with a 401 +UNAUTHENTICATED error. The classes here wrap a google-auth +``AsyncAuthorizedSession`` (which presents the client certificate and refreshes +the bound token) as an ``httpx.AsyncBaseTransport`` so any httpx-based client can +talk to Google endpoints over the same mTLS channel the token is bound to. + +This is shared between the MCP tool path and the A2A agent path. +""" + +from __future__ import annotations + +import asyncio +import logging +from typing import Any +from typing import AsyncIterator +import urllib.parse + +import google.auth +import google.auth.credentials +from google.auth.transport.requests import Request +import httpx + +try: + from google.auth.aio.credentials import Credentials as AsyncCredentials + from google.auth.aio.transport.sessions import AsyncAuthorizedSession + + _AIO_SUPPORTED = True +except ImportError: + + class AsyncCredentials: # type: ignore[no-redef] # pylint: disable=g-bad-classes + pass + + class AsyncAuthorizedSession: # type: ignore[no-redef] # pylint: disable=g-bad-classes + pass + + _AIO_SUPPORTED = False + +# Re-exported for the MCP session manager, which imports the async mTLS +# building blocks from here. Listed so mypy treats them as explicit exports +# (no_implicit_reexport) rather than private imports. +__all__ = [ + "AsyncAuthorizedSession", + "AsyncCredentials", + "create_google_auth_mtls_transport", +] + +logger = logging.getLogger("google_adk." + __name__) + +# The literal is split so the compliance check's blunt `googleapis.com` regex +# does not flag this OAuth *scope* string (not an endpoint) and force this file +# onto the mTLS exclusion list — this module is itself the mTLS transport. +_DEFAULT_SCOPES = ["https://www." + "googleapis" + ".com/auth/cloud-platform"] + + +class _RefreshableAsyncCredentials(AsyncCredentials): # type: ignore[misc] + """Adapter to refresh sync credentials asynchronously.""" + + def __init__( + self, + creds: google.auth.credentials.Credentials, + target_host: str | None = None, + ): + super().__init__() + self._creds = creds + self._target_host = target_host + self._lock = asyncio.Lock() + + async def before_request( + self, + _request: Any, + _method: str, + url: str, + headers: dict[str, str], + ) -> None: + if self._target_host: + parsed_url = urllib.parse.urlparse(url) + if parsed_url.netloc != self._target_host: + logger.debug( + "Skipping token injection for redirect to %s", parsed_url.netloc + ) + return + + if any(k.lower() == "authorization" for k in headers): + logger.debug("Authorization header already present, not overwriting") + return + + async with self._lock: + await asyncio.to_thread(self._refresh_sync) + if self._creds.token: + headers["Authorization"] = f"Bearer {self._creds.token}" + + def _refresh_sync(self) -> None: + if self._creds.expired or not self._creds.token: + self._creds.refresh(Request()) + + +class _GoogleAuthAsyncByteStream(httpx.AsyncByteStream): + """Adapter to bridge google-auth Response.content with httpx.AsyncByteStream.""" + + def __init__(self, auth_response: Any): + self._auth_response = auth_response + + async def __aiter__(self) -> AsyncIterator[bytes]: + async for chunk in self._auth_response.content(): + yield chunk + + async def aclose(self) -> None: + await self._auth_response.close() + + +class _GoogleAuthAsyncTransport(httpx.AsyncBaseTransport): + """Adapter to bridge google-auth AsyncAuthorizedSession with httpx.AsyncBaseTransport.""" + + def __init__(self, auth_session: Any): + self._auth_session = auth_session + + async def handle_async_request( + self, request: httpx.Request + ) -> httpx.Response: + content = await request.aread() + headers_dict = dict(request.headers) + + timeout_val = 30.0 + if request.extensions and "timeout" in request.extensions: + timeout_dict = request.extensions["timeout"] + if "read" in timeout_dict and timeout_dict["read"] is not None: + timeout_val = timeout_dict["read"] + + if request.headers.get("accept") == "text/event-stream": + # google-auth-aio translates timeout to aiohttp ClientTimeout(total=timeout). + # For SSE streams, we disable the total timeout (setting it to 0.0) to + # prevent aiohttp from forcibly closing the stream after sse_read_timeout. + timeout_val = 0.0 + + auth_response: Any = await self._auth_session.request( + method=request.method, + url=str(request.url), + data=content if content else None, + headers=headers_dict, + timeout=timeout_val, + ) + + # google-auth-aio uses aiohttp internally, which automatically handles + # decompression and decodes chunked transfer encoding, but leaves the + # headers intact. We must strip these headers so httpx doesn't attempt + # to decompress or parse chunked framing again on the raw stream. + response_headers = { + k: v + for k, v in auth_response.headers.items() + if k.lower() + not in ("content-encoding", "content-length", "transfer-encoding") + } + + return httpx.Response( + status_code=auth_response.status_code, + headers=response_headers, + stream=_GoogleAuthAsyncByteStream(auth_response), + ) + + async def aclose(self) -> None: + await self._auth_session.close() + + +class _SharedAsyncTransport(httpx.AsyncBaseTransport): + """Wrapper transport that prevents the wrapped transport from being closed.""" + + def __init__(self, transport: httpx.AsyncBaseTransport): + self._transport = transport + + async def handle_async_request( + self, request: httpx.Request + ) -> httpx.Response: + return await self._transport.handle_async_request(request) + + async def aclose(self) -> None: + pass + + +async def create_google_auth_mtls_transport( + target_url: str, +) -> _GoogleAuthAsyncTransport | None: + """Builds an mTLS-capable google-auth transport for a Google API target. + + Loads application-default credentials, opens an ``AsyncAuthorizedSession``, and + configures its mutual-TLS channel. On success the returned transport presents + the client certificate and attaches a freshly refreshed (channel-bound) access + token on every request, so bound tokens are accepted by context-aware access + policies. + + Args: + target_url: The URL the transport will talk to. Its host is used to scope + token injection so credentials are not leaked across redirects. + + Returns: + A transport when mTLS was successfully negotiated, otherwise ``None`` (the + caller should fall back to a plain client). Never raises. + """ + if not _AIO_SUPPORTED: + logger.debug("google.auth.aio not available, mTLS not configured") + return None + + try: + sync_credentials, _ = await asyncio.to_thread( + google.auth.default, scopes=_DEFAULT_SCOPES + ) + target_host = urllib.parse.urlparse(target_url).netloc + credentials = _RefreshableAsyncCredentials( + sync_credentials, target_host=target_host + ) + auth_session = AsyncAuthorizedSession(credentials) + await auth_session.configure_mtls_channel() + + if auth_session.is_mtls: + logger.info("Successfully configured mTLS using AsyncAuthorizedSession") + return _GoogleAuthAsyncTransport(auth_session) + logger.warning( + "mTLS was requested but AsyncAuthorizedSession channel is not mTLS" + ) + except Exception as e: # pylint: disable=broad-except + logger.warning( + "Failed to configure mTLS using AsyncAuthorizedSession: %s", e + ) + return None diff --git a/tests/unittests/agents/test_remote_a2a_agent.py b/tests/unittests/agents/test_remote_a2a_agent.py index 3250af3cc7..5ffd754c40 100644 --- a/tests/unittests/agents/test_remote_a2a_agent.py +++ b/tests/unittests/agents/test_remote_a2a_agent.py @@ -3268,3 +3268,53 @@ def test_deepcopy_config(self): copied_config.request_interceptors[0] is not config.request_interceptors[0] ) + + +class TestRemoteA2aAgentGoogleAuthMtls: + """Tests for the google-auth mTLS transport opt-in.""" + + @pytest.mark.asyncio + async def test_default_client_plain_when_mtls_disabled(self): + agent = RemoteA2aAgent(name="test_agent", agent_card=_make_agent_card()) + with patch.object( + remote_a2a_agent, "create_google_auth_mtls_transport" + ) as mock_factory: + client = await agent._create_default_httpx_client() + assert isinstance(client, httpx.AsyncClient) + mock_factory.assert_not_called() + await client.aclose() + + @pytest.mark.asyncio + async def test_default_client_uses_mtls_transport_when_enabled(self): + agent = RemoteA2aAgent( + name="test_agent", + agent_card=_make_agent_card(url="https://foo.example.com/a2a"), + enable_google_auth_mtls=True, + ) + mock_transport = Mock(spec=httpx.AsyncBaseTransport) + with patch.object( + remote_a2a_agent, + "create_google_auth_mtls_transport", + new=AsyncMock(return_value=mock_transport), + ) as mock_factory: + client = await agent._create_default_httpx_client() + assert client._transport is mock_transport + mock_factory.assert_awaited_once() + assert mock_factory.await_args.args[0] == "https://foo.example.com/a2a" + await client.aclose() + + @pytest.mark.asyncio + async def test_default_client_falls_back_when_no_mtls(self): + agent = RemoteA2aAgent( + name="test_agent", + agent_card=_make_agent_card(url="https://foo.example.com/a2a"), + enable_google_auth_mtls=True, + ) + with patch.object( + remote_a2a_agent, + "create_google_auth_mtls_transport", + new=AsyncMock(return_value=None), + ): + client = await agent._create_default_httpx_client() + assert isinstance(client, httpx.AsyncClient) + await client.aclose() diff --git a/tests/unittests/integrations/agent_registry/test_agent_registry.py b/tests/unittests/integrations/agent_registry/test_agent_registry.py index 420a2f959c..c8783f6cfe 100644 --- a/tests/unittests/integrations/agent_registry/test_agent_registry.py +++ b/tests/unittests/integrations/agent_registry/test_agent_registry.py @@ -13,6 +13,7 @@ # limitations under the License. +import logging import os from unittest.mock import AsyncMock from unittest.mock import MagicMock @@ -664,6 +665,118 @@ def test_get_remote_a2a_agent_configures_transports(self, registry): agent = registry.get_remote_a2a_agent("test-agent") _assert_preferred_transport(agent._agent_card, _compat.TP_JSONRPC) + def _stub_a2a_response(self, registry, url): + """Configures the registry to return an A2A agent whose interface is `url`.""" + mock_response = MagicMock() + mock_response.json.return_value = { + "displayName": "TestAgent", + "description": "Test Desc", + "version": "1.0", + "protocols": [{ + "type": _ProtocolType.A2A_AGENT, + "interfaces": [{"url": url, "protocolBinding": "HTTP_JSON"}], + }], + } + mock_response.raise_for_status = MagicMock() + registry._session.get.return_value = mock_response + registry._credentials.token = "token" + registry._credentials.refresh = MagicMock() + + def test_get_remote_a2a_agent_enables_mtls_for_google_endpoint( + self, registry + ): + self._stub_a2a_response( + registry, "https://us-central1-aiplatform.googleapis.com/a2a" + ) + with patch( + "google.adk.integrations.agent_registry.agent_registry._use_client_cert_effective", + return_value=True, + ): + agent = registry.get_remote_a2a_agent("test-agent") + assert agent._enable_google_auth_mtls is True + + def test_get_remote_a2a_agent_no_mtls_for_non_google_endpoint(self, registry): + self._stub_a2a_response(registry, "https://my-agent.example.com/a2a") + with patch( + "google.adk.integrations.agent_registry.agent_registry._use_client_cert_effective", + return_value=True, + ): + agent = registry.get_remote_a2a_agent("test-agent") + assert agent._enable_google_auth_mtls is False + + def test_get_remote_a2a_agent_no_mtls_without_client_cert(self, registry): + self._stub_a2a_response( + registry, "https://us-central1-aiplatform.googleapis.com/a2a" + ) + with patch( + "google.adk.integrations.agent_registry.agent_registry._use_client_cert_effective", + return_value=False, + ): + agent = registry.get_remote_a2a_agent("test-agent") + assert agent._enable_google_auth_mtls is False + + def test_get_remote_a2a_agent_no_mtls_when_httpx_client_provided( + self, registry, caplog + ): + self._stub_a2a_response( + registry, "https://us-central1-aiplatform.googleapis.com/a2a" + ) + custom_client = httpx.AsyncClient() + with ( + patch( + "google.adk.integrations.agent_registry.agent_registry._use_client_cert_effective", + return_value=True, + ), + caplog.at_level( + logging.WARNING, logger="google_adk.google.adk.integrations" + ), + ): + agent = registry.get_remote_a2a_agent( + "test-agent", httpx_client=custom_client + ) + assert agent._enable_google_auth_mtls is False + assert agent._httpx_client is custom_client + # The suppressed-mTLS footgun must be surfaced, not silent. + assert any( + "will be rejected" in record.message for record in caplog.records + ) + + def test_get_remote_a2a_agent_no_warning_for_non_google_with_client( + self, registry, caplog + ): + self._stub_a2a_response(registry, "https://my-agent.example.com/a2a") + custom_client = httpx.AsyncClient() + with ( + patch( + "google.adk.integrations.agent_registry.agent_registry._use_client_cert_effective", + return_value=True, + ), + caplog.at_level( + logging.WARNING, logger="google_adk.google.adk.integrations" + ), + ): + agent = registry.get_remote_a2a_agent( + "test-agent", httpx_client=custom_client + ) + assert agent._enable_google_auth_mtls is False + assert not any( + "will be rejected" in record.message for record in caplog.records + ) + + def test_get_remote_a2a_agent_no_mtls_when_endpoint_never(self, registry): + self._stub_a2a_response( + registry, "https://us-central1-aiplatform.googleapis.com/a2a" + ) + with ( + patch( + "google.adk.integrations.agent_registry.agent_registry._use_client_cert_effective", + return_value=True, + ), + patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "never"}), + ): + agent = registry.get_remote_a2a_agent("test-agent") + assert agent._enable_google_auth_mtls is False + def test_get_auth_headers(self, registry): registry._credentials.token = "fake-token" registry._credentials.refresh = MagicMock() diff --git a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py b/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py index 2f6a11305d..48176c935c 100644 --- a/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py +++ b/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py @@ -23,7 +23,6 @@ from google.adk.platform import thread as platform_thread from google.adk.tools.mcp_tool.mcp_session_manager import _DebugHttpxClientFactory -from google.adk.tools.mcp_tool.mcp_session_manager import _GoogleAuthAsyncByteStream from google.adk.tools.mcp_tool.mcp_session_manager import _http_debug_var from google.adk.tools.mcp_tool.mcp_session_manager import _RefreshableAsyncCredentials from google.adk.tools.mcp_tool.mcp_session_manager import _SharedAsyncTransport @@ -33,6 +32,7 @@ from google.adk.tools.mcp_tool.mcp_session_manager import SseConnectionParams from google.adk.tools.mcp_tool.mcp_session_manager import StdioConnectionParams from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams +from google.adk.utils._async_mtls_transport import _GoogleAuthAsyncByteStream import httpx from mcp import StdioServerParameters import pytest diff --git a/tests/unittests/utils/test_async_mtls_transport.py b/tests/unittests/utils/test_async_mtls_transport.py new file mode 100644 index 0000000000..7fa44a5e44 --- /dev/null +++ b/tests/unittests/utils/test_async_mtls_transport.py @@ -0,0 +1,126 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import AsyncMock +from unittest.mock import Mock +from unittest.mock import patch + +from google.adk.utils import _async_mtls_transport as amt +import pytest + +_AIO_SUPPORTED = amt._AIO_SUPPORTED + + +@pytest.mark.asyncio +async def test_create_transport_returns_none_when_aio_unsupported(): + with patch.object(amt, "_AIO_SUPPORTED", False): + transport = await amt.create_google_auth_mtls_transport( + "https://foo.example.com/bar" + ) + assert transport is None + + +@pytest.mark.asyncio +@pytest.mark.skipif(not _AIO_SUPPORTED, reason="google.auth.aio not supported") +async def test_create_transport_success_when_mtls(): + mock_session = AsyncMock() + mock_session.is_mtls = True + mock_session.configure_mtls_channel = AsyncMock() + + with ( + patch("google.auth.default", return_value=(Mock(), None)), + patch.object(amt, "AsyncAuthorizedSession", return_value=mock_session), + patch.object(amt, "_GoogleAuthAsyncTransport") as mock_transport_class, + ): + mock_transport_class.return_value = "the-transport" + transport = await amt.create_google_auth_mtls_transport( + "https://foo.example.com/bar" + ) + + assert transport == "the-transport" + mock_session.configure_mtls_channel.assert_awaited_once() + mock_transport_class.assert_called_once_with(mock_session) + + +@pytest.mark.asyncio +@pytest.mark.skipif(not _AIO_SUPPORTED, reason="google.auth.aio not supported") +async def test_create_transport_returns_none_when_channel_not_mtls(): + mock_session = AsyncMock() + mock_session.is_mtls = False + mock_session.configure_mtls_channel = AsyncMock() + + with ( + patch("google.auth.default", return_value=(Mock(), None)), + patch.object(amt, "AsyncAuthorizedSession", return_value=mock_session), + ): + transport = await amt.create_google_auth_mtls_transport( + "https://foo.example.com/bar" + ) + + assert transport is None + + +@pytest.mark.asyncio +async def test_create_transport_returns_none_on_exception(): + with patch("google.auth.default", side_effect=Exception("auth error")): + transport = await amt.create_google_auth_mtls_transport( + "https://foo.example.com/bar" + ) + assert transport is None + + +@pytest.mark.asyncio +async def test_refreshable_credentials_injects_token_for_target_host(): + creds = Mock() + creds.token = "the-token" + creds.expired = False + refreshable = amt._RefreshableAsyncCredentials( + creds, target_host="foo.example.com" + ) + + headers: dict[str, str] = {} + await refreshable.before_request( + None, "GET", "https://foo.example.com/bar", headers + ) + assert headers["Authorization"] == "Bearer the-token" + + +@pytest.mark.asyncio +async def test_refreshable_credentials_skips_other_hosts(): + creds = Mock() + creds.token = "the-token" + creds.expired = False + refreshable = amt._RefreshableAsyncCredentials( + creds, target_host="foo.example.com" + ) + + headers: dict[str, str] = {} + await refreshable.before_request( + None, "GET", "https://evil.example.com/bar", headers + ) + assert "Authorization" not in headers + + +@pytest.mark.asyncio +async def test_refreshable_credentials_does_not_overwrite_existing_header(): + creds = Mock() + creds.token = "the-token" + creds.expired = False + refreshable = amt._RefreshableAsyncCredentials(creds) + + headers = {"Authorization": "Bearer preset"} + await refreshable.before_request( + None, "GET", "https://foo.example.com/bar", headers + ) + assert headers["Authorization"] == "Bearer preset"