wehub-resource-sync
301 lines
8.7 KiB
Python
301 lines
8.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Adapted from vllm-mlx (https://github.com/vllm-project/vllm-mlx).
|
|
"""
|
|
MCP Client Manager for handling multiple MCP server connections.
|
|
"""
|
|
|
|
import asyncio
|
|
import logging
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from .client import MCPClient
|
|
from .tools import merge_tools, mcp_tools_to_openai, openai_call_to_mcp
|
|
from .types import (
|
|
MCPConfig,
|
|
MCPServerStatus,
|
|
MCPTool,
|
|
MCPToolResult,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class MCPClientManager:
|
|
"""
|
|
Manages multiple MCP server connections.
|
|
|
|
Provides a unified interface for:
|
|
- Connecting to multiple MCP servers
|
|
- Discovering and aggregating tools
|
|
- Executing tool calls
|
|
- Managing connection lifecycle
|
|
"""
|
|
|
|
def __init__(self, config: MCPConfig):
|
|
"""
|
|
Initialize MCP Client Manager.
|
|
|
|
Args:
|
|
config: MCP configuration with server definitions
|
|
"""
|
|
self.config = config
|
|
self._clients: Dict[str, MCPClient] = {}
|
|
self._started = False
|
|
self._lock = asyncio.Lock()
|
|
|
|
# Create clients for each server
|
|
for name, server_config in config.servers.items():
|
|
self._clients[name] = MCPClient(server_config)
|
|
|
|
@property
|
|
def is_started(self) -> bool:
|
|
"""Check if manager has been started."""
|
|
return self._started
|
|
|
|
async def start(self):
|
|
"""
|
|
Start the manager and connect to all enabled servers.
|
|
|
|
Connections are made in parallel for faster startup.
|
|
"""
|
|
async with self._lock:
|
|
if self._started:
|
|
return
|
|
|
|
logger.info(f"Starting MCP client manager with {len(self._clients)} servers")
|
|
|
|
# Connect to all servers in parallel
|
|
tasks = [
|
|
client.connect()
|
|
for client in self._clients.values()
|
|
if client.config.enabled
|
|
]
|
|
|
|
if tasks:
|
|
results = await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
# Log results
|
|
for client, result in zip(
|
|
[c for c in self._clients.values() if c.config.enabled],
|
|
results,
|
|
):
|
|
if isinstance(result, Exception):
|
|
logger.error(f"Failed to connect to '{client.name}': {result}")
|
|
elif result:
|
|
logger.info(f"Connected to '{client.name}'")
|
|
|
|
self._started = True
|
|
|
|
# Log summary
|
|
connected = sum(1 for c in self._clients.values() if c.is_connected)
|
|
total_tools = sum(len(c.tools) for c in self._clients.values())
|
|
logger.info(
|
|
f"MCP manager started: {connected}/{len(self._clients)} servers, "
|
|
f"{total_tools} tools available"
|
|
)
|
|
|
|
async def stop(self):
|
|
"""Stop the manager and disconnect from all servers."""
|
|
async with self._lock:
|
|
if not self._started:
|
|
return
|
|
|
|
logger.info("Stopping MCP client manager")
|
|
|
|
# Disconnect from all servers in parallel
|
|
tasks = [client.disconnect() for client in self._clients.values()]
|
|
if tasks:
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
self._started = False
|
|
logger.info("MCP client manager stopped")
|
|
|
|
def get_all_tools(self) -> List[MCPTool]:
|
|
"""
|
|
Get all tools from all connected servers.
|
|
|
|
Returns:
|
|
List of MCPTool instances
|
|
"""
|
|
tools = []
|
|
for client in self._clients.values():
|
|
if client.is_connected:
|
|
tools.extend(client.tools)
|
|
return tools
|
|
|
|
def get_all_tools_openai(self) -> List[Dict[str, Any]]:
|
|
"""
|
|
Get all tools in OpenAI function calling format.
|
|
|
|
Returns:
|
|
List of OpenAI-compatible tool definitions
|
|
"""
|
|
return mcp_tools_to_openai(self.get_all_tools())
|
|
|
|
def get_merged_tools(
|
|
self,
|
|
user_tools: Optional[List[Dict[str, Any]]] = None,
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
Get MCP tools merged with user-provided tools.
|
|
|
|
User tools take precedence on name conflicts.
|
|
|
|
Args:
|
|
user_tools: Optional user-provided tools in OpenAI format
|
|
|
|
Returns:
|
|
Combined list of tools in OpenAI format
|
|
"""
|
|
return merge_tools(self.get_all_tools(), user_tools)
|
|
|
|
def get_server_status(self) -> List[MCPServerStatus]:
|
|
"""
|
|
Get status of all servers.
|
|
|
|
Returns:
|
|
List of MCPServerStatus for each server
|
|
"""
|
|
return [client.get_status() for client in self._clients.values()]
|
|
|
|
def get_client(self, server_name: str) -> Optional[MCPClient]:
|
|
"""
|
|
Get client for a specific server.
|
|
|
|
Args:
|
|
server_name: Name of the server
|
|
|
|
Returns:
|
|
MCPClient instance or None if not found
|
|
"""
|
|
return self._clients.get(server_name)
|
|
|
|
async def execute_tool(
|
|
self,
|
|
full_name: str,
|
|
arguments: Dict[str, Any],
|
|
timeout: Optional[float] = None,
|
|
) -> MCPToolResult:
|
|
"""
|
|
Execute a tool by its full name (server__tool).
|
|
|
|
Args:
|
|
full_name: Full tool name with server prefix
|
|
arguments: Tool arguments
|
|
timeout: Optional timeout in seconds
|
|
|
|
Returns:
|
|
MCPToolResult with the result or error
|
|
"""
|
|
# Parse full name
|
|
server_name, tool_name, _ = openai_call_to_mcp({
|
|
"function": {"name": full_name, "arguments": "{}"}
|
|
})
|
|
|
|
# If no server prefix, try to find the tool
|
|
if not server_name:
|
|
server_name = self._find_tool_server(full_name)
|
|
tool_name = full_name
|
|
|
|
if not server_name:
|
|
return MCPToolResult(
|
|
tool_name=full_name,
|
|
content=None,
|
|
is_error=True,
|
|
error_message=f"Tool '{full_name}' not found in any connected server",
|
|
)
|
|
|
|
# Get client
|
|
client = self._clients.get(server_name)
|
|
if not client:
|
|
return MCPToolResult(
|
|
tool_name=full_name,
|
|
content=None,
|
|
is_error=True,
|
|
error_message=f"Server '{server_name}' not found",
|
|
)
|
|
|
|
if not client.is_connected:
|
|
return MCPToolResult(
|
|
tool_name=full_name,
|
|
content=None,
|
|
is_error=True,
|
|
error_message=f"Server '{server_name}' is not connected",
|
|
)
|
|
|
|
# Execute tool
|
|
return await client.call_tool(
|
|
tool_name,
|
|
arguments,
|
|
timeout=timeout or self.config.default_timeout,
|
|
)
|
|
|
|
async def execute_tool_call(
|
|
self,
|
|
tool_call: Dict[str, Any],
|
|
timeout: Optional[float] = None,
|
|
) -> MCPToolResult:
|
|
"""
|
|
Execute a tool call from OpenAI format.
|
|
|
|
Args:
|
|
tool_call: OpenAI tool call object
|
|
timeout: Optional timeout in seconds
|
|
|
|
Returns:
|
|
MCPToolResult with the result or error
|
|
"""
|
|
server_name, tool_name, arguments = openai_call_to_mcp(tool_call)
|
|
|
|
if server_name:
|
|
full_name = f"{server_name}__{tool_name}"
|
|
else:
|
|
full_name = tool_name
|
|
|
|
return await self.execute_tool(full_name, arguments, timeout)
|
|
|
|
def _find_tool_server(self, tool_name: str) -> Optional[str]:
|
|
"""
|
|
Find which server has a tool by name.
|
|
|
|
Args:
|
|
tool_name: Tool name (without server prefix)
|
|
|
|
Returns:
|
|
Server name or None if not found
|
|
"""
|
|
for client in self._clients.values():
|
|
if client.is_connected:
|
|
for tool in client.tools:
|
|
if tool.name == tool_name:
|
|
return client.name
|
|
return None
|
|
|
|
async def refresh_tools(self):
|
|
"""Refresh tools from all connected servers."""
|
|
tasks = [
|
|
client.refresh_tools()
|
|
for client in self._clients.values()
|
|
if client.is_connected
|
|
]
|
|
if tasks:
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
async def reconnect(self, server_name: Optional[str] = None):
|
|
"""
|
|
Reconnect to server(s).
|
|
|
|
Args:
|
|
server_name: Specific server to reconnect, or None for all
|
|
"""
|
|
if server_name:
|
|
client = self._clients.get(server_name)
|
|
if client:
|
|
await client.disconnect()
|
|
await client.connect()
|
|
else:
|
|
# Reconnect all
|
|
for client in self._clients.values():
|
|
await client.disconnect()
|
|
await client.connect()
|