refactor MCP config

Author: xingyaoww

openhands/core/config/mcp_config.py

@@ -31,50 +31,43 @@ def validate_servers(self) -> None:
                 raise ValueError(f'Invalid URL {url}: {str(e)}')
 
 
+class MCPStdioConfigEntry(BaseModel):
+    """Configuration for a single MCP stdio entry.
+
+    Attributes:
+        command: The command to run.
+        args: List of arguments for the command.
+        env: Dictionary of environment variables.
+    """
+
+    command: str
+    args: list[str] = Field(default_factory=list)
+    env: dict[str, str] = Field(default_factory=dict)
+
+    model_config = {'extra': 'forbid'}
+
+
 class MCPStdioConfig(BaseModel):
     """Configuration for MCP stdio settings.
 
     Attributes:
-        commands: List of commands to run.
-        args: List of arguments for each command.
-        envs: List of environment variable tuples for each command.
+        tools: Dictionary of tool configurations, where keys are tool names.
     """
 
-    commands: List[str] = Field(default_factory=list)
-    args: List[List[str]] = Field(default_factory=list)
-    envs: List[List[tuple[str, str]]] = Field(default_factory=list)
+    tools: dict[str, MCPStdioConfigEntry] = Field(default_factory=dict)
 
     model_config = {'extra': 'forbid'}
 
     def validate_stdio(self) -> None:
-        """Validate that commands, args, and envs are properly configured."""
-
-        # Check if number of commands matches number of args lists
-        if len(self.commands) != len(self.args):
-            raise ValueError(
-                f'Number of commands ({len(self.commands)}) does not match '
-                f'number of args lists ({len(self.args)})'
-            )
-
-        # Check if number of commands matches number of envs lists
-        if len(self.commands) != len(self.envs):
-            raise ValueError(
-                f'Number of commands ({len(self.commands)}) does not match '
-                f'number of envs lists ({len(self.envs)})'
-            )
-
-        # Validate each environment variable tuple
-        for i, env_list in enumerate(self.envs):
-            for j, env_tuple in enumerate(env_list):
-                if not isinstance(env_tuple, tuple) or len(env_tuple) != 2:
-                    raise ValueError(
-                        f'Environment variable at index {j} for command {i} must be a tuple of (key, value)'
-                    )
-                key, value = env_tuple
-                if not isinstance(key, str) or not isinstance(value, str):
-                    raise ValueError(
-                        f'Environment variable key and value at index {j} for command {i} must be strings'
-                    )
+        """Validate that tools are properly configured."""
+        # Tool names validation
+        for tool_name in self.tools:
+            if not tool_name.strip():
+                raise ValueError('Tool names cannot be empty')
+            if not tool_name.replace('-', '').isalnum():
+                raise ValueError(
+                    f'Invalid tool name: {tool_name}. Tool names must be alphanumeric (hyphens allowed)'
+                )
 
 
 class MCPConfig(BaseModel):
@@ -105,11 +98,11 @@ def from_toml_section(cls, data: dict) -> dict[str, 'MCPConfig']:
 
         try:
             # Create SSE config if present
-            sse_config = MCPSSEConfig(**data.get('mcp-sse', {}))
+            sse_config = MCPSSEConfig.model_validate(data.get('mcp-sse', {}))
             sse_config.validate_servers()
 
             # Create stdio config if present
-            stdio_config = MCPStdioConfig(**data.get('mcp-stdio', {}))
+            stdio_config = MCPStdioConfig.model_validate(data.get('mcp-stdio', {}))
             stdio_config.validate_stdio()
 
             # Create the main MCP config

openhands/mcp/utils.py

@@ -1,6 +1,4 @@
-from typing import List
-
-from openhands.core.config.mcp_config import MCPConfig
+from openhands.core.config.mcp_config import MCPConfig, MCPStdioConfigEntry
 from openhands.core.logger import openhands_logger as logger
 from openhands.mcp.client import MCPClient
 
@@ -35,12 +33,10 @@ def convert_mcp_clients_to_tools(mcp_clients: list[MCPClient] | None) -> list[di
 
 
 async def create_mcp_clients(
-    sse_mcp_server: List[str],
-    commands: List[str],
-    args: List[List[str]],
-    envs: List[List[tuple[str, str]]],
-) -> List[MCPClient]:
-    mcp_clients: List[MCPClient] = []
+    sse_mcp_server: list[str],
+    stdio_mcp_tool_configs: dict[str, MCPStdioConfigEntry],
+) -> list[MCPClient]:
+    mcp_clients: list[MCPClient] = []
     # Initialize SSE connections
     if sse_mcp_server:
         for server_url in sse_mcp_server:
@@ -65,21 +61,24 @@ async def create_mcp_clients(
                     )
 
     # Initialize stdio connections
-    if commands:
-        for i, (command, command_args, command_envs) in enumerate(
-            zip(commands, args, envs)
-        ):
+    if stdio_mcp_tool_configs:
+        for name, tool in stdio_mcp_tool_configs.items():
             logger.info(
-                f'Initializing MCP agent for {command} with stdio connection...'
+                f'Initializing MCP tool [{name}] for [{tool.command}] with stdio connection...'
             )
-
             client = MCPClient()
             try:
