""" MCP 管理插件 管理 MCP (Model Context Protocol) 服务器，将 MCP 工具自动注册到 ToolRegistry， 使 AI 可以调用各种 MCP 服务器提供的工具。 功能： - 自动连接配置的 MCP 服务器 - 将 MCP 工具转换为 OpenAI 格式并注册 - 支持热重载（禁用/启用插件时自动管理连接） - 提供管理命令（查看状态、重连等） """ import tomllib from pathlib import Path from typing import Any, Dict, List from loguru import logger from utils.plugin_base import PluginBase from utils.decorators import on_text_message from utils.tool_registry import get_tool_registry from .mcp_client import MCPManager, MCPServerConfig class MCPManagerPlugin(PluginBase): """MCP 管理插件""" description = "MCP 服务器管理，自动注册 MCP 工具到 AI" author = "ShiHao" version = "1.0.0" # 高优先级加载，确保在其他插件之前注册工具 load_priority = 90 def __init__(self): super().__init__() self.config = None self.mcp_manager: MCPManager = None self._registered_tools: List[str] = [] # 已注册的工具名列表 async def async_init(self): """插件异步初始化""" # 读取配置 config_path = Path(__file__).parent / "config.toml" with open(config_path, "rb") as f: self.config = tomllib.load(f) mcp_config = self.config.get("mcp", {}) if not mcp_config.get("enabled", True): logger.info("MCPManager: MCP 功能已禁用") return # 初始化 MCP 管理器 self.mcp_manager = MCPManager( tool_timeout=mcp_config.get("tool_timeout", 60), server_start_timeout=mcp_config.get("server_start_timeout", 30) ) # 连接所有配置的服务器 servers = mcp_config.get("servers", []) if not servers: logger.info("MCPManager: 未配置任何 MCP 服务器") return if mcp_config.get("auto_connect", True): await self._connect_all_servers(servers) logger.success(f"MCPManager 插件已加载，已连接 {len(self.mcp_manager.clients)} 个 MCP 服务器") async def _connect_all_servers(self, servers: List[Dict]): """连接所有配置的服务器""" for server_data in servers: try: config = MCPServerConfig.from_dict(server_data) success = await self.mcp_manager.add_server(config) if success: # 注册工具到 ToolRegistry await self._register_server_tools(config.name) except Exception as e: logger.error(f"MCPManager: 连接服务器 {server_data.get('name', '未知')} 失败: {e}") async def _register_server_tools(self, server_name: str): """将服务器的工具注册到 ToolRegistry""" client = self.mcp_manager.clients.get(server_name) if not client: return registry = get_tool_registry() prefix = client.config.tool_prefix for tool in client.tools.values(): schema = tool.to_openai_schema(prefix) tool_name = schema["function"]["name"] # 创建工具执行器 async def executor(name: str, arguments: Dict, bot, from_wxid: str, _tn=tool_name) -> Dict: return await self.mcp_manager.call_tool(_tn, arguments) # 注册到 ToolRegistry success = registry.register( name=tool_name, plugin_name=f"MCP:{server_name}", schema=schema, executor=executor, timeout=self.mcp_manager.tool_timeout, priority=40 # 比普通插件优先级稍低 ) if success: self._registered_tools.append(tool_name) logger.debug(f"MCPManager: 注册工具 {tool_name}") logger.info(f"MCPManager: 从 {server_name} 注册了 {len(client.tools)} 个工具") async def on_disable(self): """插件禁用时清理""" await super().on_disable() # 注销所有已注册的工具 registry = get_tool_registry() for tool_name in self._registered_tools: registry.unregister(tool_name) self._registered_tools.clear() # 关闭所有 MCP 服务器连接 if self.mcp_manager: await self.mcp_manager.shutdown() self.mcp_manager = None logger.info("MCPManager: 已清理所有 MCP 连接和工具注册") # ==================== 管理命令 ==================== @on_text_message(priority=80) async def handle_admin_commands(self, bot, message: dict): """处理管理命令""" content = message.get("Content", "").strip() # 只响应管理员的命令 # TODO: 从主配置读取管理员列表 if not content.startswith("/mcp"): return True parts = content.split() if len(parts) < 2: return True cmd = parts[1].lower() reply_to = message.get("FromWxid", "") if cmd == "status": await self._cmd_status(bot, reply_to) return False elif cmd == "list": await self._cmd_list_tools(bot, reply_to) return False elif cmd == "reload": await self._cmd_reload(bot, reply_to) return False return True async def _cmd_status(self, bot, reply_to: str): """查看 MCP 服务器状态""" if not self.mcp_manager: await bot.send_text(reply_to, "MCP 功能未启用") return servers = self.mcp_manager.list_servers() if not servers: await bot.send_text(reply_to, "没有已连接的 MCP 服务器") return lines = ["📡 MCP 服务器状态："] for s in servers: status = "✅" if s["connected"] else "❌" lines.append(f"{status} {s['name']}: {s['tools_count']} 个工具") await bot.send_text(reply_to, "\n".join(lines)) async def _cmd_list_tools(self, bot, reply_to: str): """列出所有 MCP 工具""" if not self.mcp_manager: await bot.send_text(reply_to, "MCP 功能未启用") return servers = self.mcp_manager.list_servers() if not servers: await bot.send_text(reply_to, "没有已连接的 MCP 服务器") return lines = ["🔧 MCP 工具列表："] for s in servers: if s["tools"]: lines.append(f"\n【{s['name']}】") for tool in s["tools"][:10]: # 限制显示数量 lines.append(f" • {tool}") if len(s["tools"]) > 10: lines.append(f" ... 还有 {len(s['tools']) - 10} 个") await bot.send_text(reply_to, "\n".join(lines)) async def _cmd_reload(self, bot, reply_to: str): """重新加载 MCP 服务器""" await bot.send_text(reply_to, "正在重新加载 MCP 服务器...") # 清理现有连接 if self.mcp_manager: registry = get_tool_registry() for tool_name in self._registered_tools: registry.unregister(tool_name) self._registered_tools.clear() await self.mcp_manager.shutdown() # 重新读取配置 config_path = Path(__file__).parent / "config.toml" with open(config_path, "rb") as f: self.config = tomllib.load(f) mcp_config = self.config.get("mcp", {}) # 重新初始化 self.mcp_manager = MCPManager( tool_timeout=mcp_config.get("tool_timeout", 60), server_start_timeout=mcp_config.get("server_start_timeout", 30) ) servers = mcp_config.get("servers", []) await self._connect_all_servers(servers) await bot.send_text( reply_to, f"MCP 重新加载完成，已连接 {len(self.mcp_manager.clients)} 个服务器" ) # ==================== LLM 工具接口（备用） ==================== def get_llm_tools(self) -> List[Dict]: """ 返回 MCP 工具列表（备用接口） 注意：工具已通过 ToolRegistry 注册，此方法仅供参考 """ if not self.mcp_manager: return [] return self.mcp_manager.get_all_tools() async def execute_llm_tool( self, tool_name: str, arguments: Dict[str, Any], bot, from_wxid: str ) -> Dict[str, Any]: """ 执行 MCP 工具（备用接口） 注意：工具已通过 ToolRegistry 注册，此方法仅供备用 """ if not self.mcp_manager: return {"success": False, "error": "MCP 功能未启用"} return await self.mcp_manager.call_tool(tool_name, arguments)