fix(mcp_manager): implement concurrent server connections with timeout handling

This commit is contained in:
Boody
2026-07-09 02:45:40 +03:00
committed by RaresKeY
parent df2fad2881
commit 5971b091db
3 changed files with 205 additions and 10 deletions
+26 -10
View File
@@ -9,7 +9,9 @@ import json
import logging import logging
import os import os
import re import re
import asyncio
from typing import Any, Dict, List, Optional, Set, Tuple from typing import Any, Dict, List, Optional, Set, Tuple
from src.database import McpServer, SessionLocal
from src.runtime_paths import get_app_root from src.runtime_paths import get_app_root
@@ -409,17 +411,29 @@ class McpManager:
for sid in ids: for sid in ids:
await self.disconnect_server(sid) 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() db = SessionLocal()
try: try:
servers = db.query(McpServer).filter(McpServer.is_enabled == True).all() servers = db.query(McpServer).filter(McpServer.is_enabled == True).all()
for srv in servers:
args = json.loads(srv.args) if srv.args else [] tasks = [
env = json.loads(srv.env) if srv.env else {} asyncio.create_task(self._connect_with_timeout(srv))
await self.connect_server( 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, server_id=srv.id,
name=srv.name, name=srv.name,
transport=srv.transport, transport=srv.transport,
@@ -427,9 +441,11 @@ class McpManager:
args=args, args=args,
env=env, env=env,
url=srv.url, url=srv.url,
) ),
finally: timeout=20,
db.close() )
except asyncio.TimeoutError:
logger.warning("Timed out connecting to %s", srv.name)
async def call_tool(self, qualified_name: str, arguments: Dict) -> Dict: async def call_tool(self, qualified_name: str, arguments: Dict) -> Dict:
"""Call an MCP tool by its qualified name (mcp__{server_id}__{tool_name}). """Call an MCP tool by its qualified name (mcp__{server_id}__{tool_name}).
+1
View File
@@ -26,6 +26,7 @@ try:
import sqlalchemy # noqa: F401 import sqlalchemy # noqa: F401
import sqlalchemy.orm # noqa: F401 import sqlalchemy.orm # noqa: F401
import core.database # noqa: F401 import core.database # noqa: F401
import src.database
except ImportError: except ImportError:
pass # not installed - the stubs below will handle it pass # not installed - the stubs below will handle it
+178
View File
@@ -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