Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 29 additions & 16 deletions python/packages/core/agent_framework/_mcp.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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,
Expand Down
46 changes: 46 additions & 0 deletions python/packages/core/tests/core/test_mcp.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]

Expand Down
Loading