diff --git a/nanobot/agent/tools/mcp.py b/nanobot/agent/tools/mcp.py index b1060155..bd47d77f 100644 --- a/nanobot/agent/tools/mcp.py +++ b/nanobot/agent/tools/mcp.py @@ -913,9 +913,9 @@ async def connect_mcp_servers( registered_count = 0 matched_enabled_tools: set[str] = set() available_raw_names = [tool_def.name for tool_def in tools.tools] - available_wrapped_names = [_sanitize_name(f"mcp_{name}_{tool_def.name}") for tool_def in tools.tools] + available_wrapped_names = [_sanitize_mcp_tool_name(f"mcp_{name}_{tool_def.name}") for tool_def in tools.tools] for tool_def in tools.tools: - wrapped_name = _sanitize_name(f"mcp_{name}_{tool_def.name}") + wrapped_name = _sanitize_mcp_tool_name(f"mcp_{name}_{tool_def.name}") if ( not allow_all_tools and tool_def.name not in enabled_tools @@ -1294,11 +1294,10 @@ def _attach_reconnect_handlers( ) for server_name in server_names: - prefix = _tool_prefix(server_name) for tool_name in list(registry.tool_names): - if not tool_name.startswith(prefix): - continue tool = registry.get(tool_name) + if not _tool_belongs_to_server(tool, tool_name, server_name): + continue if isinstance(tool, _MCPWrapperBase): tool.set_reconnect_handler(reconnect) @@ -1351,11 +1350,17 @@ def _tool_prefix(server_name: str) -> str: return _sanitize_name(f"mcp_{server_name}_") +def _tool_belongs_to_server(tool: Tool | None, tool_name: str, server_name: str) -> bool: + if isinstance(tool, _MCPWrapperBase): + return getattr(tool, "_server_name", None) == server_name + return tool_name.startswith(_tool_prefix(server_name)) + + def _unregister_server_tools(state: Any, registry: ToolRegistry, server_name: str) -> int: - prefix = _tool_prefix(server_name) removed = 0 for tool_name in list(registry.tool_names): - if tool_name.startswith(prefix): + tool = registry.get(tool_name) + if _tool_belongs_to_server(tool, tool_name, server_name): registry.unregister(tool_name) removed += 1 return removed diff --git a/tests/tools/test_mcp_tool.py b/tests/tools/test_mcp_tool.py index 547efdb6..97dc9280 100644 --- a/tests/tools/test_mcp_tool.py +++ b/tests/tools/test_mcp_tool.py @@ -1332,3 +1332,35 @@ def test_mcp_tool_name_limits_long_name(): assert len(name) <= 64 assert name.startswith("mcp_") + + +def test_long_server_name_tools_are_matched_by_server_name() -> None: + server_name = "very-long-server-name-" * 4 + tool_def = SimpleNamespace( + name="search", + description="search tool", + inputSchema={"type": "object", "properties": {}}, + ) + other_tool_def = SimpleNamespace( + name="search", + description="other search tool", + inputSchema={"type": "object", "properties": {}}, + ) + wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), server_name, tool_def) + other_wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "other", other_tool_def) + registry = ToolRegistry() + registry.register(wrapper) + registry.register(other_wrapper) + + assert len(wrapper.name) == 64 + assert not wrapper.name.startswith(mcp_mod._tool_prefix(server_name)) + + mcp_mod._attach_reconnect_handlers(SimpleNamespace(), registry, {server_name}) + assert wrapper._reconnect is not None + assert other_wrapper._reconnect is None + + removed = mcp_mod._unregister_server_tools(SimpleNamespace(), registry, server_name) + + assert removed == 1 + assert wrapper.name not in registry.tool_names + assert other_wrapper.name in registry.tool_names