diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 9abafdad27..54cdaa1d87 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -240,6 +240,34 @@ def _mcp_config_candidate_names(*, local_name: str, normalized_name: str, remote return tuple(names) +def _make_mcp_tool_caller( + mcp_tool: MCPTool, remote_tool_name: str +) -> Callable[..., Coroutine[Any, Any, str | list[Content]]]: + """Build the callable backing a generated MCP ``FunctionTool``. + + The remote tool name is captured in this factory's closure rather than declared as a + parameter of the returned callable. Model-supplied arguments are splatted into that + callable, so any name it declares could be bound - and overridden - by the model. Keeping + the call target out of the signature means it can only ever be reached through ``**kwargs``, + which is forwarded as tool arguments and can never redirect the call to another tool. + """ + + async def _call_tool_with_runtime_kwargs( + ctx: FunctionInvocationContext, + **kwargs: Any, + ) -> str | list[Content]: + trusted_meta = ctx.kwargs.get("_meta") + call_kwargs = dict(ctx.kwargs) + call_kwargs.update(kwargs) + if trusted_meta is not None: + call_kwargs["_meta"] = trusted_meta + else: + call_kwargs.pop("_meta", None) + return await mcp_tool.call_tool(remote_tool_name, **call_kwargs) + + return _call_tool_with_runtime_kwargs + + def _validate_mcp_meta_key(key: str) -> None: """Validate an MCP ``_meta`` key against the 2025-06-18 key-name format.""" if not _MCP_META_KEY_PATTERN.fullmatch(key): @@ -1894,24 +1922,9 @@ async def _load_tools_locked(self) -> None: ) ) - async def _call_tool_with_runtime_kwargs( - ctx: FunctionInvocationContext, - *, - _remote_tool_name: str = tool.name, - **kwargs: Any, - ) -> str | list[Content]: - trusted_meta = ctx.kwargs.get("_meta") - call_kwargs = dict(ctx.kwargs) - call_kwargs.update(kwargs) - if trusted_meta is not None: - call_kwargs["_meta"] = trusted_meta - else: - call_kwargs.pop("_meta", None) - return await self.call_tool(_remote_tool_name, **call_kwargs) - # Create FunctionTools out of each tool func: FunctionTool = FunctionTool( - func=_call_tool_with_runtime_kwargs, + func=_make_mcp_tool_caller(self, tool.name), name=local_name, description=tool.description or "", approval_mode=approval_mode, diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 028cac027b..4c27e596f8 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -4744,6 +4744,52 @@ async def test_mcp_tool_call_tool_requires_loaded_tools() -> None: await tool.call_tool("remote_tool") +async def test_generated_mcp_function_ignores_model_supplied_remote_tool_name() -> None: + """A model-supplied argument must not be able to redirect the call to another remote tool.""" + tool = MCPTool(name="test_tool") # type: ignore[abstract] + tool.session = Mock(spec=ClientSession) + tool.session.list_tools = AsyncMock( # ty: ignore[unresolved-attribute] + return_value=types.ListToolsResult( + tools=[ + types.Tool( + name="search_docs", + description="Search docs.", + inputSchema={ + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, + ), + types.Tool( + name="delete_repo", + description="Delete a repository.", + inputSchema={ + "type": "object", + "properties": {"repo": {"type": "string"}}, + "required": ["repo"], + }, + ), + ] + ) + ) + tool.session.call_tool = AsyncMock( # ty: ignore[unresolved-attribute] + return_value=types.CallToolResult(content=[types.TextContent(type="text", text="ok")]) + ) + + await tool.load_tools() + + search_docs = next(func for func in tool.functions if func.name == "search_docs") + await search_docs.invoke( + arguments={"query": "quarterly report", "_remote_tool_name": "delete_repo", "repo": "corp/prod"} + ) + + tool.session.call_tool.assert_awaited_once() # ty: ignore[unresolved-attribute] + await_args = tool.session.call_tool.await_args # ty: ignore[unresolved-attribute] + assert await_args is not None + assert await_args.args[0] == "search_docs" + assert await_args.kwargs["arguments"] == {"query": "quarterly report"} + + async def test_mcp_tool_get_prompt_requires_loaded_prompts() -> None: tool = MCPTool(name="test_tool", load_prompts=False) # type: ignore[abstract]