From 7248475fe2dfa0465c407b5dc7ce557df66f30cf Mon Sep 17 00:00:00 2001 From: teyrebaz33 Date: Thu, 12 Mar 2026 02:31:09 +0300 Subject: [PATCH] feat(mcp): per-server tool filtering via include/exclude and enabled flag Add optional config keys under each mcp_servers entry: - tools.include: whitelist, only listed tools are registered - tools.exclude: blacklist, all tools except listed are registered - enabled: false: skip server entirely, no connection attempt Backward-compatible: no config keys = all tools registered as before. Tests: TestMCPSelectiveToolLoading (4 tests), 134 passed total. --- tests/tools/test_mcp_tool.py | 76 ++++++++++++++++++++++++++++++++++++ tools/mcp_tool.py | 28 ++++++++++++- 2 files changed, 102 insertions(+), 2 deletions(-) diff --git a/tests/tools/test_mcp_tool.py b/tests/tools/test_mcp_tool.py index 0f7fc18a5..4a77afad0 100644 --- a/tests/tools/test_mcp_tool.py +++ b/tests/tools/test_mcp_tool.py @@ -2450,3 +2450,79 @@ async def selective_register(name, cfg): _servers.pop("ok1", None) _servers.pop("ok2", None) _servers.pop("fail1", None) + + +class TestMCPSelectiveToolLoading: + """Tests for tools.include / tools.exclude / enabled config keys.""" + + def _make_server(self, name, tool_names): + from tools.mcp_tool import MCPServerTask + server = MCPServerTask(name) + server.session = MagicMock() + server._tools = [_make_mcp_tool(n, n) for n in tool_names] + return server + + def _run_discover(self, name, tool_names, config): + """Run _discover_and_register_server directly and return registered names.""" + import asyncio + from tools.mcp_tool import _discover_and_register_server + server = self._make_server(name, tool_names) + + async def fake_connect(n, c): + return server + + async def run(): + with patch("tools.mcp_tool._connect_server", side_effect=fake_connect), patch("tools.mcp_tool._servers", {}): + return await _discover_and_register_server(name, config) + + return asyncio.run(run()) + + def test_include_filter_registers_only_listed_tools(self): + """tools.include whitelist: only specified tools are registered.""" + tool_names = ["create_service", "delete_service", "list_services"] + config = {"url": "https://mcp.example.com", "tools": {"include": ["create_service", "list_services"]}} + result = self._run_discover("ink", tool_names, config) + assert "mcp_ink_create_service" in result + assert "mcp_ink_list_services" in result + assert "mcp_ink_delete_service" not in result + + def test_exclude_filter_skips_listed_tools(self): + """tools.exclude blacklist: all tools except specified are registered.""" + tool_names = ["create_service", "delete_service", "list_services"] + config = {"url": "https://mcp.example.com", "tools": {"exclude": ["delete_service"]}} + result = self._run_discover("ink2", tool_names, config) + assert "mcp_ink2_create_service" in result + assert "mcp_ink2_list_services" in result + assert "mcp_ink2_delete_service" not in result + + def test_no_filter_registers_all_tools(self): + """No tools filter: all tools registered (backward compatible).""" + tool_names = ["create_service", "delete_service", "list_services"] + config = {"url": "https://mcp.example.com"} + result = self._run_discover("ink3", tool_names, config) + assert "mcp_ink3_create_service" in result + assert "mcp_ink3_delete_service" in result + assert "mcp_ink3_list_services" in result + + def test_enabled_false_skips_server(self): + """enabled: false skips the server entirely.""" + fresh_servers = {} + fake_config = { + "ink": { + "url": "https://mcp.example.com", + "enabled": False, + } + } + connect_called = [] + + async def fake_connect(name, config): + connect_called.append(name) + return self._make_server(name, ["create_service"]) + + with patch("tools.mcp_tool._MCP_AVAILABLE", True), patch("tools.mcp_tool._servers", fresh_servers), patch("tools.mcp_tool._load_mcp_config", return_value=fake_config), patch("tools.mcp_tool._connect_server", side_effect=fake_connect): + from tools.mcp_tool import discover_mcp_tools + result = discover_mcp_tools() + + assert connect_called == [] + assert "mcp_ink_create_service" not in result + diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py index 94495430b..0b524652e 100644 --- a/tools/mcp_tool.py +++ b/tools/mcp_tool.py @@ -1235,7 +1235,27 @@ async def _discover_and_register_server(name: str, config: dict) -> List[str]: registered_names: List[str] = [] toolset_name = f"mcp-{name}" + # Selective tool loading: honour include/exclude lists from config. + # Rules (matching issue #690 spec): + # tools.include — whitelist: only these tool names are registered + # tools.exclude — blacklist: all tools EXCEPT these are registered + # include and exclude are mutually exclusive; include takes precedence + # Neither set → register all tools (backward-compatible default) + tools_filter = config.get("tools") or {} + include_set = set(tools_filter.get("include") or []) + exclude_set = set(tools_filter.get("exclude") or []) + + def _should_register(tool_name: str) -> bool: + if include_set: + return tool_name in include_set + if exclude_set: + return tool_name not in exclude_set + return True + for mcp_tool in server._tools: + if not _should_register(mcp_tool.name): + logger.debug("MCP server '%s': skipping tool '%s' (filtered by config)", name, mcp_tool.name) + continue schema = _convert_mcp_schema(name, mcp_tool) tool_name_prefixed = schema["name"] @@ -1316,9 +1336,13 @@ def discover_mcp_tools() -> List[str]: logger.debug("No MCP servers configured") return [] - # Only attempt servers that aren't already connected + # Only attempt servers that aren't already connected and are enabled + # (enabled: false skips the server entirely without removing its config) with _lock: - new_servers = {k: v for k, v in servers.items() if k not in _servers} + new_servers = { + k: v for k, v in servers.items() + if k not in _servers and v.get("enabled", True) is not False + } if not new_servers: return _existing_tool_names()