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:
+106 -73
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,57 +194,65 @@ 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 = []
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 tools = []
self._stacks[server_id] = stack for tool in tools_result.tools:
self._tools[server_id] = tools tools.append({
# Extract identity hints from env vars (e.g. email address, API name) "name": tool.name,
# so tool descriptions can distinguish between multiple instances of "description": tool.description or "",
# the same MCP server (e.g. two email accounts). "input_schema": tool.inputSchema if hasattr(tool, "inputSchema") else {},
identity_hints = [] # MCP tool annotations (readOnlyHint / destructiveHint) drive
for k, v in (env or {}).items(): # plan-mode read-only gating. Absent on many servers, so we
k_lower = k.lower() # fall back to a name heuristic in mcp_tool_is_readonly().
if any(x in k_lower for x in ['email_address', 'account', 'user', 'username']): "annotations": getattr(tool, "annotations", None),
identity_hints.append(v) })
identity = ", ".join(identity_hints) if identity_hints else ""
self._connections[server_id] = { # Extract identity hints from env vars (e.g. email address, API name)
"status": "connected", # so tool descriptions can distinguish between multiple instances of
"name": name, # the same MCP server (e.g. two email accounts).
"transport": "stdio", identity_hints = []
"tool_count": len(tools), for k, v in (env or {}).items():
"identity": identity, 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") 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,42 +263,46 @@ 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 = []
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 tools = []
self._stacks[server_id] = stack for tool in tools_result.tools:
self._tools[server_id] = tools tools.append({
self._connections[server_id] = { "name": tool.name,
"status": "connected", "description": tool.description or "",
"name": name, "input_schema": tool.inputSchema if hasattr(tool, 'inputSchema') else {},
"transport": "sse", # MCP tool annotations (readOnlyHint / destructiveHint) drive
"tool_count": len(tools), # 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") self._sessions[server_id] = session
return True 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: except ImportError:
logger.warning("MCP package not installed. Install with: pip install mcp") logger.warning("MCP package not installed. Install with: pip install mcp")
@@ -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:
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 +453,16 @@ 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)
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