"""MCP (Model Context Protocol) support — full client implementation. Architecture from Gemini CLI's McpClientManager: - Manages lifecycle of multiple MCP servers - Bridges MCP tools into the ToolRegistry as ToolBuilder instances - Supports stdio and SSE transports - Lazy initialization (connects on first use) Based on the MCP Python SDK. """ from __future__ import annotations import asyncio import json import logging import threading from typing import Any, Callable from code.config import MCPServerConfig logger = logging.getLogger(__name__) # ── MCP Tool Bridge ──────────────────────────────────────────────────── class MCPToolBridge: """Wraps an MCP tool as a ToolBuilder, making it indistinguishable from built-in tools. This is the key insight from Gemini CLI: MCP tools use the same ToolBuilder interface so they get the same policy, confirmation, and scheduling treatment. """ def __init__( self, server_name: str, tool_name: str, tool_schema: dict, call_fn: Callable, ): self.server_name = server_name self.tool_name = tool_name self._schema = tool_schema self._call_fn = call_fn # Generate MCP-prefixed name self.name = f"mcp_{server_name}_{tool_name}" self.description = tool_schema.get("description", f"MCP tool: {tool_name}") self.kind = "other" self.is_read_only = tool_schema.get("annotations", {}).get("readOnlyHint", False) def get_tool_definition(self) -> dict[str, Any]: return { "name": self.name, "description": self.description, "parameters": self._schema.get( "inputSchema", {"type": "object", "properties": {}} ), "kind": self.kind, "mcp_server": self.server_name, "mcp_tool": self.tool_name, } def call(self, arguments: dict[str, Any]) -> dict[str, Any]: """Execute the MCP tool call.""" return self._call_fn(self.tool_name, arguments) # ── MCP Client Manager ───────────────────────────────────────────────── class MCPClientManager: """Manages multiple MCP server connections. Lifecycle: configure → connect (lazy) → use → disconnect """ def __init__(self): self._configs: dict[str, MCPServerConfig] = {} self._clients: dict[str, Any] = {} # server_name → MCP client/session self._tools: dict[str, MCPToolBridge] = {} self._connected: set[str] = set() self._lock = threading.Lock() def configure(self, configs: list[MCPServerConfig]) -> None: """Set MCP server configurations.""" self._configs = {c.name: c for c in configs if c.enabled} def add_server(self, config: MCPServerConfig) -> None: """Add or update a single MCP server config.""" self._configs[config.name] = config # If already connected to this server, disconnect and reconnect if config.name in self._connected: self._disconnect(config.name) def remove_server(self, name: str) -> None: """Remove an MCP server.""" self._configs.pop(name, None) self._disconnect(name) def connect_all(self) -> dict[str, str | None]: """Connect to all configured MCP servers. Returns {name: error_or_None}.""" results = {} for name, config in self._configs.items(): if config.enabled: try: error = self._connect(name, config) results[name] = error except Exception as e: results[name] = str(e) logger.error("MCP connect failed for %s: %s", name, e) return results def disconnect_all(self) -> None: """Disconnect all MCP servers.""" for name in list(self._connected): self._disconnect(name) def get_tools(self) -> dict[str, MCPToolBridge]: """Get all discovered MCP tools.""" return dict(self._tools) def get_tool(self, name: str) -> MCPToolBridge | None: """Get a specific MCP tool by its full name (mcp_server_tool).""" return self._tools.get(name) def call_tool( self, server_name: str, tool_name: str, arguments: dict[str, Any] ) -> dict[str, Any]: """Call an MCP tool on a specific server.""" client = self._clients.get(server_name) if not client: # Try lazy connect config = self._configs.get(server_name) if not config: return {"error": f"MCP server '{server_name}' not configured"} error = self._connect(server_name, config) if error: return {"error": f"Failed to connect to '{server_name}': {error}"} client = self._clients.get(server_name) if not client: return {"error": f"MCP server '{server_name}' connection failed"} try: if hasattr(client, "call_tool"): result = client.call_tool(tool_name, arguments) if hasattr(result, "content"): # Extract text content from MCP result texts = [] for item in result.content: if hasattr(item, "text"): texts.append(item.text) else: texts.append(str(item)) return {"output": "\n".join(texts)} return {"output": str(result)} return {"output": str(client)} except Exception as e: return {"error": f"MCP tool call failed: {e}"} def list_resources(self, server_name: str | None = None) -> list[dict]: """List available MCP resources.""" resources = [] servers = [server_name] if server_name else list(self._connected) for name in servers: client = self._clients.get(name) if client and hasattr(client, "list_resources"): try: result = client.list_resources() for r in result.resources if hasattr(result, "resources") else []: resources.append( { "server": name, "uri": getattr(r, "uri", str(r)), "name": getattr(r, "name", ""), "description": getattr(r, "description", ""), } ) except Exception as e: logger.warning("Failed to list resources from %s: %s", name, e) return resources def get_status(self) -> dict[str, Any]: """Get status of all MCP servers.""" status = {} for name, config in self._configs.items(): status[name] = { "enabled": config.enabled, "connected": name in self._connected, "transport": config.transport, "tool_count": sum( 1 for t in self._tools.values() if t.server_name == name ), } return status # ── Internal ─────────────────────────────────────────────────────── def _connect(self, name: str, config: MCPServerConfig) -> str | None: """Connect to a single MCP server. Returns error message or None.""" if name in self._connected: return None try: # Try to import MCP SDK try: from mcp import ClientSession, StdioServerParameters # noqa: F401 from mcp.client.stdio import stdio_client # noqa: F401 except ImportError: logger.info("MCP SDK not installed. Skipping MCP server %s", name) return "mcp SDK not installed (pip install mcp)" if config.transport == "stdio": server_params = StdioServerParameters( command=config.command or "", args=config.args, env=config.env or None, ) # Run async connect in a new thread with event loop client_session, error = self._run_async_connect(name, server_params) if error or not client_session: return error or f"Failed to connect to {name}" self._clients[name] = client_session self._connected.add(name) # Discover tools self._discover_tools(name, client_session) logger.info( "MCP server '%s' connected with %d tools", name, sum(1 for t in self._tools.values() if t.server_name == name), ) return None elif config.transport == "sse": return f"SSE transport not yet supported for {name}" return f"Unknown transport: {config.transport}" except Exception as e: logger.error("MCP connect error for %s: %s", name, e) return str(e) def _run_async_connect(self, name: str, server_params): """Run async MCP connection in a sync context.""" from mcp import ClientSession from mcp.client.stdio import stdio_client result = {"session": None, "error": None} async def _connect_async(): try: async with stdio_client(server_params) as (read, write): async with ClientSession(read, write) as session: await session.initialize() result["session"] = _AsyncSessionWrapper(session) # Keep the context managers alive by blocking await asyncio.Future() # Block forever except Exception as e: result["error"] = str(e) loop = asyncio.new_event_loop() thread = threading.Thread( target=loop.run_until_complete, args=(_connect_async(),), daemon=True ) thread.start() thread.join(timeout=10) if result["error"] and not result["session"]: return None, result["error"] if result["session"]: return result["session"], None return None, "Connection timeout" def _discover_tools(self, server_name: str, client) -> None: """Discover and register tools from an MCP server.""" # Remove old tools for this server old_prefix = f"mcp_{server_name}_" for k in list(self._tools.keys()): if k.startswith(old_prefix): del self._tools[k] # Discover new tools (sync wrapper) try: tools_result = client.list_tools() if hasattr(client, "list_tools") else None tools = tools_result.tools if hasattr(tools_result, "tools") else [] for tool in tools: schema = { "description": getattr(tool, "description", ""), "inputSchema": getattr( tool, "inputSchema", {"type": "object", "properties": {}} ), "annotations": getattr(tool, "annotations", {}), } tool_name = getattr(tool, "name", "unknown") bridge = MCPToolBridge( server_name=server_name, tool_name=tool_name, tool_schema=schema, call_fn=lambda sn=server_name, tn=tool_name, args=None: self.call_tool(sn, tn, args or {}), ) self._tools[bridge.name] = bridge logger.debug("Discovered MCP tool: %s", bridge.name) except Exception as e: logger.warning("Failed to discover tools from %s: %s", server_name, e) def _disconnect(self, name: str) -> None: """Disconnect from an MCP server.""" self._clients.pop(name, None) self._connected.discard(name) # Remove tools prefix = f"mcp_{name}_" for k in list(self._tools.keys()): if k.startswith(prefix): del self._tools[k] class _AsyncSessionWrapper: """Wraps an MCP ClientSession for sync access.""" def __init__(self, session): self._session = session def __getattr__(self, name): attr = getattr(self._session, name) if asyncio.iscoroutinefunction(attr): def sync_wrapper(*args, **kwargs): try: loop = asyncio.get_event_loop() if loop.is_running(): # Can't await in running loop, return a placeholder import concurrent.futures with concurrent.futures.ThreadPoolExecutor() as pool: future = pool.submit(asyncio.run, attr(*args, **kwargs)) return future.result(timeout=30) else: return loop.run_until_complete(attr(*args, **kwargs)) except RuntimeError: return asyncio.run(attr(*args, **kwargs)) return sync_wrapper return attr # ── Global Manager ───────────────────────────────────────────────────── _global_manager: MCPClientManager | None = None def get_mcp_manager() -> MCPClientManager: """Get or create the global MCP client manager.""" global _global_manager if _global_manager is None: _global_manager = MCPClientManager() return _global_manager __all__ = ["MCPToolBridge", "MCPClientManager", "get_mcp_manager"]