refactor(mcp): simplify MCP config & fix timeout (#7820)

Co-authored-by: ducphamle2 <ducphamle212@gmail.com>
Co-authored-by: trungbach <trunga2k29@gmail.com>
Co-authored-by: quangdz1704 <Ntq.1704@gmail.com>
Co-authored-by: Duc Pham <44611780+ducphamle2@users.noreply.github.com>
Co-authored-by: openhands <openhands@all-hands.dev>
This commit is contained in:
Xingyao Wang
2025-04-15 23:04:21 -04:00
committed by GitHub
parent 7e14a512e0
commit 07e400b73d
7 changed files with 148 additions and 46 deletions

View File

@@ -1,3 +1,4 @@
import asyncio
from contextlib import AsyncExitStack
from typing import Dict, List, Optional
@@ -36,17 +37,29 @@ class MCPClient(BaseModel):
await self.disconnect()
try:
streams_context = sse_client(
url=server_url,
)
streams = await self.exit_stack.enter_async_context(streams_context)
self.session = await self.exit_stack.enter_async_context(
ClientSession(*streams)
)
# Use asyncio.wait_for to enforce the timeout
async def connect_with_timeout():
streams_context = sse_client(
url=server_url,
timeout=timeout, # Pass the timeout to sse_client
)
streams = await self.exit_stack.enter_async_context(streams_context)
self.session = await self.exit_stack.enter_async_context(
ClientSession(*streams)
)
await self._initialize_and_list_tools()
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:

View File

@@ -38,12 +38,12 @@ def convert_mcp_clients_to_tools(mcp_clients: list[MCPClient] | None) -> list[di
async def create_mcp_clients(
sse_mcp_server: list[str],
mcp_servers: list[str],
) -> list[MCPClient]:
mcp_clients: list[MCPClient] = []
# Initialize SSE connections
if sse_mcp_server:
for server_url in sse_mcp_server:
if mcp_servers:
for server_url in mcp_servers:
logger.info(
f'Initializing MCP agent for {server_url} with SSE connection...'
)
@@ -78,11 +78,11 @@ async def fetch_mcp_tools_from_config(mcp_config: MCPConfig) -> list[dict]:
try:
logger.debug(f'Creating MCP clients with config: {mcp_config}')
mcp_clients = await create_mcp_clients(
mcp_config.sse.mcp_servers,
mcp_config.mcp_servers,
)
if not mcp_clients:
logger.warning('No MCP clients were successfully connected')
logger.debug('No MCP clients were successfully connected')
return []
mcp_tools = convert_mcp_clients_to_tools(mcp_clients)