| import asyncio |
| from typing import Dict, Any, List |
| from mcp import ClientSession, StdioServerParameters |
| from mcp.client.sse import sse_client |
| from mcp.client.stdio import stdio_client |
| from mcp.client.streamable_http import streamablehttp_client |
| from utils.logger import logger |
|
|
|
|
| class MCPConnectionManager: |
| def __init__(self): |
| self.connected_servers: Dict[str, Dict[str, Any]] = {} |
| |
| async def connect_sse_server(self, server_name: str, server_config: Dict[str, Any], timeout: int = 15) -> Dict[str, Any]: |
| url = server_config["url"] |
| headers = server_config.get("headers", {}) |
| |
| async with asyncio.timeout(timeout): |
| try: |
| async with sse_client(url, headers=headers) as (read, write): |
| async with ClientSession(read, write) as session: |
| await session.initialize() |
| tools_result = await session.list_tools() |
| |
| tools_info = [ |
| { |
| "name": tool.name, |
| "description": tool.description, |
| "input_schema": tool.inputSchema |
| } |
| for tool in tools_result.tools |
| ] |
| |
| server_info = { |
| "status": "connected", |
| "transport": "sse", |
| "url": url, |
| "tools": tools_info |
| } |
| |
| self.connected_servers[server_name] = server_info |
| logger.info(f"Connected to {server_name} via SSE ({len(tools_info)} tools)") |
| return server_info |
| |
| except TypeError as e: |
| if "unexpected keyword argument" in str(e): |
| async with sse_client(url) as (read, write): |
| async with ClientSession(read, write) as session: |
| await session.initialize() |
| tools_result = await session.list_tools() |
| |
| tools_info = [ |
| { |
| "name": tool.name, |
| "description": tool.description, |
| "input_schema": tool.inputSchema |
| } |
| for tool in tools_result.tools |
| ] |
| |
| server_info = { |
| "status": "connected", |
| "transport": "sse", |
| "url": url, |
| "tools": tools_info |
| } |
| |
| self.connected_servers[server_name] = server_info |
| logger.info(f"Connected to {server_name} via SSE ({len(tools_info)} tools)") |
| return server_info |
| else: |
| raise |
| |
| async def connect_http_server(self, server_name: str, server_config: Dict[str, Any], timeout: int = 15) -> Dict[str, Any]: |
| url = server_config["url"] |
| |
| async with asyncio.timeout(timeout): |
| async with streamablehttp_client(url) as (read_stream, write_stream, _): |
| async with ClientSession(read_stream, write_stream) as session: |
| await session.initialize() |
| tools_result = await session.list_tools() |
| |
| tools_info = [ |
| { |
| "name": tool.name, |
| "description": tool.description, |
| "input_schema": tool.inputSchema |
| } |
| for tool in tools_result.tools |
| ] |
| |
| server_info = { |
| "status": "connected", |
| "transport": "http", |
| "url": url, |
| "tools": tools_info |
| } |
| |
| self.connected_servers[server_name] = server_info |
| logger.info(f"Connected to {server_name} via HTTP ({len(tools_info)} tools)") |
| return server_info |
| |
| async def connect_stdio_server(self, server_name: str, server_config: Dict[str, Any], timeout: int = 15) -> Dict[str, Any]: |
| server_params = StdioServerParameters( |
| command=server_config["command"], |
| args=server_config.get("args", []), |
| env=server_config.get("env", {}) |
| ) |
| |
| async with asyncio.timeout(timeout): |
| async with stdio_client(server_params) as (read, write): |
| async with ClientSession(read, write) as session: |
| await session.initialize() |
| tools_result = await session.list_tools() |
| |
| tools_info = [ |
| { |
| "name": tool.name, |
| "description": tool.description, |
| "input_schema": tool.inputSchema |
| } |
| for tool in tools_result.tools |
| ] |
| |
| server_info = { |
| "status": "connected", |
| "transport": "stdio", |
| "tools": tools_info |
| } |
| |
| self.connected_servers[server_name] = server_info |
| logger.info(f"Connected to {server_name} via stdio ({len(tools_info)} tools)") |
| return server_info |
| |
| def get_server_info(self, server_name: str) -> Dict[str, Any]: |
| return self.connected_servers.get(server_name, {}) |
| |
| def get_all_servers(self) -> Dict[str, Dict[str, Any]]: |
| return self.connected_servers.copy() |
