From cb12cbdd3dc2518b4687cdd82557702b54f74410 Mon Sep 17 00:00:00 2001 From: Raven95676 Date: Wed, 16 Jul 2025 23:44:51 +0800 Subject: [PATCH 1/2] fix: managing MCP connections with AsyncExitStack --- astrbot/core/provider/func_tool_manager.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/astrbot/core/provider/func_tool_manager.py b/astrbot/core/provider/func_tool_manager.py index 8d79b1de8..f8cacff75 100644 --- a/astrbot/core/provider/func_tool_manager.py +++ b/astrbot/core/provider/func_tool_manager.py @@ -130,7 +130,7 @@ class MCPClient: timeout=cfg.get("timeout", 5), sse_read_timeout=cfg.get("sse_read_timeout", 60 * 5), ) - streams = await self._streams_context.__aenter__() + streams = await self.exit_stack.enter_async_context(self._streams_context) # Create a new client session self.session = await self.exit_stack.enter_async_context( @@ -148,7 +148,7 @@ class MCPClient: sse_read_timeout=sse_read_timeout, terminate_on_close=cfg.get("terminate_on_close", True), ) - read_s, write_s, _ = await self._streams_context.__aenter__() + read_s, write_s, _ = await self.exit_stack.enter_async_context(self._streams_context) # Create a new client session self.session = await self.exit_stack.enter_async_context( @@ -414,7 +414,7 @@ class FuncCall: try: # 关闭MCP连接 await self.mcp_client_dict[name].cleanup() - del self.mcp_client_dict[name] + self.mcp_client_dict.pop(name) except Exception as e: logger.info(f"清空 MCP 客户端资源 {name}: {e}。") # 移除关联的FuncTool From f86c8e8cab8b51125afa5939bf0075e39504e490 Mon Sep 17 00:00:00 2001 From: Raven95676 Date: Thu, 17 Jul 2025 23:17:23 +0800 Subject: [PATCH 2/2] perf: ensure MCP client termination in cleanup process --- astrbot/core/provider/func_tool_manager.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/astrbot/core/provider/func_tool_manager.py b/astrbot/core/provider/func_tool_manager.py index f8cacff75..cea2e4f38 100644 --- a/astrbot/core/provider/func_tool_manager.py +++ b/astrbot/core/provider/func_tool_manager.py @@ -356,12 +356,14 @@ class FuncCall: 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: import traceback traceback.print_exc() logger.error(f"初始化 MCP 客户端 {name} 失败: {e}") + finally: + # 无论如何都能清理 + await self._terminate_mcp_client(name) async def _init_mcp_client(self, name: str, config: dict) -> None: """初始化单个MCP客户端""" @@ -629,8 +631,3 @@ class FuncCall: def __repr__(self): return str(self.func_list) - - async def terminate(self): - for name in self.mcp_client_dict.keys(): - await self._terminate_mcp_client(name) - logger.debug(f"清理 MCP 客户端 {name} 资源")