(refactor): Update MCP Client to use FastMCP (OpenHands#8931)

Author: malhotra5

openhands/mcp/client.py

@@ -1,13 +1,12 @@
-import asyncio
-import datetime
-from contextlib import AsyncExitStack
 from typing import Optional
 
-from mcp import ClientSession
-from mcp.client.sse import sse_client
-from mcp.client.streamable_http import streamablehttp_client
+from fastmcp import Client
+from fastmcp.client.transports import SSETransport, StreamableHttpTransport
+from mcp import McpError
+from mcp.types import CallToolResult
 from pydantic import BaseModel, Field
 
+from openhands.core.config.mcp_config import MCPSHTTPServerConfig, MCPSSEServerConfig
 from openhands.core.logger import openhands_logger as logger
 from openhands.mcp.tool import MCPClientTool
 
@@ -17,198 +16,95 @@ class MCPClient(BaseModel):
     A collection of tools that connects to an MCP server and manages available tools through the Model Context Protocol.
     """
 
-    session: Optional[ClientSession] = None
-    exit_stack: AsyncExitStack = AsyncExitStack()
+    client: Optional[Client] = None
     description: str = 'MCP client tools for server interaction'
     tools: list[MCPClientTool] = Field(default_factory=list)
     tool_map: dict[str, MCPClientTool] = Field(default_factory=dict)
 
     class Config:
         arbitrary_types_allowed = True
 
-    async def connect_sse(
-        self,
-        server_url: str,
-        api_key: str | None = None,
-        conversation_id: str | None = None,
-        timeout: float = 30.0,
-    ) -> None:
-        """Connect to an MCP server using SSE transport.
-
-        Args:
-            server_url: The URL of the SSE server to connect to.
-            timeout: Connection timeout in seconds. Default is 30 seconds.
-        """
-        if not server_url:
-            raise ValueError('Server URL is required.')
-        if self.session:
-            await self.disconnect()
-
-        try:
-            # Use asyncio.wait_for to enforce the timeout
-            async def connect_with_timeout():
-                headers = (
-                    {
-                        'Authorization': f'Bearer {api_key}',
-                        's': api_key,  # We need this for action execution server's MCP Router
-                        'X-Session-API-Key': api_key,  # We need this for Remote Runtime
-                    }
-                    if api_key
-                    else {}
-                )
-
-                if conversation_id:
-                    headers['X-OpenHands-Conversation-ID'] = conversation_id
-
-                # Convert float timeout to datetime.timedelta for consistency
-                timeout_delta = datetime.timedelta(seconds=timeout)
-
-                streams_context = sse_client(
-                    url=server_url,
-                    headers=headers if headers else None,
-                    timeout=timeout,
-                )
-                streams = await self.exit_stack.enter_async_context(streams_context)
-                # For SSE client, we only get read_stream and write_stream (2 values)
-                read_stream, write_stream = streams
-                self.session = await self.exit_stack.enter_async_context(
-                    ClientSession(
-                        read_stream, write_stream, read_timeout_seconds=timeout_delta
-                    )
-                )
-                await self._initialize_and_list_tools()
-
-            # Apply timeout to the entire connection process
-            await asyncio.wait_for(connect_with_timeout(), timeout=timeout)
-        except asyncio.TimeoutError:
-            logger.error(
-                f'Connection to {server_url} timed out after {timeout} seconds'
-            )
-            await self.disconnect()  # Clean up resources
-            raise  # Re-raise the TimeoutError
-        except Exception as e:
-            logger.error(f'Error connecting to {server_url}: {str(e)}')
-            await self.disconnect()  # Clean up resources
-            raise
-
     async def _initialize_and_list_tools(self) -> None:
         """Initialize session and populate tool map."""
-        if not self.session:
+        if not self.client:
             raise RuntimeError('Session not initialized.')
 
-        await self.session.initialize()
-        response = await self.session.list_tools()
+        async with self.client:
+            tools = await self.client.list_tools()
 
         # Clear existing tools
         self.tools = []
 
         # Create proper tool objects for each server tool
-        for tool in response.tools:
+        for tool in tools:
             server_tool = MCPClientTool(
                 name=tool.name,
                 description=tool.description,
                 inputSchema=tool.inputSchema,
-                session=self.session,
+                session=self.client,
             )
             self.tool_map[tool.name] = server_tool
             self.tools.append(server_tool)
 
-        logger.info(
-            f'Connected to server with tools: {[tool.name for tool in response.tools]}'
-        )
+        logger.info(f'Connected to server with tools: {[tool.name for tool in tools]}')
 
-    async def call_tool(self, tool_name: str, args: dict):
-        """Call a tool on the MCP server."""
-        if tool_name not in self.tool_map:
-            raise ValueError(f'Tool {tool_name} not found.')
-        # The MCPClientTool is primarily for metadata; use the session to call the actual tool.
-        if not self.session:
-            raise RuntimeError('Client session is not available.')
-        return await self.session.call_tool(name=tool_name, arguments=args)
-
-    async def connect_shttp(
+    async def connect_http(
         self,
-        server_url: str,
-        api_key: str | None = None,
+        server: MCPSSEServerConfig | MCPSHTTPServerConfig,
         conversation_id: str | None = None,
         timeout: float = 30.0,
-    ) -> None:
-        """Connect to an MCP server using StreamableHTTP transport.
-
-        Args:
-            server_url: The URL of the StreamableHTTP server to connect to.
-            api_key: Optional API key for authentication.
-            conversation_id: Optional conversation ID for session tracking.
-            timeout: Connection timeout in seconds. Default is 30 seconds.
-        """
+    ):
+        """Connect to MCP server using SHTTP or SSE transport"""
+        server_url = server.url
+        api_key = server.api_key
+
         if not server_url:
             raise ValueError('Server URL is required.')
-        if self.session:
-            await self.disconnect()
 
         try:
-            # Use asyncio.wait_for to enforce the timeout
-            async def connect_with_timeout():
-                headers = (
-                    {
-                        'Authorization': f'Bearer {api_key}',
-                        's': api_key,  # We need this for action execution server's MCP Router
-                        'X-Session-API-Key': api_key,  # We need this for Remote Runtime
-                    }
-                    if api_key
-                    else {}
-                )
-
-                if conversation_id:
-                    headers['X-OpenHands-Conversation-ID'] = conversation_id
+            headers = (
+                {
+                    'Authorization': f'Bearer {api_key}',
+                    's': api_key,  # We need this for action execution server's MCP Router
+                    'X-Session-API-Key': api_key,  # We need this for Remote Runtime
+                }
+                if api_key
+                else {}
+            )
 
-                # Convert float timeout to datetime.timedelta
-                timeout_delta = datetime.timedelta(seconds=timeout)
-                sse_read_timeout_delta = datetime.timedelta(
-                    seconds=timeout * 10
-                )  # 10x longer for read timeout
+            if conversation_id:
+                headers['X-OpenHands-Conversation-ID'] = conversation_id
 
-                streams_context = streamablehttp_client(
+            # Instantiate custom transports due to custom headers
+            if isinstance(server, MCPSHTTPServerConfig):
+                transport = StreamableHttpTransport(
                     url=server_url,
                     headers=headers if headers else None,
-                    timeout=timeout_delta,
-                    sse_read_timeout=sse_read_timeout_delta,
                 )
-                streams = await self.exit_stack.enter_async_context(streams_context)
-                # For StreamableHTTP client, we get read_stream, write_stream, and get_session_id (3 values)
-                read_stream, write_stream, _ = streams
-                self.session = await self.exit_stack.enter_async_context(
-                    ClientSession(
-                        read_stream, write_stream, read_timeout_seconds=timeout_delta
-                    )
+            else:
+                transport = SSETransport(
+                    url=server_url,
+                    headers=headers if headers else None,
                 )
-                await self._initialize_and_list_tools()
 
-            # Apply timeout to the entire connection process
-            await asyncio.wait_for(connect_with_timeout(), timeout=timeout)
-        except asyncio.TimeoutError:
-            logger.error(
-                f'Connection to {server_url} timed out after {timeout} seconds'
-            )
-            await self.disconnect()  # Clean up resources
-            raise  # Re-raise the TimeoutError
+            self.client = Client(transport, timeout=timeout)
+
+            await self._initialize_and_list_tools()
+        except McpError as e:
+            logger.error(f'McpError connecting to {server_url}: {e}')
+            raise  # Re-raise the error
+
         except Exception as e:
-            logger.error(f'Error connecting to {server_url}: {str(e)}')
-            await self.disconnect()  # Clean up resources
+            logger.error(f'Error connecting to {server_url}: {e}')
             raise
 
-    async def disconnect(self) -> None:
-        """Disconnect from the MCP server and clean up resources."""
-        if self.session:
-            try:
-                # Close the session first
-                if hasattr(self.session, 'close'):
-                    await self.session.close()
-                # Then close the exit stack
-                await self.exit_stack.aclose()
-            except Exception as e:
-                logger.error(f'Error during disconnect: {str(e)}')
-            finally:
-                self.session = None
-                self.tools = []
-                logger.info('Disconnected from MCP server')
+    async def call_tool(self, tool_name: str, args: dict) -> CallToolResult:
+        """Call a tool on the MCP server."""
+        if tool_name not in self.tool_map:
+            raise ValueError(f'Tool {tool_name} not found.')
+        # The MCPClientTool is primarily for metadata; use the session to call the actual tool.
+        if not self.client:
+            raise RuntimeError('Client session is not available.')
+
+        async with self.client:
+            return await self.client.call_tool_mcp(name=tool_name, arguments=args)

openhands/mcp/utils.py

@@ -72,38 +72,22 @@ async def create_mcp_clients(
     mcp_clients = []
 
     for server in servers:
-        is_sse = isinstance(server, MCPSSEServerConfig)
-        connection_type = 'SSE' if is_sse else 'SHTTP'
+        is_shttp = isinstance(server, MCPSHTTPServerConfig)
+        connection_type = 'SHTTP' if is_shttp else 'SSE'
         logger.info(
             f'Initializing MCP agent for {server} with {connection_type} connection...'
         )
         client = MCPClient()
 
         try:
-            if is_sse:
-                await client.connect_sse(
-                    server.url,
-                    api_key=server.api_key,
-                    conversation_id=conversation_id,
-                )
-            else:
-                await client.connect_shttp(
-                    server.url,
-                    api_key=server.api_key,
-                    conversation_id=conversation_id,
-                )
+            await client.connect_http(server, conversation_id=conversation_id)
 
             # Only add the client to the list after a successful connection
             mcp_clients.append(client)
 
         except Exception as e:
             logger.error(f'Failed to connect to {server}: {str(e)}', exc_info=True)
-            try:
-                await client.disconnect()
-            except Exception as disconnect_error:
-                logger.error(
-                    f'Error during disconnect after failed connection: {str(disconnect_error)}'
-                )
+
     return mcp_clients
 
 
@@ -143,13 +127,6 @@ async def fetch_mcp_tools_from_config(
         # Convert tools to the format expected by the agent
         mcp_tools = convert_mcp_clients_to_tools(mcp_clients)
 
-        # Always disconnect clients to clean up resources
-        for mcp_client in mcp_clients:
-            try:
-                await mcp_client.disconnect()
-            except Exception as disconnect_error:
-                logger.error(f'Error disconnecting MCP client: {str(disconnect_error)}')
-
     except Exception as e:
         logger.error(f'Error fetching MCP tools: {str(e)}')
         return []

openhands/runtime/impl/action_execution/action_execution_client.py

@@ -471,11 +471,6 @@ async def call_tool_mcp(self, action: MCPAction) -> Observation:
         # Call the tool and return the result
         # No need for try/finally since disconnect() is now just resetting state
         result = await call_tool_mcp_handler(mcp_clients, action)
-
-        # Reset client state (no active connections to worry about)
-        for client in mcp_clients:
-            await client.disconnect()
-
         return result
 
     def close(self) -> None: