"""MCP插件注册表 - 管理运行时插件实例""" import asyncio import time import httpx from typing import Dict, Optional, Any, List, Tuple from collections import OrderedDict from app.mcp.http_client import HTTPMCPClient, MCPError from app.models.mcp_plugin import MCPPlugin from app.logger import get_logger logger = get_logger(__name__) class MCPPluginRegistry: """MCP插件注册表 - 管理运行时插件实例（多用户优化版）""" def __init__(self, max_clients: int = 1000, client_ttl: int = 3600): """ 初始化注册表 Args: max_clients: 最大缓存客户端数量 client_ttl: 客户端过期时间（秒），默认1小时 """ # 存储格式: {plugin_id: (client, last_access_time)} self._clients: OrderedDict[str, Tuple[HTTPMCPClient, float]] = OrderedDict() # 细粒度锁：每个用户一个锁 self._user_locks: Dict[str, asyncio.Lock] = {} self._locks_lock = asyncio.Lock() # 保护locks字典本身 # 配置参数 self._max_clients = max_clients self._client_ttl = client_ttl # 共享HTTP客户端池（用于所有MCP HTTP请求） self._shared_http_client = httpx.AsyncClient( limits=httpx.Limits( max_keepalive_connections=100, max_connections=200, keepalive_expiry=30.0 ), timeout=httpx.Timeout(connect=10.0, read=60.0, write=10.0, pool=5.0), headers={ "User-Agent": "MuMuAINovel-MCP-Client/1.0" } ) # 启动后台清理任务 self._cleanup_task = None self._start_cleanup_task() def _start_cleanup_task(self): """启动后台清理任务""" if self._cleanup_task is None: self._cleanup_task = asyncio.create_task(self._cleanup_loop()) logger.info("✅ MCP插件注册表后台清理任务已启动") async def _cleanup_loop(self): """后台清理过期客户端""" while True: try: await asyncio.sleep(300) # 每5分钟清理一次 await self._cleanup_expired_clients() except asyncio.CancelledError: break except Exception as e: logger.error(f"清理任务异常: {e}") async def _cleanup_expired_clients(self): """清理过期的客户端""" now = time.time() expired_ids = [] # 收集过期的plugin_id for plugin_id, (client, last_access) in list(self._clients.items()): if now - last_access > self._client_ttl: expired_ids.append(plugin_id) if expired_ids: logger.info(f"🧹 清理 {len(expired_ids)} 个过期的MCP客户端") for plugin_id in expired_ids: # 提取user_id来获取对应的锁 user_id = plugin_id.split(':', 1)[0] user_lock = await self._get_user_lock(user_id) async with user_lock: if plugin_id in self._clients: await self._unload_plugin_unsafe(plugin_id) async def _get_user_lock(self, user_id: str) -> asyncio.Lock: """ 获取用户专属的锁（细粒度锁） Args: user_id: 用户ID Returns: 该用户的锁对象 """ async with self._locks_lock: if user_id not in self._user_locks: self._user_locks[user_id] = asyncio.Lock() return self._user_locks[user_id] def _touch_client(self, plugin_id: str): """ 更新客户端的最后访问时间（LRU） Args: plugin_id: 插件ID """ if plugin_id in self._clients: client, _ = self._clients[plugin_id] self._clients[plugin_id] = (client, time.time()) # 移到末尾（LRU） self._clients.move_to_end(plugin_id) async def _evict_lru_client(self): """驱逐最久未使用的客户端（当达到max_clients限制时）""" if len(self._clients) >= self._max_clients: # 获取最旧的plugin_id oldest_id = next(iter(self._clients)) logger.info(f"📤 达到最大客户端数量限制，驱逐: {oldest_id}") await self._unload_plugin_unsafe(oldest_id) async def load_plugin(self, plugin: MCPPlugin) -> bool: """ 从配置加载插件 Args: plugin: 插件配置 Returns: 是否加载成功 """ # 使用细粒度锁（只锁定当前用户） user_lock = await self._get_user_lock(plugin.user_id) async with user_lock: try: plugin_id = f"{plugin.user_id}:{plugin.plugin_name}" # 如果已加载，先卸载 if plugin_id in self._clients: await self._unload_plugin_unsafe(plugin_id) # 检查是否需要驱逐LRU客户端 await self._evict_lru_client() # 目前只支持HTTP类型 if plugin.plugin_type == "http": if not plugin.server_url: logger.error(f"HTTP插件缺少server_url: {plugin.plugin_name}") return False # 使用共享HTTP连接池创建客户端 client = HTTPMCPClient( url=plugin.server_url, headers=plugin.headers or {}, env=plugin.env or {}, timeout=plugin.config.get('timeout', 60.0) if plugin.config else 60.0, http_client=self._shared_http_client # 传入共享连接池 ) # 存储客户端和当前时间戳 self._clients[plugin_id] = (client, time.time()) logger.info(f"✅ 加载MCP插件: {plugin_id}") return True else: logger.warning(f"暂不支持的插件类型: {plugin.plugin_type}") return False except Exception as e: logger.error(f"加载插件失败 {plugin.plugin_name}: {e}") return False async def unload_plugin(self, user_id: str, plugin_name: str): """ 卸载插件 Args: user_id: 用户ID plugin_name: 插件名称 """ # 使用细粒度锁（只锁定当前用户） user_lock = await self._get_user_lock(user_id) async with user_lock: plugin_id = f"{user_id}:{plugin_name}" await self._unload_plugin_unsafe(plugin_id) async def _unload_plugin_unsafe(self, plugin_id: str): """卸载插件（不加锁，内部使用）""" if plugin_id in self._clients: client, _ = self._clients[plugin_id] # 解包 (client, timestamp) try: await client.close() except Exception as e: logger.error(f"关闭插件客户端失败 {plugin_id}: {e}") del self._clients[plugin_id] logger.info(f"卸载MCP插件: {plugin_id}") async def reload_plugin(self, plugin: MCPPlugin) -> bool: """ 重新加载插件 Args: plugin: 插件配置 Returns: 是否重载成功 """ await self.unload_plugin(plugin.user_id, plugin.plugin_name) return await self.load_plugin(plugin) def get_client(self, user_id: str, plugin_name: str) -> Optional[HTTPMCPClient]: """ 获取插件客户端（支持LRU访问时间更新） Args: user_id: 用户ID plugin_name: 插件名称 Returns: 客户端实例或None """ plugin_id = f"{user_id}:{plugin_name}" entry = self._clients.get(plugin_id) if entry: # 更新访问时间（LRU） self._touch_client(plugin_id) return entry[0] # 返回客户端对象 return None async def call_tool( self, user_id: str, plugin_name: str, tool_name: str, arguments: Dict[str, Any] ) -> Any: """ 调用插件工具 Args: user_id: 用户ID plugin_name: 插件名称 tool_name: 工具名称 arguments: 工具参数 Returns: 工具执行结果 Raises: ValueError: 插件不存在或未启用 MCPError: 工具调用失败 """ client = self.get_client(user_id, plugin_name) if not client: raise ValueError(f"插件未加载: {plugin_name}") try: result = await client.call_tool(tool_name, arguments) logger.info(f"✅ 工具调用成功: {plugin_name}.{tool_name}") # logger.info(f"✅ 工具返回内容: {result}") return result except Exception as e: logger.error(f"❌ 工具调用失败: {plugin_name}.{tool_name}, 错误: {e}") raise async def get_plugin_tools( self, user_id: str, plugin_name: str ) -> List[Dict[str, Any]]: """ 获取插件的工具列表 Args: user_id: 用户ID plugin_name: 插件名称 Returns: 工具列表 """ client = self.get_client(user_id, plugin_name) if not client: raise ValueError(f"插件未加载: {plugin_name}") try: tools = await client.list_tools() return tools except Exception as e: logger.error(f"获取工具列表失败: {plugin_name}, 错误: {e}") raise async def test_plugin( self, user_id: str, plugin_name: str ) -> Dict[str, Any]: """ 测试插件连接 Args: user_id: 用户ID plugin_name: 插件名称 Returns: 测试结果 """ client = self.get_client(user_id, plugin_name) if not client: raise ValueError(f"插件未加载: {plugin_name}") return await client.test_connection() async def cleanup_all(self): """清理所有插件和资源""" # 停止后台清理任务 if self._cleanup_task: self._cleanup_task.cancel() try: await self._cleanup_task except asyncio.CancelledError: pass # 清理所有客户端 plugin_ids = list(self._clients.keys()) for plugin_id in plugin_ids: user_id = plugin_id.split(':', 1)[0] user_lock = await self._get_user_lock(user_id) async with user_lock: await self._unload_plugin_unsafe(plugin_id) # 关闭共享HTTP客户端 try: await self._shared_http_client.aclose() except Exception as e: logger.error(f"关闭共享HTTP客户端失败: {e}") logger.info("✅ 已清理所有MCP插件和资源") # 全局注册表实例 mcp_registry = MCPPluginRegistry()