mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
3f8f91d064
* fix(mcp): isolate pooled sessions by ownership domain Separate deployment and personal stdio sessions even when they share the same runtime server name and thread scope. Preserve deployment call compatibility while carrying personal ownership through Agent wrappers, durable task calls, and disconnect eviction. Add focused regression coverage for personal discovery wiring and domain-scoped disconnect cleanup. * test(mcp): address ownership-domain review feedback Accept the new wrapper keyword in the prefix-overlap test double and reuse MCPPoolDomain for internal task-caller ownership parameters. * style(mcp): format ownership-domain changes * style(mcp): apply repository ruff formatting
239 lines
8.8 KiB
Python
239 lines
8.8 KiB
Python
from pathlib import Path
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
from langchain_core.tools import StructuredTool
|
|
from pydantic import BaseModel, Field
|
|
|
|
from app.gateway.routers.mcp import McpServerConfigResponse
|
|
from deerflow.config.extensions_config import ExtensionsConfig, McpServerConfig
|
|
from deerflow.mcp.tools import _make_session_pool_tool, get_mcp_tools
|
|
from deerflow.tools.mcp_metadata import get_mcp_routing
|
|
|
|
|
|
class _Args(BaseModel):
|
|
query: str = Field(..., description="query")
|
|
|
|
|
|
def _tool(name: str) -> StructuredTool:
|
|
async def _call(query: str) -> str:
|
|
return query
|
|
|
|
return StructuredTool(
|
|
name=name,
|
|
description="Search",
|
|
args_schema=_Args,
|
|
coroutine=_call,
|
|
)
|
|
|
|
|
|
def test_mcp_tool_name_prefix_is_an_explicit_default_true_field() -> None:
|
|
assert "tool_name_prefix" in McpServerConfig.model_fields
|
|
assert McpServerConfig().tool_name_prefix is True
|
|
assert McpServerConfig(tool_name_prefix=False).model_dump()["tool_name_prefix"] is False
|
|
|
|
|
|
def test_gateway_mcp_config_preserves_tool_name_prefix() -> None:
|
|
response = McpServerConfigResponse.model_validate(McpServerConfig(tool_name_prefix=False).model_dump())
|
|
|
|
assert response.tool_name_prefix is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mcp_tool_name_prefix_can_be_disabled_per_server_without_disabling_stdio_pooling() -> None:
|
|
extensions_config = ExtensionsConfig.model_validate(
|
|
{
|
|
"mcpServers": {
|
|
"semantic_scholar": {
|
|
"type": "stdio",
|
|
"command": "uvx",
|
|
"args": ["s2-mcp-server"],
|
|
"tool_name_prefix": False,
|
|
"routing": {"mode": "prefer", "priority": 10},
|
|
"tools": {
|
|
"semantic_scholar_search_papers": {
|
|
"routing": {"priority": 99},
|
|
}
|
|
},
|
|
},
|
|
"github": {
|
|
"type": "http",
|
|
"url": "https://example.test/mcp",
|
|
},
|
|
}
|
|
}
|
|
)
|
|
servers_config = {
|
|
"semantic_scholar": {
|
|
"transport": "stdio",
|
|
"command": "uvx",
|
|
"args": ["s2-mcp-server"],
|
|
},
|
|
"github": {
|
|
"transport": "http",
|
|
"url": "https://example.test/mcp",
|
|
},
|
|
}
|
|
raw_tools = {
|
|
"semantic_scholar": _tool("semantic_scholar_search_papers"),
|
|
"github": _tool("search_repositories"),
|
|
}
|
|
|
|
class FakeClient:
|
|
def __init__(
|
|
self,
|
|
connections,
|
|
*,
|
|
callbacks=None,
|
|
tool_interceptors=None,
|
|
tool_name_prefix=False,
|
|
) -> None:
|
|
self.connections = connections
|
|
self.callbacks = callbacks
|
|
self.tool_interceptors = tool_interceptors or []
|
|
self.tool_name_prefix = tool_name_prefix
|
|
|
|
async def get_tools(self, *, server_name=None):
|
|
tool = raw_tools[server_name]
|
|
name = f"{server_name}_{tool.name}" if self.tool_name_prefix else tool.name
|
|
return [_tool(name)]
|
|
|
|
async def fake_load_mcp_tools(
|
|
session,
|
|
*,
|
|
connection,
|
|
callbacks=None,
|
|
tool_interceptors=None,
|
|
server_name=None,
|
|
tool_name_prefix=False,
|
|
):
|
|
assert session is None
|
|
assert connection is servers_config[server_name]
|
|
tool = raw_tools[server_name]
|
|
name = f"{server_name}_{tool.name}" if tool_name_prefix else tool.name
|
|
return [_tool(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", new_callable=AsyncMock, return_value={}),
|
|
patch("deerflow.mcp.tools.build_mcp_tool_interceptors", return_value=[]),
|
|
patch("langchain_mcp_adapters.client.MultiServerMCPClient", FakeClient),
|
|
patch("langchain_mcp_adapters.tools.load_mcp_tools", side_effect=fake_load_mcp_tools),
|
|
patch("deerflow.mcp.tools._make_session_pool_tool", side_effect=lambda tool, *_args, **_kwargs: tool) as wrap_tool,
|
|
):
|
|
tools = await get_mcp_tools()
|
|
|
|
assert {tool.name for tool in tools} == {
|
|
"semantic_scholar_search_papers",
|
|
"github_search_repositories",
|
|
}
|
|
semantic_scholar_tool = next(tool for tool in tools if tool.name == "semantic_scholar_search_papers")
|
|
routing = get_mcp_routing(semantic_scholar_tool)
|
|
assert routing is not None
|
|
assert routing["priority"] == 99
|
|
wrap_tool.assert_called_once()
|
|
assert wrap_tool.call_args.args[1] == "semantic_scholar"
|
|
assert wrap_tool.call_args.kwargs["tool_name_prefix"] is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_personal_discovery_routes_stdio_wrapper_to_personal_pool_domain() -> None:
|
|
extensions_config = ExtensionsConfig.model_validate(
|
|
{
|
|
"mcpServers": {
|
|
"github": {
|
|
"type": "stdio",
|
|
"command": "uvx",
|
|
"args": ["github-mcp"],
|
|
}
|
|
}
|
|
}
|
|
)
|
|
servers_config = {
|
|
"github": {
|
|
"transport": "stdio",
|
|
"command": "uvx",
|
|
"args": ["github-mcp"],
|
|
}
|
|
}
|
|
|
|
class FakeClient:
|
|
def __init__(self, _connections, *, callbacks=None, tool_interceptors=None, **_kwargs) -> None:
|
|
self.callbacks = callbacks
|
|
self.tool_interceptors = tool_interceptors or []
|
|
|
|
async def get_tools(self, *, server_name=None):
|
|
assert server_name == "github"
|
|
return [_tool("github_search")]
|
|
|
|
with (
|
|
patch("deerflow.mcp.user_config.load_user_mcp_config", return_value=extensions_config),
|
|
patch("deerflow.mcp.personal_access.authorized_personal_config", new_callable=AsyncMock, return_value=extensions_config),
|
|
patch("deerflow.mcp.tools.build_servers_config", return_value=servers_config),
|
|
patch("deerflow.mcp.tools.get_initial_oauth_headers", new_callable=AsyncMock, return_value={}),
|
|
patch("deerflow.mcp.tools.build_mcp_tool_interceptors", return_value=[]),
|
|
patch("langchain_mcp_adapters.client.MultiServerMCPClient", FakeClient),
|
|
patch("deerflow.mcp.tools._make_session_pool_tool", side_effect=lambda tool, *_args, **_kwargs: tool) as wrap_tool,
|
|
):
|
|
tools = await get_mcp_tools(extensions_config, personal_user_id="alice")
|
|
|
|
assert [tool.name for tool in tools] == ["github_search"]
|
|
wrap_tool.assert_called_once()
|
|
assert wrap_tool.call_args.kwargs["ownership_domain"] == "personal"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unprefixed_stdio_tool_keeps_server_like_original_name(tmp_path: Path) -> None:
|
|
"""Pooling must not strip a prefix that belongs to the MCP tool itself."""
|
|
original_tool = _tool("github_search")
|
|
mock_session = AsyncMock()
|
|
mock_session.call_tool = AsyncMock(return_value=MagicMock(content=[], isError=False, structuredContent=None))
|
|
mock_pool = MagicMock()
|
|
mock_pool.get_session = AsyncMock(return_value=mock_session)
|
|
|
|
with (
|
|
patch("deerflow.mcp.tools.get_session_pool", return_value=mock_pool),
|
|
patch("deerflow.mcp.tools.get_paths", return_value=MagicMock()),
|
|
patch(
|
|
"deerflow.mcp.tools._prepare_stdio_workspace",
|
|
return_value=(tmp_path, tmp_path / "tmp", {}),
|
|
),
|
|
):
|
|
wrapped = _make_session_pool_tool(
|
|
original_tool,
|
|
"github",
|
|
{"transport": "stdio", "command": "mcp-server", "args": []},
|
|
tool_name_prefix=False,
|
|
)
|
|
await wrapped.coroutine(query="repositories")
|
|
|
|
mock_session.call_tool.assert_awaited_once_with("github_search", {"query": "repositories"})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_personal_stdio_wrapper_uses_personal_pool_domain(tmp_path: Path) -> None:
|
|
original_tool = _tool("github_search")
|
|
mock_session = AsyncMock()
|
|
mock_session.call_tool = AsyncMock(return_value=MagicMock(content=[], isError=False, structuredContent=None))
|
|
mock_pool = MagicMock()
|
|
mock_pool.get_session = AsyncMock(return_value=mock_session)
|
|
|
|
with (
|
|
patch("deerflow.mcp.tools.get_session_pool", return_value=mock_pool),
|
|
patch("deerflow.mcp.tools.get_paths", return_value=MagicMock()),
|
|
patch(
|
|
"deerflow.mcp.tools._prepare_stdio_workspace",
|
|
return_value=(tmp_path, tmp_path / "tmp", {}),
|
|
),
|
|
):
|
|
wrapped = _make_session_pool_tool(
|
|
original_tool,
|
|
"github",
|
|
{"transport": "stdio", "command": "mcp-server", "args": []},
|
|
ownership_domain="personal",
|
|
)
|
|
await wrapped.coroutine(query="repositories")
|
|
|
|
assert mock_pool.get_session.await_args.kwargs["domain"] == "personal"
|