fix(mcp_manager): remove timeout from MCP connection attempts and handle registration cleanup

app.py

@@ -1048,7 +1048,7 @@ async def _startup_event():
         except BaseException as e:
             logger.warning(f"Built-in MCP registration failed (non-critical): {type(e).__name__}: {e}")
         try:
-            await asyncio.wait_for(mcp_manager.connect_all_enabled(), timeout=20)
+            await mcp_manager.connect_all_enabled()
         except asyncio.TimeoutError:
             logger.warning("User MCP startup timed out (non-critical)")
         except BaseException as e:

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
 
@@ -192,57 +194,65 @@ class McpManager:
             )
 
             stack = AsyncExitStack()
+            registered = False
+
             try:
                 transport = await stack.enter_async_context(stdio_client(server_params))
                 read_stream, write_stream = transport
                 session = await stack.enter_async_context(ClientSession(read_stream, write_stream))
 
                 await session.initialize()
-
-                # Discover tools
                 tools_result = await session.list_tools()
-            except Exception:
-                await stack.aclose()
-                raise
-            tools = []
-            for tool in tools_result.tools:
-                tools.append({
-                    "name": tool.name,
-                    "description": tool.description or "",
-                    "input_schema": tool.inputSchema if hasattr(tool, 'inputSchema') else {},
-                    # MCP tool annotations (readOnlyHint / destructiveHint) drive
-                    # plan-mode read-only gating. Absent on many servers, so we
-                    # fall back to a name heuristic in mcp_tool_is_readonly().
-                    "annotations": getattr(tool, 'annotations', None),
-                })
 
-            self._sessions[server_id] = session
-            self._stacks[server_id] = stack
-            self._tools[server_id] = tools
-            # Extract identity hints from env vars (e.g. email address, API name)
-            # so tool descriptions can distinguish between multiple instances of
-            # the same MCP server (e.g. two email accounts).
-            identity_hints = []
-            for k, v in (env or {}).items():
-                k_lower = k.lower()
-                if any(x in k_lower for x in ['email_address', 'account', 'user', 'username']):
-                    identity_hints.append(v)
-            identity = ", ".join(identity_hints) if identity_hints else ""
+                tools = []
+                for tool in tools_result.tools:
+                    tools.append({
+                        "name": tool.name,
+                        "description": tool.description or "",
+                        "input_schema": tool.inputSchema if hasattr(tool, "inputSchema") else {},
+                        # MCP tool annotations (readOnlyHint / destructiveHint) drive
+                        # plan-mode read-only gating. Absent on many servers, so we
+                        # fall back to a name heuristic in mcp_tool_is_readonly().
+                        "annotations": getattr(tool, "annotations", None),
+                    })
 
-            self._connections[server_id] = {
-                "status": "connected",
-                "name": name,
-                "transport": "stdio",
-                "tool_count": len(tools),
-                "identity": identity,
-            }
+                # Extract identity hints from env vars (e.g. email address, API name)
+                # so tool descriptions can distinguish between multiple instances of
+                # the same MCP server (e.g. two email accounts).
+                identity_hints = []
+                for k, v in (env or {}).items():
+                    k_lower = k.lower()
+                    if any(x in k_lower for x in ["email_address", "account", "user", "username"]):
+                        identity_hints.append(v)
+                identity = ", ".join(identity_hints) if identity_hints else ""
+
+                self._sessions[server_id] = session
+                self._stacks[server_id] = stack
+                self._tools[server_id] = tools
+                self._connections[server_id] = {
+                    "status": "connected",
+                    "name": name,
+                    "transport": "stdio",
+                    "tool_count": len(tools),
+                    "identity": identity,
+                }
+
+                registered = True
+
+            finally:
+                if not registered:
+                    await stack.aclose()
 
             logger.info(f"MCP server connected: {name} ({server_id}) - {len(tools)} tools via stdio")
             return True
 
         except ImportError:
             logger.warning("MCP package not installed. Install with: pip install mcp")
-            self._connections[server_id] = {"status": "error", "error": "mcp package not installed", "name": name}
+            self._connections[server_id] = {
+                "status": "error",
+                "error": "mcp package not installed",
+                "name": name,
+            }
             return False
 
     async def _connect_sse(self, server_id: str, name: str, url: str) -> bool:
@@ -253,42 +263,46 @@ class McpManager:
             from contextlib import AsyncExitStack
 
             stack = AsyncExitStack()
+            registered = False
+
             try:
                 transport = await stack.enter_async_context(sse_client(url))
                 read_stream, write_stream = transport
                 session = await stack.enter_async_context(ClientSession(read_stream, write_stream))
 
                 await session.initialize()
-
-                # Discover tools
                 tools_result = await session.list_tools()
-            except Exception:
-                await stack.aclose()
-                raise
-            tools = []
-            for tool in tools_result.tools:
-                tools.append({
-                    "name": tool.name,
-                    "description": tool.description or "",
-                    "input_schema": tool.inputSchema if hasattr(tool, 'inputSchema') else {},
-                    # MCP tool annotations (readOnlyHint / destructiveHint) drive
-                    # plan-mode read-only gating. Absent on many servers, so we
-                    # fall back to a name heuristic in mcp_tool_is_readonly().
-                    "annotations": getattr(tool, 'annotations', None),
-                })
 
-            self._sessions[server_id] = session
-            self._stacks[server_id] = stack
-            self._tools[server_id] = tools
-            self._connections[server_id] = {
-                "status": "connected",
-                "name": name,
-                "transport": "sse",
-                "tool_count": len(tools),
-            }
+                tools = []
+                for tool in tools_result.tools:
+                    tools.append({
+                        "name": tool.name,
+                        "description": tool.description or "",
+                        "input_schema": tool.inputSchema if hasattr(tool, 'inputSchema') else {},
+                        # MCP tool annotations (readOnlyHint / destructiveHint) drive
+                        # plan-mode read-only gating. Absent on many servers, so we
+                        # fall back to a name heuristic in mcp_tool_is_readonly().
+                        "annotations": getattr(tool, 'annotations', None),
+                    })
 
-            logger.info(f"MCP server connected: {name} ({server_id}) - {len(tools)} tools via SSE")
-            return True
+                self._sessions[server_id] = session
+                self._stacks[server_id] = stack
+                self._tools[server_id] = tools
+                self._connections[server_id] = {
+                    "status": "connected",
+                    "name": name,
+                    "transport": "sse",
+                    "tool_count": len(tools),
+                }
+
+                registered = True
+
+                logger.info(f"MCP server connected: {name} ({server_id}) - {len(tools)} tools via SSE")
+                return True
+
+            finally:
+                if not registered:
+                    await stack.aclose()
 
         except ImportError:
             logger.warning("MCP package not installed. Install with: pip install mcp")
@@ -409,17 +423,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 +453,16 @@ 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)
+            self._connections[srv.id] = {
+                "status": "timeout",
+                "error": f"Timed out after 20 seconds",
+                "name": 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}).

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
 

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