Merge pull request 'Backport tool enable / disable from main' (#439) from backport-tool-selection into development

Reviewed-on: https://git.biggo.com/Funmula/dive-mcp-host/pulls/439

Author: LouisFunmula

dive_mcp_host/host/tools/mcp_server.py

@@ -65,12 +65,23 @@
 logger = getLogger(__name__)
 
 
+class ToolInfo(types.Tool):
+    """Custom tool info with extra info."""
+
+    enable: bool
+
+    @classmethod
+    def from_tool(cls, tool: types.Tool, enable: bool) -> Self:
+        """Create from mcp Tool type."""
+        return cls(**tool.model_dump(), enable=enable)
+
+
 class McpServerInfo(BaseModel):
     """MCP server capability and tool list."""
 
     name: str
     """The name of the MCP server."""
-    tools: list[types.Tool]
+    tools: list[ToolInfo]
     """The tools provided by the MCP server."""
     initialize_result: types.InitializeResult | None
     """The result of the initialize method.
@@ -239,19 +250,35 @@ def log_buffer(self) -> LogBuffer:
     @property
     def server_info(self) -> McpServerInfo:
         """Get the server info."""
+        tools: list[ToolInfo] = []
+        if self._tool_results:
+            for tool in self._tool_results.tools:
+                enable: bool = True
+                if tool.name in self.config.exclude_tools:
+                    enable = False
+                tools.append(ToolInfo.from_tool(tool, enable))
+
         return McpServerInfo(
             name=self.name,
             initialize_result=self._initialize_result,
-            tools=self._tool_results.tools if self._tool_results is not None else [],
+            tools=tools,
             client_status=self._client_status,
             error=self._exception,
         )
 
+    def _get_enabled_tools(self) -> list[McpTool]:
+        result: list[McpTool] = []
+        for tool in self._mcp_tools:
+            if tool.name in self.config.exclude_tools:
+                continue
+            result.append(tool)
+        return result
+
     @property
     def mcp_tools(self) -> list[McpTool]:
         """Get the tools."""
         if self._client_status == ClientState.RUNNING:
-            return self._mcp_tools
+            return self._get_enabled_tools()
         return []
 
     def session(

dive_mcp_host/httpd/conf/mcp_servers.py

@@ -27,12 +27,21 @@ class MCPServerConfig(BaseModel):
     ) = "stdio"
     enabled: bool = True
     command: str | None = None
-    args: list[str] | None = None
-    env: dict[str, str] | None = None
+    args: list[str] | None = Field(default_factory=list)
+    env: dict[str, str] | None = Field(default_factory=dict)
     url: str | None = None
-    headers: dict[str, SecretStr] | None = None
     extra_data: dict[str, Any] | None = Field(default=None, alias="extraData")
     proxy: ProxyUrl | None = None
+    headers: dict[str, SecretStr] | None = Field(default_factory=dict)
+    exclude_tools: list[str] = Field(default_factory=list)
+
+    def model_post_init(self, _: Any) -> None:
+        """Post-initialization hook."""
+        if self.transport in ["sse", "websocket"]:
+            if self.url is None:
+                raise ValueError("url is required for sse and websocket transport")
+        elif self.transport == "stdio" and self.command is None:
+            raise ValueError("command is required for stdio transport")
 
     @field_serializer("headers", when_used="json")
     def dump_headers(self, v: dict[str, SecretStr] | None) -> dict[str, str] | None:
@@ -43,7 +52,9 @@ def dump_headers(self, v: dict[str, SecretStr] | None) -> dict[str, str] | None:
 class Config(BaseModel):
     """Model of mcp_config.json."""
 
-    mcp_servers: dict[str, MCPServerConfig] = Field(alias="mcpServers")
+    mcp_servers: dict[str, MCPServerConfig] = Field(
+        alias="mcpServers", default_factory=dict
+    )
 
 
 type McpServerConfigCallback = Callable[[Config], Config | Coroutine[Any, Any, Config]]

dive_mcp_host/httpd/routers/config.py

@@ -3,14 +3,13 @@
 from fastapi import APIRouter, Depends, Request
 from pydantic import BaseModel, Field
 
-from dive_mcp_host.httpd.conf.mcp_servers import Config
+from dive_mcp_host.httpd.conf.mcp_servers import Config as McpServers
 from dive_mcp_host.httpd.dependencies import get_app
 from dive_mcp_host.httpd.server import DiveHostAPI
 
 from .models import (
     EmbedConfig,
     McpServerError,
-    McpServers,
     ModelFullConfigs,
     ModelInterfaceDefinition,
     ModelSettingsDefinition,
@@ -79,24 +78,26 @@ async def get_mcp_server(
     )
 
 
+# Frontend prefers to use this API for all MCP server config interactions.
+# Doesn't matter if they only want to change a small thing in a single MCP server.
+# Just overwrite the entire config every time.
 @config.post("/mcpserver")
 async def post_mcp_server(
-    servers: McpServers,
+    new_config: McpServers,
     app: DiveHostAPI = Depends(get_app),
     force: bool = False,
 ) -> SaveConfigResult:
     """Save MCP server configurations.
 
     Args:
-        servers (McpServers): The server configurations to save.
+        new_config (McpServers): The server configurations to save.
         app (DiveHostAPI): The DiveHostAPI instance.
         force (bool): If True, reload all mcp servers even if they are not changed.
 
     Returns:
         SaveConfigResult: Result of the save operation with any errors.
     """
     # Update conifg
-    new_config = Config.model_validate(servers.model_dump(by_alias=True))
     if not await app.mcp_server_config_manager.update_all_configs(new_config):
         raise ValueError("Failed to update MCP server configurations")
 

dive_mcp_host/httpd/routers/models.py

@@ -1,9 +1,8 @@
 from enum import StrEnum
-from typing import Annotated, Any, Literal, Self, TypeVar
+from typing import Any, Literal, Self, TypeVar
 
 from pydantic import (
     BaseModel,
-    BeforeValidator,
     ConfigDict,
     Field,
     RootModel,
@@ -29,45 +28,6 @@ class ResultResponse(BaseModel):
     message: str | None = None
 
 
-Transport = Literal["stdio", "sse", "websocket", "streamable"]
-
-
-class McpServerConfig(BaseModel):
-    """MCP server configuration with transport and connection settings."""
-
-    transport: Annotated[
-        Transport, BeforeValidator(lambda v: "stdio" if v == "command" else v)
-    ]
-    enabled: bool | None
-    command: str | None = None
-    args: list[str] | None = Field(default_factory=list)
-    env: dict[str, str] | None = Field(default_factory=dict)
-    url: str | None = None
-    headers: dict[str, SecretStr] | None = Field(default_factory=dict)
-    extra_data: dict[str, Any] | None = Field(default=None, alias="extraData")
-
-    def model_post_init(self, _: Any) -> None:
-        """Post-initialization hook."""
-        if self.transport in ["sse", "websocket"]:
-            if self.url is None:
-                raise ValueError("url is required for sse and websocket transport")
-        elif self.transport == "stdio" and self.command is None:
-            raise ValueError("command is required for stdio transport")
-
-    @field_serializer("headers", when_used="json")
-    def dump_api_key(self, v: dict[str, SecretStr] | None) -> dict[str, str] | None:
-        """Serialize the api_key field to plain text."""
-        return {k: v.get_secret_value() for k, v in v.items()} if v else None
-
-
-class McpServers(BaseModel):
-    """Collection of MCP server configurations."""
-
-    mcp_servers: dict[str, McpServerConfig] = Field(
-        alias="mcpServers", default_factory=dict
-    )
-
-
 class McpServerError(BaseModel):
     """Represents an error from an MCP server."""
 
@@ -127,6 +87,7 @@ class SimpleToolInfo(BaseModel):
 
     name: str
     description: str
+    enabled: bool = True
 
 
 class McpTool(BaseModel):

dive_mcp_host/httpd/routers/tools.py

@@ -5,6 +5,7 @@
 from pydantic import ValidationError
 
 from dive_mcp_host.host.tools.model_types import ClientState
+from dive_mcp_host.httpd.conf.mcp_servers import Config
 from dive_mcp_host.httpd.dependencies import get_app
 from dive_mcp_host.httpd.routers.models import (
     McpTool,
@@ -42,7 +43,7 @@ async def initialized(
 
 
 @tools.get("/")
-async def list_tools(
+async def list_tools(  # noqa: PLR0912, C901
     app: DiveHostAPI = Depends(get_app),
 ) -> ToolsResult:
     """Lists all available MCP tools.
@@ -54,16 +55,20 @@ async def list_tools(
 
     # get full list of servers from config
     if (config := await app.mcp_server_config_manager.get_current_config()) is not None:
-        all_servers = set(config.mcp_servers.keys())
+        all_server_configs = config
     else:
-        all_servers = set()
+        all_server_configs = Config()
 
     # get tools from dive host
     for server_name, server_info in app.dive_host["default"].mcp_server_info.items():
         result[server_name] = McpTool(
             name=server_name,
             tools=[
-                SimpleToolInfo(name=tool.name, description=tool.description or "")
+                SimpleToolInfo(
+                    name=tool.name,
+                    description=tool.description or "",
+                    enabled=tool.enable,
+                )
                 for tool in server_info.tools
             ],
             description="",
@@ -73,7 +78,7 @@ async def list_tools(
         )
     logger.debug("active mcp servers: %s", result.keys())
     # find missing servers
-    missing_servers = all_servers - set(result.keys())
+    missing_servers = set(all_server_configs.mcp_servers.keys()) - set(result.keys())
     logger.debug("disabled mcp servers: %s", missing_servers)
 
     # get missing from local cache
@@ -90,6 +95,15 @@ async def list_tools(
         for server_name in missing_servers:
             if server_info := cached_tools.root.get(server_name, None):
                 server_info.enabled = False
+
+                # Sync tool 'enabled' with 'exclude_tools'
+                if mcp_config := all_server_configs.mcp_servers.get(server_name):
+                    for tool in server_info.tools:
+                        if tool.name in mcp_config.exclude_tools:
+                            tool.enabled = False
+                        else:
+                            tool.enabled = True
+
                 result[server_name] = server_info
             else:
                 logger.warning("Server %s not found in cached tools", server_name)

dive_mcp_host/httpd/server.py

@@ -235,6 +235,7 @@ async def load_host_config(self) -> HostConfig:
                 transport=server_config.transport or "stdio",
                 headers=server_config.headers or {},
                 proxy=server_config.proxy or None,
+                exclude_tools=server_config.exclude_tools,
             )
 
         logger.debug("got %s mcp servers in config", len(mcp_servers))