mirror of
https://github.com/pewdiepie-archdaemon/odysseus.git
synced 2026-11-08 14:12:22 +01:00
Compare commits
2
Commits
df2fad2881
...
c1d6287e32
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c1d6287e32 | ||
|
|
5971b091db |
@@ -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
@@ -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}).
|
||||||
|
|||||||
@@ -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
|
||||||
|
|
||||||
|
|||||||
@@ -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
|
||||||
Reference in New Issue
Block a user