Odysseus v1.0
This commit is contained in:
pewdiepie-archdaemon
409
src/mcp_manager.py
Normal file
409
src/mcp_manager.py
Normal file
@@ -0,0 +1,409 @@
|
||||
"""
|
||||
mcp_manager.py
|
||||
|
||||
Manages connections to MCP (Model Context Protocol) tool servers.
|
||||
Each server exposes tools that are made available to the agent loop.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class McpManager:
|
||||
"""Manages MCP server connections and tool routing."""
|
||||
|
||||
def __init__(self):
|
||||
# server_id -> connection state
|
||||
self._connections: Dict[str, Dict[str, Any]] = {}
|
||||
# server_id -> list of tool schemas
|
||||
self._tools: Dict[str, List[Dict]] = {}
|
||||
# server_id -> MCP ClientSession
|
||||
self._sessions: Dict[str, Any] = {}
|
||||
# server_id -> exit stack (for cleanup)
|
||||
self._stacks: Dict[str, Any] = {}
|
||||
|
||||
async def connect_server(
|
||||
self,
|
||||
server_id: str,
|
||||
name: str,
|
||||
transport: str,
|
||||
command: Optional[str] = None,
|
||||
args: Optional[List[str]] = None,
|
||||
env: Optional[Dict[str, str]] = None,
|
||||
url: Optional[str] = None,
|
||||
) -> bool:
|
||||
"""Connect to an MCP server via stdio or SSE transport."""
|
||||
try:
|
||||
if transport == "stdio":
|
||||
return await self._connect_stdio(server_id, name, command, args or [], env or {})
|
||||
elif transport == "sse":
|
||||
return await self._connect_sse(server_id, name, url)
|
||||
else:
|
||||
logger.error(f"Unknown MCP transport: {transport}")
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to connect MCP server {name} ({server_id}): {e}")
|
||||
self._connections[server_id] = {"status": "error", "error": str(e), "name": name}
|
||||
return False
|
||||
|
||||
async def _connect_stdio(self, server_id: str, name: str, command: str, args: List[str], env: Dict[str, str]) -> bool:
|
||||
"""Connect to an MCP server via stdio transport."""
|
||||
try:
|
||||
from mcp import ClientSession, StdioServerParameters
|
||||
from mcp.client.stdio import stdio_client
|
||||
from contextlib import AsyncExitStack
|
||||
|
||||
server_params = StdioServerParameters(
|
||||
command=command,
|
||||
args=args,
|
||||
env={**os.environ, **env} if env else None,
|
||||
)
|
||||
|
||||
stack = AsyncExitStack()
|
||||
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()
|
||||
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 {},
|
||||
})
|
||||
|
||||
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 ""
|
||||
|
||||
self._connections[server_id] = {
|
||||
"status": "connected",
|
||||
"name": name,
|
||||
"transport": "stdio",
|
||||
"tool_count": len(tools),
|
||||
"identity": identity,
|
||||
}
|
||||
|
||||
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}
|
||||
return False
|
||||
|
||||
async def _connect_sse(self, server_id: str, name: str, url: str) -> bool:
|
||||
"""Connect to an MCP server via SSE transport."""
|
||||
try:
|
||||
from mcp import ClientSession
|
||||
from mcp.client.sse import sse_client
|
||||
from contextlib import AsyncExitStack
|
||||
|
||||
stack = AsyncExitStack()
|
||||
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()
|
||||
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 {},
|
||||
})
|
||||
|
||||
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),
|
||||
}
|
||||
|
||||
logger.info(f"MCP server connected: {name} ({server_id}) - {len(tools)} tools via SSE")
|
||||
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}
|
||||
return False
|
||||
|
||||
async def disconnect_server(self, server_id: str):
|
||||
"""Disconnect from an MCP server."""
|
||||
stack = self._stacks.pop(server_id, None)
|
||||
if stack:
|
||||
try:
|
||||
await stack.aclose()
|
||||
except Exception as e:
|
||||
logger.warning(f"Error closing MCP server {server_id}: {e}")
|
||||
|
||||
self._sessions.pop(server_id, None)
|
||||
self._tools.pop(server_id, None)
|
||||
self._connections.pop(server_id, None)
|
||||
logger.info(f"MCP server disconnected: {server_id}")
|
||||
|
||||
async def disconnect_all(self):
|
||||
"""Disconnect from all MCP servers."""
|
||||
ids = list(self._sessions.keys())
|
||||
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
|
||||
|
||||
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(
|
||||
server_id=srv.id,
|
||||
name=srv.name,
|
||||
transport=srv.transport,
|
||||
command=srv.command,
|
||||
args=args,
|
||||
env=env,
|
||||
url=srv.url,
|
||||
)
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
async def call_tool(self, qualified_name: str, arguments: Dict) -> Dict:
|
||||
"""Call an MCP tool by its qualified name (mcp__{server_id}__{tool_name}).
|
||||
|
||||
Returns a result dict compatible with agent_tools format.
|
||||
"""
|
||||
parts = qualified_name.split("__", 2)
|
||||
if len(parts) != 3 or parts[0] != "mcp":
|
||||
return {"error": f"Invalid MCP tool name: {qualified_name}", "exit_code": 1}
|
||||
|
||||
server_id = parts[1]
|
||||
tool_name = parts[2]
|
||||
|
||||
session = self._sessions.get(server_id)
|
||||
if not session:
|
||||
return {"error": f"MCP server not connected: {server_id}", "exit_code": 1}
|
||||
|
||||
try:
|
||||
result = await self._do_call(session, tool_name, arguments)
|
||||
except Exception as e:
|
||||
# Auto-reconnect for builtin servers whose subprocess may have died
|
||||
if self.is_builtin(server_id):
|
||||
logger.warning(f"MCP call failed for {qualified_name}, attempting reconnect: {e}")
|
||||
reconnected = await self._reconnect_builtin(server_id)
|
||||
if reconnected:
|
||||
session = self._sessions.get(server_id)
|
||||
if session:
|
||||
try:
|
||||
result = await self._do_call(session, tool_name, arguments)
|
||||
except Exception as e2:
|
||||
logger.error(f"MCP tool call failed after reconnect: {qualified_name}: {e2}")
|
||||
return {"error": str(e2), "exit_code": 1}
|
||||
else:
|
||||
return {"error": f"Reconnected but no session for {server_id}", "exit_code": 1}
|
||||
else:
|
||||
logger.error(f"MCP reconnect failed for {server_id}")
|
||||
return {"error": f"MCP server crashed and reconnect failed: {server_id}", "exit_code": 1}
|
||||
else:
|
||||
logger.error(f"MCP tool call failed: {qualified_name}: {e}")
|
||||
return {"error": str(e), "exit_code": 1}
|
||||
|
||||
return result
|
||||
|
||||
async def _do_call(self, session, tool_name: str, arguments: Dict) -> Dict:
|
||||
"""Execute a single MCP tool call and return result dict."""
|
||||
result = await session.call_tool(tool_name, arguments)
|
||||
output_parts = []
|
||||
images = []
|
||||
for content in result.content:
|
||||
if hasattr(content, 'text'):
|
||||
output_parts.append(content.text)
|
||||
elif getattr(content, 'type', '') == 'image' and hasattr(content, 'data'):
|
||||
# Image content (e.g. Playwright screenshots)
|
||||
mime = getattr(content, 'mimeType', 'image/png')
|
||||
images.append({"data": content.data, "mimeType": mime})
|
||||
output_parts.append(f"[Screenshot captured ({mime})]")
|
||||
elif hasattr(content, 'data'):
|
||||
output_parts.append(str(content.data))
|
||||
|
||||
output = "\n".join(output_parts)
|
||||
is_error = getattr(result, 'isError', False)
|
||||
|
||||
result_dict = {
|
||||
"stdout": output if not is_error else "",
|
||||
"stderr": output if is_error else "",
|
||||
"exit_code": 1 if is_error else 0,
|
||||
}
|
||||
if images:
|
||||
result_dict["images"] = images
|
||||
return result_dict
|
||||
|
||||
async def _reconnect_builtin(self, server_id: str) -> bool:
|
||||
"""Tear down and reconnect a crashed builtin MCP server."""
|
||||
import sys
|
||||
from src.builtin_mcp import _BUILTIN_SERVERS
|
||||
|
||||
if server_id not in _BUILTIN_SERVERS:
|
||||
return False
|
||||
|
||||
script_rel, name = _BUILTIN_SERVERS[server_id]
|
||||
base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
script_path = os.path.join(base_dir, script_rel)
|
||||
|
||||
# Clean up old connection
|
||||
await self.disconnect_server(server_id)
|
||||
|
||||
try:
|
||||
ok = await self.connect_server(
|
||||
server_id=server_id,
|
||||
name=name,
|
||||
transport="stdio",
|
||||
command=sys.executable,
|
||||
args=[script_path],
|
||||
env={"PYTHONPATH": base_dir},
|
||||
)
|
||||
if ok:
|
||||
logger.info(f"Reconnected builtin MCP server: {name}")
|
||||
return ok
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to reconnect builtin MCP server {name}: {e}")
|
||||
return False
|
||||
|
||||
def get_all_openai_schemas(self, disabled_map: Optional[Dict[str, set]] = None) -> List[Dict]:
|
||||
"""Return all MCP tools in OpenAI function-calling format.
|
||||
|
||||
Tool names are namespaced as mcp__{server_id}__{tool_name}.
|
||||
disabled_map: optional {server_id: set_of_disabled_tool_names} to filter out.
|
||||
"""
|
||||
schemas = []
|
||||
for server_id, tools in self._tools.items():
|
||||
# Skip builtin Python servers — they use the code-block tool format
|
||||
# But include NPX-based builtins (like browser) which need function calling
|
||||
if self.is_builtin(server_id) and server_id != "builtin_browser":
|
||||
continue
|
||||
conn = self._connections.get(server_id, {})
|
||||
server_name = conn.get("name", server_id)
|
||||
disabled = (disabled_map or {}).get(server_id, set())
|
||||
|
||||
identity = conn.get("identity", "")
|
||||
label = f"{server_name} ({identity})" if identity else server_name
|
||||
|
||||
for tool in tools:
|
||||
if tool["name"] in disabled:
|
||||
continue
|
||||
qualified = f"mcp__{server_id}__{tool['name']}"
|
||||
schema = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": qualified,
|
||||
"description": f"[MCP:{label}] {tool['description']}",
|
||||
"parameters": tool.get("input_schema", {"type": "object", "properties": {}}),
|
||||
},
|
||||
}
|
||||
schemas.append(schema)
|
||||
|
||||
return schemas
|
||||
|
||||
def get_all_tools(self, disabled_map: Optional[Dict[str, set]] = None) -> List[Dict]:
|
||||
"""Return a flat list of all discovered tools with server info."""
|
||||
result = []
|
||||
for server_id, tools in self._tools.items():
|
||||
conn = self._connections.get(server_id, {})
|
||||
disabled = (disabled_map or {}).get(server_id, set())
|
||||
for tool in tools:
|
||||
result.append({
|
||||
"server_id": server_id,
|
||||
"server_name": conn.get("name", server_id),
|
||||
"name": tool["name"],
|
||||
"qualified_name": f"mcp__{server_id}__{tool['name']}",
|
||||
"description": tool.get("description", ""),
|
||||
"is_disabled": tool["name"] in disabled,
|
||||
})
|
||||
return result
|
||||
|
||||
def is_builtin(self, server_id: str) -> bool:
|
||||
"""Check if a server is a built-in (auto-registered) server."""
|
||||
return server_id.startswith("builtin_") or server_id in {
|
||||
"image_gen",
|
||||
"memory",
|
||||
"rag",
|
||||
"email",
|
||||
}
|
||||
|
||||
def get_server_status(self, server_id: str) -> Dict:
|
||||
"""Get connection status for a server."""
|
||||
return self._connections.get(server_id, {"status": "disconnected"})
|
||||
|
||||
def get_all_statuses(self) -> Dict[str, Dict]:
|
||||
"""Get connection statuses for all servers."""
|
||||
return dict(self._connections)
|
||||
|
||||
_cached_prompt_desc = None
|
||||
_cached_prompt_desc_key = None
|
||||
|
||||
def get_tool_descriptions_for_prompt(self, disabled_map: Optional[Dict[str, set]] = None) -> str:
|
||||
"""Generate text describing MCP tools for the agent system prompt. Cached."""
|
||||
cache_key = (frozenset((k, frozenset(v)) for k, v in (disabled_map or {}).items()), len(self._tools))
|
||||
if self._cached_prompt_desc is not None and self._cached_prompt_desc_key == cache_key:
|
||||
return self._cached_prompt_desc
|
||||
tools = self.get_all_tools(disabled_map)
|
||||
if not tools:
|
||||
return ""
|
||||
|
||||
lines = ["\n\nYou also have access to external MCP tool servers. These tools are called via native function calling:"]
|
||||
by_server = {}
|
||||
for t in tools:
|
||||
# Skip builtin Python servers — they're already in the agent prompt
|
||||
# But include NPX-based builtins (like browser) which aren't hardcoded
|
||||
if self.is_builtin(t["server_id"]) and t["server_id"] != "builtin_browser":
|
||||
continue
|
||||
if t.get("is_disabled"):
|
||||
continue
|
||||
sn = t["server_name"]
|
||||
if sn not in by_server:
|
||||
by_server[sn] = []
|
||||
by_server[sn].append(t)
|
||||
|
||||
if not by_server:
|
||||
return ""
|
||||
|
||||
for server_name, server_tools in by_server.items():
|
||||
# Include identity (e.g. email address) if available
|
||||
sid = server_tools[0]["server_id"] if server_tools else ""
|
||||
identity = self._connections.get(sid, {}).get("identity", "")
|
||||
label = f"{server_name} ({identity})" if identity else server_name
|
||||
lines.append(f"\n**{label}:**")
|
||||
for t in server_tools:
|
||||
# Truncate long descriptions
|
||||
desc = t['description'][:120] + '...' if len(t['description']) > 120 else t['description']
|
||||
lines.append(f" - {t['qualified_name']}: {desc}")
|
||||
|
||||
result = "\n".join(lines)
|
||||
self._cached_prompt_desc = result
|
||||
self._cached_prompt_desc_key = cache_key
|
||||
return result
|
||||
Reference in New Issue
Block a user