From 5971b091dbf20d65dce1a8f53ff71efb772b26b5 Mon Sep 17 00:00:00 2001 From: Boody Date: Thu, 9 Jul 2026 02:45:40 +0300 Subject: [PATCH] fix(mcp_manager): implement concurrent server connections with timeout handling --- src/mcp_manager.py | 36 +++-- tests/conftest.py | 1 + tests/test_multiple_mcp_servers_timeout.py | 178 +++++++++++++++++++++ 3 files changed, 205 insertions(+), 10 deletions(-) create mode 100644 tests/test_multiple_mcp_servers_timeout.py diff --git a/src/mcp_manager.py b/src/mcp_manager.py index 90430c78f..ae72f6aea 100644 --- a/src/mcp_manager.py +++ b/src/mcp_manager.py @@ -9,7 +9,9 @@ import json import logging import os import re +import asyncio from typing import Any, Dict, List, Optional, Set, Tuple +from src.database import McpServer, SessionLocal from src.runtime_paths import get_app_root @@ -409,17 +411,29 @@ class McpManager: for sid in ids: await self.disconnect_server(sid) - async def connect_all_enabled(self): - """Connect to all enabled MCP servers from the database.""" - from src.database import McpServer, SessionLocal + async def connect_all_enabled(self): db = SessionLocal() try: servers = db.query(McpServer).filter(McpServer.is_enabled == True).all() - for srv in servers: - args = json.loads(srv.args) if srv.args else [] - env = json.loads(srv.env) if srv.env else {} - await self.connect_server( + + tasks = [ + asyncio.create_task(self._connect_with_timeout(srv)) + for srv in servers + ] + + await asyncio.gather(*tasks) + finally: + db.close() + + + async def _connect_with_timeout(self, srv): + args = json.loads(srv.args) if srv.args else [] + env = json.loads(srv.env) if srv.env else {} + + try: + await asyncio.wait_for( + self.connect_server( server_id=srv.id, name=srv.name, transport=srv.transport, @@ -427,9 +441,11 @@ class McpManager: args=args, env=env, url=srv.url, - ) - finally: - db.close() + ), + timeout=20, + ) + except asyncio.TimeoutError: + logger.warning("Timed out connecting to %s", srv.name) async def call_tool(self, qualified_name: str, arguments: Dict) -> Dict: """Call an MCP tool by its qualified name (mcp__{server_id}__{tool_name}). diff --git a/tests/conftest.py b/tests/conftest.py index e78db01cf..2f6defa1c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -26,6 +26,7 @@ try: import sqlalchemy # noqa: F401 import sqlalchemy.orm # noqa: F401 import core.database # noqa: F401 + import src.database except ImportError: pass # not installed - the stubs below will handle it diff --git a/tests/test_multiple_mcp_servers_timeout.py b/tests/test_multiple_mcp_servers_timeout.py new file mode 100644 index 000000000..4c0986d0b --- /dev/null +++ b/tests/test_multiple_mcp_servers_timeout.py @@ -0,0 +1,178 @@ +import asyncio +import json +import time +from types import SimpleNamespace + +import pytest + + +class FakeQuery: + def __init__(self, servers): + self._servers = servers + + def filter(self, *_args, **_kwargs): + return self + + def all(self): + return self._servers + + +class FakeDB: + def __init__(self, servers): + self._servers = servers + + def query(self, *_args, **_kwargs): + return FakeQuery(self._servers) + + def close(self): + pass + + +@pytest.mark.asyncio +async def test_connect_all_enabled_runs_concurrently(monkeypatch): + from src.mcp_manager import McpManager + + manager = McpManager() + + servers = [ + SimpleNamespace( + id=1, + name="server1", + transport="stdio", + command="cmd1", + args=json.dumps([]), + env=json.dumps({}), + url=None, + ), + SimpleNamespace( + id=2, + name="server2", + transport="stdio", + command="cmd2", + args=json.dumps([]), + env=json.dumps({}), + url=None, + ), + SimpleNamespace( + id=3, + name="server3", + transport="stdio", + command="cmd3", + args=json.dumps([]), + env=json.dumps({}), + url=None, + ), + ] + + # Patch the SessionLocal used by connect_all_enabled(). + import src.mcp_manager as mcp_manager + + monkeypatch.setattr( + mcp_manager, + "SessionLocal", + lambda: FakeDB(servers), + ) + + async def fake_connect_with_timeout(_server): + await asyncio.sleep(1) + + # We're testing that connect_all_enabled launches these concurrently, + # not the implementation of connect_server(). + monkeypatch.setattr( + manager, + "_connect_with_timeout", + fake_connect_with_timeout, + ) + + start = time.perf_counter() + + await manager.connect_all_enabled() + + elapsed = time.perf_counter() - start + + # Sequential would take ~3 seconds. + # Concurrent should take about 1 second. + assert 0.9 <= elapsed < 2.0 + +@pytest.mark.asyncio +async def test_connect_all_enabled_timeout_does_not_block_other_servers(monkeypatch): + import src.mcp_manager as mcp_manager + from src.mcp_manager import McpManager + + manager = McpManager() + + servers = [ + SimpleNamespace( + id=1, + name="fast1", + transport="stdio", + command="cmd1", + args=json.dumps([]), + env=json.dumps({}), + url=None, + ), + SimpleNamespace( + id=2, + name="slow", + transport="stdio", + command="cmd2", + args=json.dumps([]), + env=json.dumps({}), + url=None, + ), + SimpleNamespace( + id=3, + name="fast2", + transport="stdio", + command="cmd3", + args=json.dumps([]), + env=json.dumps({}), + url=None, + ), + ] + + monkeypatch.setattr( + mcp_manager, + "SessionLocal", + lambda: FakeDB(servers), + ) + + completed = [] + + async def fake_connect_server(server_id, **kwargs): + if server_id == 2: + # Simulate a hung connection. + await asyncio.sleep(30) + else: + await asyncio.sleep(0.1) + completed.append(server_id) + + monkeypatch.setattr( + manager, + "connect_server", + fake_connect_server, + ) + + # + # Don't actually wait 20 seconds during the test. + # Replace asyncio.wait_for used by mcp_manager with a much shorter timeout. + # + real_wait_for = asyncio.wait_for + + async def short_wait_for(awaitable, timeout): + return await real_wait_for(awaitable, timeout=0.2) + + monkeypatch.setattr( + mcp_manager.asyncio, + "wait_for", + short_wait_for, + ) + + start = time.perf_counter() + + await manager.connect_all_enabled() + + elapsed = time.perf_counter() - start + + assert set(completed) == {1, 3} + assert elapsed < 1 \ No newline at end of file