Compare commits

...
2 Commits
4 changed files with 286 additions and 74 deletions
+1 -1
View File
@@ -1048,7 +1048,7 @@ async def _startup_event():
except BaseException as e: except BaseException as e:
logger.warning(f"Built-in MCP registration failed (non-critical): {type(e).__name__}: {e}") logger.warning(f"Built-in MCP registration failed (non-critical): {type(e).__name__}: {e}")
try: try:
await asyncio.wait_for(mcp_manager.connect_all_enabled(), timeout=20) await mcp_manager.connect_all_enabled()
except asyncio.TimeoutError: except asyncio.TimeoutError:
logger.warning("User MCP startup timed out (non-critical)") logger.warning("User MCP startup timed out (non-critical)")
except BaseException as e: except BaseException as e:
+57 -24
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
@@ -192,43 +194,41 @@ class McpManager:
) )
stack = AsyncExitStack() stack = AsyncExitStack()
registered = False
try: try:
transport = await stack.enter_async_context(stdio_client(server_params)) transport = await stack.enter_async_context(stdio_client(server_params))
read_stream, write_stream = transport read_stream, write_stream = transport
session = await stack.enter_async_context(ClientSession(read_stream, write_stream)) session = await stack.enter_async_context(ClientSession(read_stream, write_stream))
await session.initialize() await session.initialize()
# Discover tools
tools_result = await session.list_tools() tools_result = await session.list_tools()
except Exception:
await stack.aclose()
raise
tools = [] tools = []
for tool in tools_result.tools: for tool in tools_result.tools:
tools.append({ tools.append({
"name": tool.name, "name": tool.name,
"description": tool.description or "", "description": tool.description or "",
"input_schema": tool.inputSchema if hasattr(tool, 'inputSchema') else {}, "input_schema": tool.inputSchema if hasattr(tool, "inputSchema") else {},
# MCP tool annotations (readOnlyHint / destructiveHint) drive # MCP tool annotations (readOnlyHint / destructiveHint) drive
# plan-mode read-only gating. Absent on many servers, so we # plan-mode read-only gating. Absent on many servers, so we
# fall back to a name heuristic in mcp_tool_is_readonly(). # fall back to a name heuristic in mcp_tool_is_readonly().
"annotations": getattr(tool, 'annotations', None), "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) # Extract identity hints from env vars (e.g. email address, API name)
# so tool descriptions can distinguish between multiple instances of # so tool descriptions can distinguish between multiple instances of
# the same MCP server (e.g. two email accounts). # the same MCP server (e.g. two email accounts).
identity_hints = [] identity_hints = []
for k, v in (env or {}).items(): for k, v in (env or {}).items():
k_lower = k.lower() k_lower = k.lower()
if any(x in k_lower for x in ['email_address', 'account', 'user', 'username']): if any(x in k_lower for x in ["email_address", "account", "user", "username"]):
identity_hints.append(v) identity_hints.append(v)
identity = ", ".join(identity_hints) if identity_hints else "" 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] = { self._connections[server_id] = {
"status": "connected", "status": "connected",
"name": name, "name": name,
@@ -237,12 +237,22 @@ class McpManager:
"identity": identity, "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") logger.info(f"MCP server connected: {name} ({server_id}) - {len(tools)} tools via stdio")
return True return True
except ImportError: except ImportError:
logger.warning("MCP package not installed. Install with: pip install mcp") 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 return False
async def _connect_sse(self, server_id: str, name: str, url: str) -> bool: async def _connect_sse(self, server_id: str, name: str, url: str) -> bool:
@@ -253,18 +263,16 @@ class McpManager:
from contextlib import AsyncExitStack from contextlib import AsyncExitStack
stack = AsyncExitStack() stack = AsyncExitStack()
registered = False
try: try:
transport = await stack.enter_async_context(sse_client(url)) transport = await stack.enter_async_context(sse_client(url))
read_stream, write_stream = transport read_stream, write_stream = transport
session = await stack.enter_async_context(ClientSession(read_stream, write_stream)) session = await stack.enter_async_context(ClientSession(read_stream, write_stream))
await session.initialize() await session.initialize()
# Discover tools
tools_result = await session.list_tools() tools_result = await session.list_tools()
except Exception:
await stack.aclose()
raise
tools = [] tools = []
for tool in tools_result.tools: for tool in tools_result.tools:
tools.append({ tools.append({
@@ -287,9 +295,15 @@ class McpManager:
"tool_count": len(tools), "tool_count": len(tools),
} }
registered = True
logger.info(f"MCP server connected: {name} ({server_id}) - {len(tools)} tools via SSE") logger.info(f"MCP server connected: {name} ({server_id}) - {len(tools)} tools via SSE")
return True return True
finally:
if not registered:
await stack.aclose()
except ImportError: except ImportError:
logger.warning("MCP package not installed. Install with: pip install mcp") 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}
@@ -409,17 +423,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:
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 [] args = json.loads(srv.args) if srv.args else []
env = json.loads(srv.env) if srv.env else {} env = json.loads(srv.env) if srv.env else {}
await self.connect_server(
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 +453,16 @@ class McpManager:
args=args, args=args,
env=env, env=env,
url=srv.url, url=srv.url,
),
timeout=20,
) )
finally: except asyncio.TimeoutError:
db.close() 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: 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