mirror of
https://github.com/OpenHands/OpenHands.git
synced 2026-03-22 13:47:19 +08:00
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:
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user