feat(agents): support remote MCP servers for ManagedAgent with runtime header callbacks

Add `RemoteMcpServer`, a server-side remote MCP tool for `ManagedAgent`. The
Managed Agents / Interactions API runs the MCP server itself, so ADK only
forwards the server URL and headers as an `MCPServerParam` and never opens an
MCP session. A `header_provider` callback (the same contract as the `LlmAgent`
`McpToolset.header_provider`) mints auth headers at request time, driven by the
runner, and is merged over any static headers so a fresh token can be generated
per turn.

Only remote (HTTP/streamable) MCP servers are supported; raw
`types.Tool.mcp_servers` remains rejected. Includes a live integration test
against Maps Grounding Lite, scoped to the Gemini Developer API backend; the
Vertex Interactions endpoint does not yet accept the `mcp_server` tool param
(consistent with google-genai documenting `types.Tool.mcp_servers` as
unsupported on Vertex AI).

Co-authored-by: Haran Rajkumar <haranrk@google.com>
PiperOrigin-RevId: 945306712
This commit is contained in:
Haran Rajkumar
2026-07-09 14:06:49 -07:00
committed by Copybara-Service
parent 283e92efc2
commit 2e2ec09a76
8 changed files with 445 additions and 16 deletions
+1
View File
@@ -63,6 +63,7 @@ _EXCLUDED_FROM_MTLS = {
'src/google/adk/tools/openapi_tool/auth/credential_exchangers/service_account_exchanger.py',
'src/google/adk/tools/pubsub/pubsub_credentials.py',
'src/google/adk/tools/spanner/spanner_credentials.py',
'tests/integration/test_managed_agent.py',
'tests/unittests/auth/test_credential_manager.py',
'tests/unittests/cli/utils/test_gcp_utils.py',
'tests/unittests/flows/llm_flows/test_functions_request_euc.py',
+35 -11
View File
@@ -14,6 +14,7 @@
from __future__ import annotations
import inspect
import logging
from typing import Any
from typing import AsyncGenerator
@@ -32,6 +33,7 @@ from pydantic import PrivateAttr
from ..events.event import Event
from ..flows.llm_flows.interactions_processor import _find_previous_interaction_state
from ..models.interactions_utils import _build_mcp_server_param
from ..models.interactions_utils import _convert_content_to_step
from ..models.interactions_utils import _create_interactions
from ..models.interactions_utils import build_interactions_request_log
@@ -39,12 +41,14 @@ from ..models.interactions_utils import convert_tools_config_to_interactions_for
from ..models.llm_request import LlmRequest
from ..models.llm_response import LlmResponse
from ..telemetry import tracer
from ..tools._remote_mcp_server import RemoteMcpServer
from ..tools.base_tool import BaseTool
from ..tools.tool_context import ToolContext
from ..utils.context_utils import Aclosing
from ..utils.env_utils import is_enterprise_mode_enabled
from .base_agent import BaseAgent
from .invocation_context import InvocationContext
from .readonly_context import ReadonlyContext
from .run_config import StreamingMode
if TYPE_CHECKING:
@@ -109,10 +113,12 @@ class ManagedAgent(BaseAgent):
"""An agent backed by the Managed Agents API (interactions.create).
This agent calls the Managed Agents API directly from its execution loop.
In this version only server-side tools are supported: ADK built-in tools and
raw ``google.genai.types.Tool`` configs (the kinds the interactions converter
understands). Client-executed tools (FunctionTool/callables) and MCP are not
yet supported.
Only server-side tools are supported: ADK built-in tools, raw
``google.genai.types.Tool`` configs (the kinds the interactions converter
understands), and server-side remote MCP servers declared as
``RemoteMcpServer`` specs (forwarded to the backend as an ``MCPServerParam``).
Client-executed tools (FunctionTool/callables) and raw
``types.Tool.mcp_servers`` configs are not supported and are rejected.
ManagedAgent supports streaming interactions only. Interactions are always
created with ``background=True`` (required by the Managed Agents workflow) and
@@ -132,10 +138,11 @@ class ManagedAgent(BaseAgent):
agent_config: Optional[CreateAgentInteractionAgentConfigParam] = None
"""Runtime configuration passed to interactions.create."""
tools: list[Union[types.Tool, BaseTool, Callable[..., Any]]] = Field(
default_factory=list
)
"""Server-side tools: ADK built-in tools or raw types.Tool configs."""
tools: list[
Union[types.Tool, BaseTool, Callable[..., Any], RemoteMcpServer]
] = Field(default_factory=list)
"""Server-side tools: ADK built-in tools, raw types.Tool configs, or
RemoteMcpServer specs for server-side remote MCP."""
_api_client: Optional[Client] = PrivateAttr(default=None)
@@ -176,8 +183,10 @@ class ManagedAgent(BaseAgent):
"""Resolve self.tools into interaction ToolParams (server-side only).
Raw types.Tool configs are passed through; ADK built-in tools are processed
into native tool configs. Client-executed tools (FunctionTool/callables) and
MCP tools are rejected.
into native tool configs. ``RemoteMcpServer`` specs are resolved to an
``MCPServerParam`` (headers minted at request time via ``header_provider``).
Client-executed tools (FunctionTool/callables) and raw
``types.Tool.mcp_servers`` configs are rejected.
"""
# Built-in tools are resolved in "managed agent" mode: the request carries
# the internal _is_managed_agent flag (and no model), so tools that normally
@@ -186,8 +195,20 @@ class ManagedAgent(BaseAgent):
llm_request = LlmRequest(config=types.GenerateContentConfig())
llm_request._is_managed_agent = True
tool_context = ToolContext(ctx)
mcp_params: list[ToolParam] = []
for tool in self.tools:
if isinstance(tool, RemoteMcpServer):
resolved_headers = dict(tool.headers or {})
if tool.header_provider is not None:
dynamic = tool.header_provider(ReadonlyContext(ctx))
if inspect.isawaitable(dynamic):
dynamic = await dynamic
if dynamic:
resolved_headers.update(dynamic) # dynamic wins on key conflict
mcp_params.append(_build_mcp_server_param(tool, resolved_headers))
continue
if isinstance(tool, types.Tool):
if tool.mcp_servers:
raise NotImplementedError(
@@ -233,7 +254,10 @@ class ManagedAgent(BaseAgent):
f'{tool.name}'
)
return convert_tools_config_to_interactions_format(llm_request.config)
return (
convert_tools_config_to_interactions_format(llm_request.config)
+ mcp_params
)
def _response_to_event(
self, ctx: InvocationContext, llm_response: LlmResponse
+29 -5
View File
@@ -60,6 +60,7 @@ from google.genai.interactions import InteractionCompletedEvent
from google.genai.interactions import InteractionCreatedEvent
from google.genai.interactions import InteractionSSEEvent
from google.genai.interactions import InteractionStatusUpdate
from google.genai.interactions import MCPServerParam
from google.genai.interactions import ModelOutputStep
from google.genai.interactions import ModelOutputStepParam
from google.genai.interactions import Step
@@ -81,6 +82,8 @@ from typing_extensions import deprecated
if TYPE_CHECKING:
from google.genai import Client
from ..tools._remote_mcp_server import RemoteMcpServer
from .llm_request import LlmRequest
from .llm_response import LlmResponse
@@ -251,12 +254,12 @@ def convert_part_to_interaction_content(part: types.Part) -> dict | None:
elif part.thought:
# part.thought is a boolean indicating this is a thought part
# ThoughtContentParam expects 'signature' (base64 encoded bytes)
result: dict[str, Any] = {'type': 'thought'}
thought_result: dict[str, Any] = {'type': 'thought'}
if part.thought_signature is not None:
result['signature'] = base64.b64encode(part.thought_signature).decode(
'utf-8'
)
return result
thought_result['signature'] = base64.b64encode(
part.thought_signature
).decode('utf-8')
return thought_result
elif part.code_execution_result is not None:
is_error = part.code_execution_result.outcome in (
types.Outcome.OUTCOME_FAILED,
@@ -521,6 +524,27 @@ def convert_tools_config_to_interactions_format(
return interaction_tools
def _build_mcp_server_param(
server: RemoteMcpServer,
resolved_headers: dict[str, str],
) -> MCPServerParam:
"""Map a RemoteMcpServer + resolved headers to an interactions MCPServerParam.
Built directly (not via ``types.McpServer``) so ``allowed_tools`` can be
carried and the "not supported in Vertex AI" restriction on
``types.Tool.mcp_servers`` is avoided. ``resolved_headers`` is the static
headers already merged with any ``header_provider`` output by the caller.
"""
param: MCPServerParam = {'type': 'mcp_server', 'url': server.url}
if server.name is not None:
param['name'] = server.name
if resolved_headers:
param['headers'] = resolved_headers
if server.allowed_tools is not None:
param['allowed_tools'] = [{'tools': list(server.allowed_tools)}]
return param
def _function_result_to_response(
result: BaseModel | dict[str, Any] | list[Any] | str,
) -> dict[str, Any]:
+2
View File
@@ -20,6 +20,7 @@ from typing import TYPE_CHECKING
# The TYPE_CHECKING block is needed for autocomplete to work.
if TYPE_CHECKING:
from ..auth.auth_tool import AuthToolArguments
from ._remote_mcp_server import RemoteMcpServer
from ._request_input_tool import request_input
from .agent_tool import AgentTool
from .api_registry import ApiRegistry
@@ -83,6 +84,7 @@ _LAZY_MAPPING = {
),
'preload_memory': ('.preload_memory_tool', 'preload_memory_tool'),
'request_input': ('._request_input_tool', 'request_input'),
'RemoteMcpServer': ('._remote_mcp_server', 'RemoteMcpServer'),
'ToolContext': ('.tool_context', 'ToolContext'),
'transfer_to_agent': ('.transfer_to_agent_tool', 'transfer_to_agent'),
'TransferToAgentTool': (
@@ -0,0 +1,68 @@
# 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 __future__ import annotations
from typing import Awaitable
from typing import Callable
from typing import TYPE_CHECKING
from pydantic import BaseModel
from pydantic import ConfigDict
if TYPE_CHECKING:
from ..agents.readonly_context import ReadonlyContext
HeaderProvider = Callable[
[ReadonlyContext], dict[str, str] | Awaitable[dict[str, str]]
]
else:
HeaderProvider = Callable[..., dict[str, str] | Awaitable[dict[str, str]]]
class RemoteMcpServer(BaseModel):
"""A remote MCP server executed server-side by the Managed Agents API.
``ManagedAgent`` forwards the server's URL and headers to
``interactions.create``; the Interactions backend opens the MCP session and
runs the tools. Only remote (HTTP/streamable) MCP servers are supported.
This is server-side MCP: unlike ``LlmAgent``'s ``McpToolset`` (which opens the
session and executes tools client-side), ADK never connects to the MCP server
here. The reused concept is the ``header_provider`` callback contract.
"""
model_config = ConfigDict(arbitrary_types_allowed=True, extra='forbid')
url: str
"""Full URL of the remote MCP server endpoint (e.g.
'https://api.example.com/mcp'). Maps to ``MCPServerParam.url``."""
name: str | None = None
"""Optional server label. Maps to ``MCPServerParam.name``."""
headers: dict[str, str] | None = None
"""Static headers sent on every turn (e.g. a fixed API key). Merged with
``header_provider`` output; ``header_provider`` wins on key conflict."""
allowed_tools: list[str] | None = None
"""Restrict which of the server's tools are exposed. Maps to
``MCPServerParam.allowed_tools``."""
header_provider: HeaderProvider | None = None
"""Runtime callback that mints headers (e.g. a fresh bearer token) at request
time. Invoked by ``ManagedAgent`` during resolution (runner-driven), once per
turn. Receives a ``ReadonlyContext`` and returns a headers dict (or an
awaitable of one). Same contract as ``LlmAgent``'s
``McpToolset.header_provider``."""
+49
View File
@@ -22,10 +22,13 @@ auth (ADC) is configured. Run explicitly:
from __future__ import annotations
import os
from google.adk.agents import ManagedAgent
from google.adk.runners import Runner
from google.adk.sessions.in_memory_session_service import InMemorySessionService
from google.adk.tools import google_search
from google.adk.tools import RemoteMcpServer
from google.adk.utils.context_utils import Aclosing
from google.genai import types
import pytest
@@ -121,3 +124,49 @@ async def test_code_execution_prime_sum():
assert (
'5117' in normalized
), f'expected the code-executed sum 5117; got: {answer!r}'
@pytest.mark.asyncio
# Server-side remote MCP (mcp_server tool param) is currently only accepted by
# the Gemini Developer API Interactions endpoint. The Vertex Interactions
# endpoint returns 400 invalid_request for it (google-genai likewise documents
# types.Tool.mcp_servers as unsupported on Vertex AI), so this live test is
# scoped to GOOGLE_AI.
@pytest.mark.parametrize('llm_backend', ['GOOGLE_AI'], indirect=True)
@pytest.mark.skipif(
not os.environ.get('GOOGLE_MAPS_API_KEY'),
reason='GOOGLE_MAPS_API_KEY not set',
)
async def test_remote_mcp_maps_grounding_lite():
agent = ManagedAgent(
name='managed_maps_agent',
agent_id=_AGENT_ID,
environment={'type': 'remote'},
tools=[
RemoteMcpServer(
name='maps_grounding_lite',
url='https://mapstools.googleapis.com/mcp',
header_provider=lambda ctx: {
'X-Goog-Api-Key': os.environ['GOOGLE_MAPS_API_KEY']
},
)
],
)
session_service = InMemorySessionService()
runner = Runner(
app_name='managed_agent_it',
agent=agent,
session_service=session_service,
)
session = await session_service.create_session(
app_name='managed_agent_it', user_id='test_user'
)
events = await _run_turn(
runner, session, 'Find a few coffee shops near Golden Gate Park.'
)
# Non-deterministic content; assert a non-empty grounded answer came back and
# no terminal error event was emitted.
assert _joined_text(events).strip()
assert not any(e.error_code for e in events)
@@ -13,6 +13,8 @@
# limitations under the License.
import asyncio
import subprocess
import sys
from typing import Any
from unittest.mock import MagicMock
@@ -22,9 +24,11 @@ from google.adk.agents.run_config import StreamingMode
from google.adk.events.event import Event
from google.adk.models.llm_response import LlmResponse
from google.adk.tools import google_search
from google.adk.tools._remote_mcp_server import RemoteMcpServer
from google.adk.tools.function_tool import FunctionTool
from google.genai import types
from google.genai import types as genai_types
from pydantic import ValidationError
import pytest
@@ -696,3 +700,219 @@ async def _drain_collect(agen):
async for e in agen:
out.append(e)
return out
def test_remote_mcp_server_constructs_and_is_exported():
import google.adk.tools as tools_pkg
server = RemoteMcpServer(
url='https://mcp.example.com/mcp',
name='example',
headers={'X-Static': 'v'},
allowed_tools=['a', 'b'],
header_provider=lambda ctx: {'Authorization': 'Bearer t'},
)
assert server.url == 'https://mcp.example.com/mcp'
assert server.name == 'example'
assert server.headers == {'X-Static': 'v'}
assert server.allowed_tools == ['a', 'b']
assert server.header_provider is not None
assert tools_pkg.RemoteMcpServer is RemoteMcpServer
assert 'RemoteMcpServer' in tools_pkg.__all__
def test_tools_import_first_has_no_cycle():
"""Importing google.adk.tools before google.adk.agents must not cycle.
Guards the leaf-module invariant for RemoteMcpServer: a future edit that adds
a runtime (non-TYPE_CHECKING) google.adk.agents import to
google.adk.tools._remote_mcp_server would reintroduce a circular import and
fail this test.
"""
subprocess.run(
[
sys.executable,
'-c',
(
'import google.adk.tools as t; t.RemoteMcpServer; '
'from google.adk.agents._managed_agent import ManagedAgent'
),
],
check=True,
)
def test_remote_mcp_server_defaults():
server = RemoteMcpServer(url='https://x/mcp')
assert server.name is None
assert server.headers is None
assert server.allowed_tools is None
assert server.header_provider is None
def test_remote_mcp_server_forbids_extra_fields():
with pytest.raises(ValidationError):
RemoteMcpServer(url='https://x/mcp', bogus='nope')
def _mcp_params(params):
return [p for p in params if p.get('type') == 'mcp_server']
def test_resolve_mcp_basic_mapping():
server = RemoteMcpServer(
url='https://mcp.example.com/mcp', name='example', allowed_tools=['a']
)
agent = ManagedAgent(
name='mgr', agent_id='agents/a', tools=[server], api_client=_FakeClient()
)
params = asyncio.run(agent._resolve_backend_tools(_ctx()))
assert {
'type': 'mcp_server',
'url': 'https://mcp.example.com/mcp',
'name': 'example',
'allowed_tools': [{'tools': ['a']}],
} in params
def test_resolve_mcp_sync_header_provider():
captured = {}
def provider(ctx):
captured['called'] = True
return {'Authorization': 'Bearer tok'}
server = RemoteMcpServer(url='https://x/mcp', header_provider=provider)
agent = ManagedAgent(
name='mgr', agent_id='agents/a', tools=[server], api_client=_FakeClient()
)
params = asyncio.run(agent._resolve_backend_tools(_ctx()))
assert captured['called'] is True
assert _mcp_params(params)[0]['headers'] == {'Authorization': 'Bearer tok'}
def test_resolve_mcp_async_header_provider():
async def provider(ctx):
return {'Authorization': 'Bearer async'}
server = RemoteMcpServer(url='https://x/mcp', header_provider=provider)
agent = ManagedAgent(
name='mgr', agent_id='agents/a', tools=[server], api_client=_FakeClient()
)
params = asyncio.run(agent._resolve_backend_tools(_ctx()))
assert _mcp_params(params)[0]['headers'] == {'Authorization': 'Bearer async'}
def test_resolve_mcp_merges_static_and_dynamic_dynamic_wins():
server = RemoteMcpServer(
url='https://x/mcp',
headers={'X-Static': 's', 'Shared': 'static'},
header_provider=lambda ctx: {'Shared': 'dynamic', 'X-Dyn': 'd'},
)
agent = ManagedAgent(
name='mgr', agent_id='agents/a', tools=[server], api_client=_FakeClient()
)
params = asyncio.run(agent._resolve_backend_tools(_ctx()))
assert _mcp_params(params)[0]['headers'] == {
'X-Static': 's',
'Shared': 'dynamic',
'X-Dyn': 'd',
}
def test_resolve_mcp_no_header_provider_static_only():
server = RemoteMcpServer(url='https://x/mcp', headers={'X-Static': 's'})
agent = ManagedAgent(
name='mgr', agent_id='agents/a', tools=[server], api_client=_FakeClient()
)
params = asyncio.run(agent._resolve_backend_tools(_ctx()))
assert _mcp_params(params)[0]['headers'] == {'X-Static': 's'}
def test_resolve_mcp_header_provider_error_propagates():
def boom(ctx):
raise RuntimeError('token mint failed')
server = RemoteMcpServer(url='https://x/mcp', header_provider=boom)
agent = ManagedAgent(
name='mgr', agent_id='agents/a', tools=[server], api_client=_FakeClient()
)
with pytest.raises(RuntimeError, match='token mint failed'):
asyncio.run(agent._resolve_backend_tools(_ctx()))
def test_resolve_mcp_mixed_with_builtin():
server = RemoteMcpServer(url='https://x/mcp')
agent = ManagedAgent(
name='mgr',
agent_id='agents/a',
tools=[google_search, server],
api_client=_FakeClient(),
)
params = asyncio.run(agent._resolve_backend_tools(_ctx()))
assert {'type': 'google_search'} in params
assert _mcp_params(params)
def test_resolve_mcp_empty_header_provider_omits_headers():
# A header_provider returning an empty dict (or None), with no static headers,
# must not add a 'headers' key to the mcp_server param.
server = RemoteMcpServer(url='https://x/mcp', header_provider=lambda ctx: {})
agent = ManagedAgent(
name='mgr', agent_id='agents/a', tools=[server], api_client=_FakeClient()
)
params = asyncio.run(agent._resolve_backend_tools(_ctx()))
assert 'headers' not in _mcp_params(params)[0]
def test_resolve_mcp_does_not_mutate_spec_headers():
original_headers = {'X-Static': 's'}
server = RemoteMcpServer(
url='https://x/mcp',
headers=original_headers,
header_provider=lambda ctx: {'Authorization': 'Bearer tok'},
)
agent = ManagedAgent(
name='mgr', agent_id='agents/a', tools=[server], api_client=_FakeClient()
)
asyncio.run(agent._resolve_backend_tools(_ctx()))
# The spec's original headers dict must be untouched by resolution.
assert server.headers == {'X-Static': 's'}
assert original_headers == {'X-Static': 's'}
def test_run_async_forwards_mcp_server_param():
client = _RecordingClient([[]])
server = RemoteMcpServer(
url='https://mcp.example.com/mcp',
header_provider=lambda ctx: {'X-Goog-Api-Key': 'k'},
)
agent = ManagedAgent(
name='mgr', agent_id='agents/a', tools=[server], api_client=client
)
asyncio.run(_drain(agent._run_async_impl(_user_ctx('hi'))))
create_kwargs = client.aio.interactions.calls[0]
mcp = [t for t in create_kwargs['tools'] if t['type'] == 'mcp_server'][0]
assert mcp['url'] == 'https://mcp.example.com/mcp'
assert mcp['headers'] == {'X-Goog-Api-Key': 'k'}
@@ -2147,3 +2147,44 @@ def test_create_interactions_surfaces_environment_id_non_stream():
responses = asyncio.run(_collect())
assert len(responses) == 1
assert responses[-1].environment_id == 'env_ns'
class TestBuildMcpServerParam:
"""Tests for _build_mcp_server_param."""
def _server(self, **kwargs):
from google.adk.tools._remote_mcp_server import RemoteMcpServer
kwargs.setdefault('url', 'https://mcp.example.com/mcp')
return RemoteMcpServer(**kwargs)
def test_minimal_url_only(self):
param = interactions_utils._build_mcp_server_param(self._server(), {})
assert param == {
'type': 'mcp_server',
'url': 'https://mcp.example.com/mcp',
}
def test_with_name(self):
param = interactions_utils._build_mcp_server_param(
self._server(name='maps'), {}
)
assert param['name'] == 'maps'
def test_with_headers(self):
param = interactions_utils._build_mcp_server_param(
self._server(), {'X-Goog-Api-Key': 'k'}
)
assert param['headers'] == {'X-Goog-Api-Key': 'k'}
def test_with_allowed_tools(self):
param = interactions_utils._build_mcp_server_param(
self._server(allowed_tools=['search_places']), {}
)
assert param['allowed_tools'] == [{'tools': ['search_places']}]
def test_omits_unset_fields(self):
param = interactions_utils._build_mcp_server_param(self._server(), {})
assert 'name' not in param
assert 'headers' not in param
assert 'allowed_tools' not in param