Manager for MCP (Model Context Protocol) server connections. Handles: - Auto-discovery of local MCP servers - Connection management for local (stdio) and remote (SSE/HTTP) servers - Tool listing and execution across all connected servers - Tool name to server session ma
| 19 | |
| 20 | |
| 21 | class MCPManager: |
| 22 | """ |
| 23 | Manager for MCP (Model Context Protocol) server connections. |
| 24 | |
| 25 | Handles: |
| 26 | - Auto-discovery of local MCP servers |
| 27 | - Connection management for local (stdio) and remote (SSE/HTTP) servers |
| 28 | - Tool listing and execution across all connected servers |
| 29 | - Tool name to server session mapping |
| 30 | """ |
| 31 | |
| 32 | def __init__(self): |
| 33 | """Initialize the MCP manager.""" |
| 34 | self.sessions: Dict[str, ClientSession] = {} # server_id -> session |
| 35 | self.tool_to_server: Dict[str, str] = {} # tool_name -> server_id |
| 36 | self.exit_stack = AsyncExitStack() |
| 37 | self._initialized = False |
| 38 | |
| 39 | async def initialize(self, |
| 40 | local_servers_dir: Optional[str] = None, |
| 41 | remote_servers: Optional[List[Dict[str, Any]]] = None): |
| 42 | """ |
| 43 | Initialize all MCP server connections. |
| 44 | |
| 45 | Args: |
| 46 | local_servers_dir: Directory containing local MCP server scripts. |
| 47 | Defaults to internagent/mas/tools/mcp/ |
| 48 | remote_servers: List of remote server configurations, each containing: |
| 49 | - id: Unique identifier |
| 50 | - url: Server URL (can be /sse or /mcp endpoint) |
| 51 | - headers: Optional HTTP headers (for authentication) |
| 52 | - protocol: Optional, 'sse' or 'http' (auto-detected if not specified) |
| 53 | """ |
| 54 | if self._initialized: |
| 55 | logger.warning("MCPManager already initialized") |
| 56 | return |
| 57 | |
| 58 | # Auto-discover and connect local servers |
| 59 | if local_servers_dir is None: |
| 60 | current_dir = Path(__file__).parent |
| 61 | local_servers_dir = current_dir / "mcp" |
| 62 | else: |
| 63 | local_servers_dir = Path(local_servers_dir) |
| 64 | |
| 65 | if local_servers_dir.exists(): |
| 66 | await self._discover_and_connect_local_servers(local_servers_dir) |
| 67 | else: |
| 68 | logger.warning(f"Local MCP servers directory not found: {local_servers_dir}") |
| 69 | |
| 70 | # Connect to remote servers if provided |
| 71 | if remote_servers: |
| 72 | logger.info(f"Attempting to connect to {len(remote_servers)} remote server(s)") |
| 73 | for server_config in remote_servers: |
| 74 | try: |
| 75 | await self._connect_remote_server(server_config) |
| 76 | except Exception as e: |
| 77 | server_id = server_config.get('id', 'unknown') |
| 78 | logger.error(f"Failed to connect to remote server {server_id}: {e}") |