From 116f1cb4715328cbe2c58e83151ec29c6e5d9514 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9C=B1=E6=BD=AE?= <zhuchaowe@users.noreply.github.com>
Date: Thu, 27 Nov 2025 16:01:30 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E5=B9=B6=E5=8F=91?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 agent/custom_mcp_manager.py | 38 +++++++++++++++++++++++++++++++------
 utils/fastapi_utils.py      |  8 +++++---
 2 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/agent/custom_mcp_manager.py b/agent/custom_mcp_manager.py
index e8ce4dc..fa31e47 100644
--- a/agent/custom_mcp_manager.py
+++ b/agent/custom_mcp_manager.py
@@ -151,16 +151,27 @@ class CustomMCPManager:
     async def init_config_async(self, config: Dict):
         tools: list = []
         mcp_servers = config['mcpServers']
+
+        # 并发连接所有MCP服务器
+        connection_tasks = []
         for server_name in mcp_servers:
             client = CustomMCPClient()
             server = mcp_servers[server_name]
-            await client.connection_server(mcp_server_name=server_name,
-                                           mcp_server=server)  # Attempt to connect to the server
+            # 创建连接任务
+            task = self._connect_and_store_client(client, server_name, server)
+            connection_tasks.append(task)
 
-            client_id = server_name + '_' + str(
-                uuid.uuid4())  # To allow the same server name be used across different running agents
-            client.client_id = client_id  # Ensure client_id is set on the client instance
-            self.clients[client_id] = client  # Add to clients dict after successful connection
+        # 并发执行所有连接任务
+        connected_clients = await asyncio.gather(*connection_tasks, return_exceptions=True)
+
+        # 处理连接结果并为每个客户端创建工具
+        for result in connected_clients:
+            if isinstance(result, Exception):
+                logger.error(f'Failed to connect MCP server: {result}')
+                continue
+
+            client, server_name = result
+            client_id = client.client_id
             for tool in client.tools:
                 """MCP tool example:
                 {
@@ -261,6 +272,21 @@ class CustomMCPManager:
 
         return tools
 
+    async def _connect_and_store_client(self, client, server_name, server):
+        """辅助方法：连接MCP服务器并存储客户端"""
+        try:
+            await client.connection_server(mcp_server_name=server_name,
+                                           mcp_server=server)  # Attempt to connect to the server
+
+            client_id = server_name + '_' + str(
+                uuid.uuid4())  # To allow the same server name be used across different running agents
+            client.client_id = client_id  # Ensure client_id is set on the client instance
+            self.clients[client_id] = client  # Add to clients dict after successful connection
+            return client, server_name
+        except Exception as e:
+            logger.error(f'Failed to connect MCP server {server_name}: {e}')
+            raise e
+
     def create_tool_class(self, register_name, register_client_id, tool_name, tool_desc, tool_parameters):
 
         class ToolClass(BaseTool):
diff --git a/utils/fastapi_utils.py b/utils/fastapi_utils.py
index c7a3160..f96957d 100644
--- a/utils/fastapi_utils.py
+++ b/utils/fastapi_utils.py
@@ -445,12 +445,14 @@ async def call_guideline_llm(chat_history: str, guidelines_text: str, terms:str,
 
 def _get_optimal_batch_size(guidelines_count: int) -> int:
     """根据guidelines数量决定最优批次数量（并发数）"""
-    if guidelines_count <= 10:
+    if guidelines_count <= 5:
         return 1
-    elif guidelines_count <= 20:
+    elif guidelines_count <= 10:
         return 2
-    elif guidelines_count <= 30:
+    elif guidelines_count <= 20:
         return 3
+    elif guidelines_count <= 30:
+        return 4
     else:
         return 5