update:1.更新mcp插件功能，目前只支持remote调用

2025-11-07 22:14:20 +08:00
parent 1e998920e3
commit 88115a45c5
26 changed files with 4088 additions and 138 deletions
@@ -0,0 +1,349 @@
+"""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()