-                await client.connect_stdio(command, command_args, command_envs)
+                await client.connect_stdio(
+                    tool.command,
+                    tool.args,
+                    [(k, v) for k, v in tool.env.items()],
+                )
                 mcp_clients.append(client)
-                logger.info(f'Connected to MCP server via stdio with command {command}')
+                logger.info(
+                    f'Connected to MCP server via stdio with command {tool.command}'
+                )
             except Exception as e:
-                logger.error(f'Failed to connect with command {command}: {str(e)}')
+                logger.error(f'Failed to connect with command {tool.command}: {str(e)}')
                 # Don't raise the exception, just log it and continue
                 # Make sure to disconnect the client to clean up resources
                 try:
@@ -88,7 +87,6 @@ async def create_mcp_clients(
                     logger.error(
                         f'Error during disconnect after failed connection: {str(disconnect_error)}'
                     )
-
     return mcp_clients
 
 
@@ -103,11 +101,10 @@ async def fetch_mcp_tools_from_config(mcp_config: MCPConfig) -> list[dict]:
     mcp_tools = []
 
     try:
+        logger.debug(f'Creating MCP clients with config: {mcp_config}')
         mcp_clients = await create_mcp_clients(
             mcp_config.sse.mcp_servers,
-            mcp_config.stdio.commands,
-            mcp_config.stdio.args,
-            mcp_config.stdio.envs,
+            mcp_config.stdio.tools,
         )
 
         if not mcp_clients:
@@ -126,4 +123,5 @@ async def fetch_mcp_tools_from_config(mcp_config: MCPConfig) -> list[dict]:
             except Exception as disconnect_error:
                 logger.error(f'Error disconnecting MCP client: {str(disconnect_error)}')
 
+    logger.debug(f'MCP tools: {mcp_tools}')
     return mcp_tools

openhands/runtime/action_execution_server.py

@@ -17,7 +17,6 @@
 import traceback
 from contextlib import asynccontextmanager
 from pathlib import Path
-from typing import Optional
 from zipfile import ZipFile
 
 from fastapi import Depends, FastAPI, HTTPException, Request, UploadFile
@@ -33,6 +32,7 @@
 from starlette.exceptions import HTTPException as StarletteHTTPException
 from uvicorn import run
 
+from openhands.core.config.mcp_config import MCPStdioConfigEntry
 from openhands.core.exceptions import BrowserUnavailableException
 from openhands.core.logger import openhands_logger as logger
 from openhands.events.action import (
@@ -73,10 +73,8 @@
 
 class ActionRequest(BaseModel):
     action: dict
-    sse_mcp_config: Optional[list[str]] = None
-    stdio_mcp_config: Optional[
-        tuple[list[str], list[list[str]], list[list[tuple[str, str]]]]
-    ] = None
+    sse_mcp_config: list[str] | None = None
+    stdio_mcp_tool_configs: dict[str, MCPStdioConfigEntry] | None = None
 
 
 ROOT_GID = 0
@@ -192,20 +190,18 @@ def __init__(
         )
         self.memory_monitor.start_monitoring()
         self.sse_mcp_servers: list[str] = []
-        self.stdio_mcp_config: tuple[
-            list[str], list[list[str]], list[list[tuple[str, str]]]
-        ] = ([], [], [])
+        self.stdio_mcp_tool_configs: dict[str, MCPStdioConfigEntry] = {}
 
     @property
     def initial_cwd(self):
         return self._initial_cwd
 
     def process_request(self, action_request: ActionRequest):
-        # update the sse_mcp_servers and stdio_mcp_config to prepare for MCP action if needed
+        # update the sse_mcp_servers and stdio_mcp_tool_configs to prepare for MCP action if needed
         if action_request.sse_mcp_config:
             self.sse_mcp_servers = action_request.sse_mcp_config
-        if action_request.stdio_mcp_config:
-            self.stdio_mcp_config = action_request.stdio_mcp_config
+        if action_request.stdio_mcp_tool_configs:
+            self.stdio_mcp_tool_configs = action_request.stdio_mcp_tool_configs
 
     async def _init_browser_async(self):
         """Initialize the browser asynchronously."""
@@ -526,21 +522,13 @@ async def browse_interactive(self, action: BrowseInteractiveAction) -> Observati
         return await browse(action, self.browser)
 
     async def call_tool_mcp(self, action: McpAction) -> Observation:
-        commands: list[str] = []
-        args: list[list[str]] = []
-        envs: list[list[tuple[str, str]]] = []
-        if self.stdio_mcp_config:
-            commands = self.stdio_mcp_config[0]
-            if len(self.stdio_mcp_config) > 1:
-                args = self.stdio_mcp_config[1]
-            if len(self.stdio_mcp_config) > 2:
-                envs = self.stdio_mcp_config[2]
-        if not self.sse_mcp_servers and not commands:
+        if not self.sse_mcp_servers and not self.stdio_mcp_tool_configs:
             raise ValueError('No MCP servers or stdio MCP config found')
 
         logger.warning(f'SSE MCP servers: {self.sse_mcp_servers}')
         mcp_clients = await create_mcp_clients(
-            self.sse_mcp_servers, commands, args, envs
+            self.sse_mcp_servers,
+            self.stdio_mcp_tool_configs,
         )
         logger.warn(f'MCP action received: {action}')
         # Find the MCP agent that has the matching tool name