diff --git a/scripts/compliance_checks.py b/scripts/compliance_checks.py index 56a23d49..6f96e5ed 100755 --- a/scripts/compliance_checks.py +++ b/scripts/compliance_checks.py @@ -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', diff --git a/src/google/adk/agents/_managed_agent.py b/src/google/adk/agents/_managed_agent.py index a7ea2271..df13f008 100644 --- a/src/google/adk/agents/_managed_agent.py +++ b/src/google/adk/agents/_managed_agent.py @@ -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 diff --git a/src/google/adk/models/interactions_utils.py b/src/google/adk/models/interactions_utils.py index 28bd4467..9f251746 100644 --- a/src/google/adk/models/interactions_utils.py +++ b/src/google/adk/models/interactions_utils.py @@ -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]: diff --git a/src/google/adk/tools/__init__.py b/src/google/adk/tools/__init__.py index 641c5c20..9cd22f30 100644 --- a/src/google/adk/tools/__init__.py +++ b/src/google/adk/tools/__init__.py @@ -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': ( diff --git a/src/google/adk/tools/_remote_mcp_server.py b/src/google/adk/tools/_remote_mcp_server.py new file mode 100644 index 00000000..8a2b43ee --- /dev/null +++ b/src/google/adk/tools/_remote_mcp_server.py @@ -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``.""" diff --git a/tests/integration/test_managed_agent.py b/tests/integration/test_managed_agent.py index dde6d56a..c8d05f70 100644 --- a/tests/integration/test_managed_agent.py +++ b/tests/integration/test_managed_agent.py @@ -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) diff --git a/tests/unittests/agents/test_managed_agent.py b/tests/unittests/agents/test_managed_agent.py index 7bb2841a..2117e441 100644 --- a/tests/unittests/agents/test_managed_agent.py +++ b/tests/unittests/agents/test_managed_agent.py @@ -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'} diff --git a/tests/unittests/models/test_interactions_utils.py b/tests/unittests/models/test_interactions_utils.py index 49360fd7..28ff8a2a 100644 --- a/tests/unittests/models/test_interactions_utils.py +++ b/tests/unittests/models/test_interactions_utils.py @@ -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