mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-24 15:07:49 +00:00
fix(mcp): route tools by source server, not name prefix (#3812)
* fix(mcp): route tools by source server to fix prefix-collision mis-routing get_mcp_tools() flattened the per-server tool lists and then re-derived each tool's server by scanning servers_config for a name that is a prefix of the (prefixed) tool name, taking the first match. When one server name is a prefix of another (e.g. "web" and "web_scraper"), a "web_scraper_search" tool matches "web" first, so it is pooled under the wrong server and invoked with the wrong stripped name — the call fails, or for mixed transports the stdio tool loses session pooling. Route each tool by the server that actually produced it (tools_by_server[i] corresponds to the i-th server in servers_config), keeping the prefix guard so unprefixed tools still fall through unwrapped. Add a regression test covering overlapping server-name prefixes. * Apply suggestions from code review Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
Aari
GitHub
parent
72f033fbbe
commit
94c852ac4a
@ -652,30 +652,28 @@ async def get_mcp_tools() -> list[BaseTool]:
|
||||
# internally which cannot be closed from a different async task, so
|
||||
# pooling them causes RuntimeError on cleanup (see #3203).
|
||||
wrapped_tools: list[BaseTool] = []
|
||||
for tool in tools:
|
||||
tool_server: str | None = None
|
||||
for name in servers_config:
|
||||
if tool.name.startswith(f"{name}_"):
|
||||
tool_server = name
|
||||
break
|
||||
|
||||
if tool_server is not None:
|
||||
transport = servers_config[tool_server].get("transport", "stdio")
|
||||
if transport == "stdio":
|
||||
server_cfg = extensions_config.mcp_servers.get(tool_server)
|
||||
# Route each tool by the server that actually produced it: tools_by_server[i]
|
||||
# corresponds to the i-th server in servers_config. Inferring the source server by
|
||||
# scanning servers_config for a name prefix is ambiguous when one server name is a
|
||||
# prefix of another (e.g. "web" vs "web_scraper" → "web_scraper_search".startswith(
|
||||
# "web_") matches "web" first), which pools the tool under the wrong server. Using the
|
||||
# source grouping makes routing exact; the prefix guard preserves the previous
|
||||
# behavior of leaving unprefixed tools unwrapped.
|
||||
for source_name, server_tools in zip(servers_config.keys(), tools_by_server, strict=True):
|
||||
transport = servers_config[source_name].get("transport", "stdio")
|
||||
server_cfg = extensions_config.mcp_servers.get(source_name)
|
||||
for tool in server_tools:
|
||||
if tool.name.startswith(f"{source_name}_") and transport == "stdio":
|
||||
_timeout = server_cfg.tool_call_timeout if server_cfg else None
|
||||
wrapped_tools.append(_make_session_pool_tool(tool, tool_server, servers_config[tool_server], tool_interceptors, tool_call_timeout=_timeout))
|
||||
wrapped_tools.append(_make_session_pool_tool(tool, source_name, servers_config[source_name], tool_interceptors, tool_call_timeout=_timeout))
|
||||
else:
|
||||
server_cfg = extensions_config.mcp_servers.get(tool_server)
|
||||
if server_cfg and server_cfg.tool_call_timeout is not None:
|
||||
if transport != "stdio" and server_cfg and server_cfg.tool_call_timeout is not None:
|
||||
logger.warning(
|
||||
"Ignoring tool_call_timeout for MCP server '%s' because transport '%s' is not stdio; configure HTTP/SSE transport-level timeouts instead.",
|
||||
tool_server,
|
||||
source_name,
|
||||
transport,
|
||||
)
|
||||
wrapped_tools.append(tool)
|
||||
else:
|
||||
wrapped_tools.append(tool)
|
||||
|
||||
# Patch tools to support sync invocation, as deerflow client streams synchronously
|
||||
for tool in wrapped_tools:
|
||||
|
||||
@ -1625,3 +1625,80 @@ def test_reset_mcp_tools_cache_from_running_loop_is_bounded():
|
||||
|
||||
assert done.is_set(), "reset_mcp_tools_cache() deadlocked inside a running loop"
|
||||
assert cm.closed is True, "owner task must run __aexit__ once the loop regains control"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_mcp_tools: routing when one server name is a prefix of another
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_tools_routed_to_source_server_with_prefix_overlap():
|
||||
"""Regression: tools must be routed to the server that produced them, not the first
|
||||
server whose name is a string prefix of the (prefixed) tool name.
|
||||
|
||||
With `tool_name_prefix=True`, a tool from server `web_scraper` is named
|
||||
`web_scraper_search`. When a server `web` is also configured, prefix-matching the tool
|
||||
name picks `web` first (`"web_scraper_search".startswith("web_")`), mis-routing the
|
||||
tool and stripping it to the wrong original name. Routing by the source grouping fixes it.
|
||||
"""
|
||||
from langchain_core.tools import StructuredTool
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from deerflow.mcp.tools import get_mcp_tools
|
||||
|
||||
class Args(BaseModel):
|
||||
query: str = Field(..., description="query")
|
||||
|
||||
web_tool = StructuredTool(
|
||||
name="web_open",
|
||||
description="d",
|
||||
args_schema=Args,
|
||||
coroutine=AsyncMock(),
|
||||
response_format="content_and_artifact",
|
||||
)
|
||||
scraper_tool = StructuredTool(
|
||||
name="web_scraper_search",
|
||||
description="d",
|
||||
args_schema=Args,
|
||||
coroutine=AsyncMock(),
|
||||
response_format="content_and_artifact",
|
||||
)
|
||||
|
||||
extensions_config = MagicMock()
|
||||
extensions_config.model_extra = {}
|
||||
|
||||
# `web` is inserted before `web_scraper`, so a first-prefix-match mis-routes
|
||||
# `web_scraper_search` to `web`.
|
||||
servers_config = {
|
||||
"web": {"transport": "stdio", "command": "npx", "args": ["web"]},
|
||||
"web_scraper": {"transport": "stdio", "command": "npx", "args": ["scraper"]},
|
||||
}
|
||||
|
||||
routed: list[tuple[str, str]] = []
|
||||
|
||||
def fake_wrap(tool, server_name, connection, interceptors, tool_call_timeout=None):
|
||||
routed.append((tool.name, server_name))
|
||||
return tool
|
||||
|
||||
async def get_tools_for_server(*, server_name: str | None = None):
|
||||
if server_name == "web":
|
||||
return [web_tool]
|
||||
if server_name == "web_scraper":
|
||||
return [scraper_tool]
|
||||
raise AssertionError(f"unexpected server_name: {server_name}")
|
||||
|
||||
with (
|
||||
patch("deerflow.mcp.tools.ExtensionsConfig.from_file", return_value=extensions_config),
|
||||
patch("deerflow.mcp.tools.build_servers_config", return_value=servers_config),
|
||||
patch("deerflow.mcp.tools.get_initial_oauth_headers", return_value={}),
|
||||
patch("deerflow.mcp.tools.build_oauth_tool_interceptor", return_value=None),
|
||||
patch("langchain_mcp_adapters.client.MultiServerMCPClient") as MockClient,
|
||||
patch("deerflow.mcp.tools._make_session_pool_tool", side_effect=fake_wrap),
|
||||
):
|
||||
MockClient.return_value.get_tools = AsyncMock(side_effect=get_tools_for_server)
|
||||
await get_mcp_tools()
|
||||
|
||||
routing = dict(routed)
|
||||
assert routing["web_scraper_search"] == "web_scraper", f"tool mis-routed to {routing.get('web_scraper_search')!r}, expected 'web_scraper'"
|
||||
assert routing["web_open"] == "web"
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user