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:
committed by
Copybara-Service
parent
283e92efc2
commit
2e2ec09a76
@@ -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',
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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]:
|
||||
|
||||
@@ -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``."""
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user