From 873b7715f40b703bd50558cca39f826aaebf2c3a Mon Sep 17 00:00:00 2001
From: Soulter <905617992@qq.com>
Date: Sun, 23 Mar 2025 16:50:23 +0800
Subject: [PATCH] =?UTF-8?q?=F0=9F=8E=88=20perf:=20=E4=BC=98=E5=8C=96=20MCP?=
 =?UTF-8?q?=20Client=20=E5=BC=82=E6=AD=A5=20Event=20=E7=AE=A1=E7=90=86?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 astrbot/core/provider/func_tool_manager.py | 36 +++++++++++++++++++---
 1 file changed, 32 insertions(+), 4 deletions(-)

diff --git a/astrbot/core/provider/func_tool_manager.py b/astrbot/core/provider/func_tool_manager.py
index b795027da..fd4a14f7f 100644
--- a/astrbot/core/provider/func_tool_manager.py
+++ b/astrbot/core/provider/func_tool_manager.py
@@ -127,6 +127,7 @@ class FuncCall:
         """MCP 服务列表"""
         self.mcp_service_queue = asyncio.Queue()
         """用于外部控制 MCP 服务的启停"""
+        self.mcp_client_event: Dict[str, asyncio.Event] = {}
 
     def empty(self) -> bool:
         return len(self.func_list) == 0
@@ -218,7 +219,11 @@ class FuncCall:
         for name in mcp_server_json_obj.keys():
             cfg = mcp_server_json_obj[name]
             if cfg.get("active", True):
-                asyncio.create_task(self._init_mcp_client(name, cfg))
+                event = asyncio.Event()
+                asyncio.create_task(
+                    self._init_mcp_client_task_wrapper(name, cfg, event)
+                )
+                self.mcp_client_event[name] = event
 
     async def mcp_service_selector(self):
         """为了避免在不同异步任务中控制 MCP 服务导致的报错，整个项目统一通过这个 Task 来控制
@@ -237,17 +242,40 @@ class FuncCall:
             data = await self.mcp_service_queue.get()
             if data["type"] == "init":
                 if "name" in data:
+                    event = asyncio.Event()
                     asyncio.create_task(
-                        self._init_mcp_client(data["name"], data["cfg"])
+                        self._init_mcp_client_task_wrapper(
+                            data["name"], data["cfg"], event
+                        )
                     )
+                    self.mcp_client_event[data["name"]] = event
                 else:
                     await self._init_mcp_clients()
             elif data["type"] == "terminate":
                 if "name" in data:
-                    await self._terminate_mcp_client(data["name"])
+                    # await self._terminate_mcp_client(data["name"])
+                    if data["name"] in self.mcp_client_event:
+                        self.mcp_client_event[data["name"]].set()
+                        self.mcp_client_event.pop(data["name"], None)
                 else:
                     for name in self.mcp_client_dict.keys():
-                        await self._terminate_mcp_client(name)
+                        # await self._terminate_mcp_client(name)
+                        # self.mcp_client_event[name].set()
+                        if name in self.mcp_client_event:
+                            self.mcp_client_event[name].set()
+                            self.mcp_client_event.pop(name, None)
+
+    async def _init_mcp_client_task_wrapper(
+        self, name: str, cfg: dict, event: asyncio.Event
+    ) -> None:
+        """初始化 MCP 客户端的包装函数，用于捕获异常"""
+        try:
+            await self._init_mcp_client(name, cfg)
+            await event.wait()
+            logger.info(f"收到 MCP 客户端 {name} 终止信号")
+            await self._terminate_mcp_client(name)
+        except Exception as e:
+            logger.error(f"初始化 MCP 客户端 {name} 失败: {e}")
 
     async def _init_mcp_client(self, name: str, config: dict) -> None:
         """初始化单个MCP客户端"""