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.
This commit is contained in:
parent
6fa197f973
commit
3198cc8fd9
2 changed files with 102 additions and 2 deletions
|
|
@ -2447,3 +2447,79 @@ class TestDiscoveryFailedCount:
|
||||||
_servers.pop("ok1", None)
|
_servers.pop("ok1", None)
|
||||||
_servers.pop("ok2", None)
|
_servers.pop("ok2", None)
|
||||||
_servers.pop("fail1", 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
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1343,7 +1343,27 @@ async def _discover_and_register_server(name: str, config: dict) -> List[str]:
|
||||||
registered_names: List[str] = []
|
registered_names: List[str] = []
|
||||||
toolset_name = f"mcp-{name}"
|
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:
|
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)
|
schema = _convert_mcp_schema(name, mcp_tool)
|
||||||
tool_name_prefixed = schema["name"]
|
tool_name_prefixed = schema["name"]
|
||||||
|
|
||||||
|
|
@ -1424,9 +1444,13 @@ def discover_mcp_tools() -> List[str]:
|
||||||
logger.debug("No MCP servers configured")
|
logger.debug("No MCP servers configured")
|
||||||
return []
|
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:
|
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:
|
if not new_servers:
|
||||||
return _existing_tool_names()
|
return _existing_tool_names()
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue