Skip to content
19 changes: 19 additions & 0 deletions src/google/adk/a2a/_compat.py
Original file line number Diff line number Diff line change
Expand Up @@ -562,6 +562,25 @@ def agent_card_url(
return getattr(card, "url", None)


def rewrite_agent_card_url(
card: AgentCard, rewriter: Callable[[str], str]
) -> None:
"""Rewrites the RPC URL(s) on an AgentCard in place using ``rewriter``.

Applies ``rewriter`` to each RPC URL the card exposes, mirroring
``agent_card_url``. 1.x rewrites every ``supported_interfaces[i].url``; 0.3.x
rewrites the top-level ``url``.
"""
if IS_A2A_V1:
for iface in card.supported_interfaces or []:
if iface.url:
iface.url = rewriter(iface.url)
else:
existing = getattr(card, "url", None)
if existing:
card.url = rewriter(existing)


# -----------------------------------------------------------------------------
# Stream-item normalization
# -----------------------------------------------------------------------------
Expand Down
42 changes: 39 additions & 3 deletions src/google/adk/agents/remote_a2a_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,8 @@
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 ..utils._mtls_utils import effective_googleapis_endpoint
from .base_agent import BaseAgent

__all__ = [
Expand Down Expand Up @@ -148,6 +150,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.
Expand All @@ -171,6 +174,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:
Expand Down Expand Up @@ -199,6 +209,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:
Expand All @@ -220,12 +231,37 @@ 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:
# A channel-bound token is only honored on the mTLS endpoint
# (*.mtls.googleapis.com); presenting it to the standard endpoint still
# 401s. Resolve to the mTLS endpoint and bind both the transport and the
# A2A client to it. effective_googleapis_endpoint is a no-op for
# non-Google hosts and when GOOGLE_API_USE_MTLS_ENDPOINT=never.
mtls_url = effective_googleapis_endpoint(str(target_url))
transport = await create_google_auth_mtls_transport(mtls_url)
if transport is not None:
if mtls_url != str(target_url):
_compat.rewrite_agent_card_url(
self._agent_card, effective_googleapis_endpoint
)
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(
Expand Down
59 changes: 55 additions & 4 deletions src/google/adk/integrations/agent_registry/agent_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -579,6 +579,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))
Expand Down Expand Up @@ -621,6 +624,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),
)


Expand All @@ -637,17 +641,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
148 changes: 5 additions & 143 deletions src/google/adk/tools/mcp_tool/mcp_session_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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__)
Expand Down Expand Up @@ -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:
Expand Down
Loading