auto-coder 0.1.287__py3-none-any.whl → 0.1.289__py3-none-any.whl

autocoder/common/mcp_hub.py

@@ -2,16 +2,18 @@ import os
 import json
 import asyncio
 import aiohttp
-from datetime import datetime
+from datetime import datetime, timedelta
 from typing import Dict, List, Optional, Any, Set, Optional
 from pathlib import Path
 from pydantic import BaseModel, Field
 
 from mcp import ClientSession
 from mcp.client.stdio import stdio_client, StdioServerParameters
+from mcp.client.sse import sse_client
 import mcp.types as mcp_types
 from loguru import logger
-import time
+from contextlib import AsyncExitStack
+from datetime import timedelta
 
 
 class McpTool(BaseModel):
@@ -55,12 +57,9 @@ class McpServer(BaseModel):
 class McpConnection:
     """Represents an active MCP server connection"""
 
-    def __init__(self, server: McpServer, session: ClientSession, transport_manager):
+    def __init__(self, server: McpServer, session: ClientSession):
         self.server = server
-        self.session = session
-        self.transport_manager = (
-            transport_manager  # Will hold transport context manager
-        )
+        self.session = session        
 
 
 MCP_PERPLEXITY_SERVER = '''
@@ -81,6 +80,7 @@ MCP_BUILD_IN_SERVERS = {
     "perplexity": json.loads(MCP_PERPLEXITY_SERVER)["perplexity"]
 }
 
+
 class McpHub:
     """
     Manages MCP server connections and interactions.
@@ -105,6 +105,7 @@ class McpHub:
             self.settings_path = Path(settings_path)
         self.connections: Dict[str, McpConnection] = {}
         self.is_connecting = False
+        self.exit_stacks: Dict[str, AsyncExitStack] = {}
 
         # Ensure settings directory exists
         self.settings_path.parent.mkdir(parents=True, exist_ok=True)
@@ -116,26 +117,31 @@ class McpHub:
     def _write_default_settings(self):
         """Write default MCP settings file"""
         default_settings = {"mcpServers": {}}
-        with open(self.settings_path, "w",encoding="utf-8") as f:
+        with open(self.settings_path, "w", encoding="utf-8") as f:
             json.dump(default_settings, f, indent=2)
 
-    async def add_server_config(self, name: str, config:Dict[str,Any]) -> None:
+    async def add_server_config(self, name: str, config: Dict[str, Any]) -> None:
         """
         Add or update MCP server configuration in settings file.
 
         Args:
             server_name_or_config: Name of the server or server configuration dictionary
         """
-        try:
-            settings = self._read_settings()            
-            settings["mcpServers"][name] = config
-            with open(self.settings_path, "w",encoding="utf-8") as f:
-                json.dump(settings, f, indent=2, ensure_ascii=False)
-            await self.initialize()
-            logger.info(f"Added/updated MCP server config: {name}")
-        except Exception as e:
-            logger.error(f"Failed to add MCP server config: {e}")
-            raise
+        settings = self._read_settings()
+        settings["mcpServers"][name] = config
+        with open(self.settings_path, "w", encoding="utf-8") as f:
+            json.dump(settings, f, indent=2, ensure_ascii=False)
+        await self.initialize()
+        if name in self.connections:
+            if self.connections[name].server.status == "connected":
+                logger.info(f"Added/updated MCP server config: {name}")
+                return True
+            else:
+                logger.error(f"Failed to add MCP server config: {name}")
+                return False
+        else:
+            logger.error(f"Failed to add MCP server config: {name}")
+            return False
 
     async def remove_server_config(self, name: str) -> None:
         """
@@ -148,7 +154,7 @@ class McpHub:
             settings = self._read_settings()
             if name in settings["mcpServers"]:
                 del settings["mcpServers"][name]
-                with open(self.settings_path, "w",encoding="utf-8") as f:
+                with open(self.settings_path, "w", encoding="utf-8") as f:
                     json.dump(settings, f, indent=2, ensure_ascii=False)
                 logger.info(f"Removed MCP server config: {name}")
                 await self.initialize()
@@ -160,18 +166,14 @@ class McpHub:
 
     async def initialize(self):
         """Initialize MCP server connections from settings"""
-        try:
-            config = self._read_settings()
-            await self.update_server_connections(config.get("mcpServers", {}))
-        except Exception as e:
-            logger.error(f"Failed to initialize MCP servers: {e}")
-            raise
+        config = self._read_settings()
+        await self.update_server_connections(config.get("mcpServers", {}))
 
     def get_servers(self) -> List[McpServer]:
         """Get list of all configured servers"""
         return [conn.server for conn in self.connections.values()]
 
-    async def connect_to_server(self, name: str, config: dict) -> None:
+    async def connect_to_server(self, name: str, config: dict) -> McpServer:
         """
         Establish connection to an MCP server with proper resource management
         """
@@ -179,79 +181,90 @@ class McpHub:
         if name in self.connections:
             await self.delete_connection(name)
 
-        try:
-            server = McpServer(
-                name=name, config=json.dumps(config), status="connecting"
-            )
+        self.exit_stacks[name] = AsyncExitStack()
+        exit_stack = self.exit_stacks[name]
+
+        server = McpServer(
+            name=name, config=json.dumps(config), status="connecting"
+        )
+        transport_config = config.get("transport", {
+            "type": "stdio",
+            "endpoint":""
+        })
 
+        if transport_config["type"] not in ["stdio", "sse"]:
+            raise ValueError(f"Invalid transport type: {transport_config['type']}")
+        
+        if transport_config["type"] == "stdio":            
             # Setup transport parameters
             server_params = StdioServerParameters(
                 command=config["command"],
                 args=config.get("args", []),
                 env={**config.get("env", {}),
-                     "PATH": os.environ.get("PATH", "")},
+                    "PATH": os.environ.get("PATH", "")},
             )
 
-            # Create transport using context manager            
-            transport_manager = stdio_client(server_params)
-            transport = await transport_manager.__aenter__()
-            try:
-                session = await ClientSession(transport[0], transport[1]).__aenter__()
-                await session.initialize()
-
-                # Store connection with transport manager
-                connection = McpConnection(server, session, transport_manager)
-                self.connections[name] = connection
-
-                # Update server status and fetch capabilities
-                server.status = "connected"                
-                server.tools = await self._fetch_tools(name)                
-                server.resources = await self._fetch_resources(name)                
-                server.resource_templates = await self._fetch_resource_templates(name)
-
-            except Exception as e:
-                # Clean up transport if session initialization fails
-
-                await transport_manager.__aexit__(None, None, None)
-                raise
-
+            # Create transport using context manager
+            transport = await exit_stack.enter_async_context(stdio_client(server_params))
+        
+        if transport_config["type"] == "sse":
+            url = transport_config["endpoint"]
+            transport = await exit_stack.enter_async_context(sse_client(url))
+        
+        
+        try:
+            stdio, write = transport
+            session = await exit_stack.enter_async_context(
+                ClientSession(stdio, write, read_timeout_seconds=timedelta(seconds=60)))
+            await session.initialize()
+            connection = McpConnection(server, session)
+            self.connections[name] = connection
+            # Update server status and fetch capabilities
+            server.status = "connected"
+            server.tools = await self._fetch_tools(name)
+            server.resources = await self._fetch_resources(name)
+            server.resource_templates = await self._fetch_resource_templates(name)
+            return server
+        
         except Exception as e:
             error_msg = str(e)
             logger.error(f"Failed to connect to server {name}: {error_msg}")
             if name in self.connections:
                 self.connections[name].server.status = "disconnected"
                 self.connections[name].server.error = error_msg
-            raise
+            return server
+
+        
+
+
+
 
     async def delete_connection(self, name: str) -> None:
         """
         Close and remove a server connection with proper cleanup
         """
         if name in self.connections:
-            try:
-                connection = self.connections[name]
-                # Clean up in reverse order of creation
-                await connection.session.__aexit__(None, None, None)
-                await connection.transport_manager.__aexit__(None, None, None)
+            
+            try:                
                 del self.connections[name]
             except Exception as e:
                 logger.error(f"Error closing connection to {name}: {e}")
-                # Continue with deletion even if cleanup fails
-                if name in self.connections:
-                    del self.connections[name]
+
+            try:
+                exit_stack = self.exit_stacks[name]
+                await exit_stack.aclose()
+                del self.exit_stacks[name]
+            except Exception as e:
+                logger.error(f"Error cleaning up to {name}: {e}")
+                        
 
     async def refresh_server_connection(self, name: str) -> None:
         """
         Refresh a server connection
         """
-        try:
-            config = self._read_settings()
-            await self.delete_connection(name)
-            await self.connect_to_server(name, config.get("mcpServers", {}).get(name, {}))
-        except Exception as e:
-            logger.error(f"Failed to refresh MCP server {name}: {e}")
-            raise
-        
+        config = self._read_settings()
+        await self.delete_connection(name)
+        await self.connect_to_server(name, config.get("mcpServers", {}).get(name, {}))
 
     async def update_server_connections(self, new_servers: Dict[str, Any]) -> None:
         """
@@ -273,13 +286,22 @@ class McpHub:
 
                 if not current_conn:
                     # New server
-                    await self.connect_to_server(name, config)
-                    logger.info(f"Connected to new MCP server: {name}")
+                    server = await self.connect_to_server(name, config)
+                    if server.status == "connected":
+                        logger.info(f"Connected to new MCP server: {name}")
+                    else:
+                        logger.error(
+                            f"Failed to connect to new MCP server: {name}")
+
                 elif current_conn.server.config != json.dumps(config):
                     # Updated configuration
-                    await self.connect_to_server(name, config)
-                    logger.info(
-                        f"Reconnected MCP server with updated config: {name}")
+                    server = await self.connect_to_server(name, config)
+                    if server.status == "connected":
+                        logger.info(
+                            f"Reconnected MCP server with updated config: {name}")
+                    else:
+                        logger.error(
+                            f"Failed to reconnected MCP server with updated config: {name}")
 
         finally:
             self.is_connecting = False

autocoder/common/mcp_server.py

@@ -13,6 +13,7 @@ from pydantic import BaseModel
 import sys
 from loguru import logger
 from autocoder.utils.llms import get_single_llm
+from autocoder.chat_auto_coder_lang import get_message_with_format
 
 @dataclass
 class McpRequest:
@@ -157,55 +158,135 @@ class McpServer:
                 print(f"We have already updated the server configuration in ~/.autocoder/mcp/settings.json.\n")
                 print(f"After installation, you can restart the auto-coder.chat using the server.\033[0m\n")
 
+    def _deep_merge_dicts(self, dict1, dict2):
+        """
+        深度合并两个字典，包括嵌套的字典
+        dict1是基础字典，dict2是优先字典（当有冲突时保留dict2的值）
+        """
+        result = dict1.copy()
+        for key, value in dict2.items():
+            if key in result and isinstance(result[key], dict) and isinstance(value, dict):
+                # 如果两个字典都包含相同的键，并且两个值都是字典，则递归合并
+                result[key] = self._deep_merge_dicts(result[key], value)
+            else:
+                # 否则，使用dict2的值覆盖或添加到结果
+                result[key] = value
+        return result
+
+    def _parse_command_line_args(self, server_name_or_config: str) -> tuple[str, dict]:
+        """Parse command-line style arguments into name and config"""
+        name = ""
+        config = {}
+        args = server_name_or_config.strip().split()
+        i = 0
+        while i < len(args):
+            if args[i] == "--name" and i + 1 < len(args):
+                name = args[i + 1]
+                i += 2
+            elif args[i] == "--command" and i + 1 < len(args):
+                config["command"] = args[i + 1]
+                i += 2
+            elif args[i] == "--args":
+                config["args"] = []
+                i += 1
+                while i < len(args) and not args[i].startswith("--"):
+                    config["args"].append(args[i])
+                    i += 1
+            elif args[i] == "--env":
+                config["env"] = {}
+                i += 1
+                while i < len(args) and not args[i].startswith("--"):
+                    if "=" in args[i]:
+                        key, value = args[i].split("=", 1)
+                        config["env"][key] = value
+                    i += 1
+            else:
+                i += 1                
+        
+        template_config = {}
+        
+        if name in MCP_BUILD_IN_SERVERS:
+            template_config = MCP_BUILD_IN_SERVERS[name]
+        else:
+            # 查找外部server
+            external_servers = get_mcp_external_servers()
+            for s in external_servers:
+                if s.name == name:
+                    if s.runtime == "python":
+                        self._install_python_package(name)
+                        template_config = {
+                            "command": "python",
+                            "args": [
+                                "-m", name.replace("-", "_")
+                            ],
+                        }
+                    elif s.runtime == "node":
+                        self._install_node_package(name)
+                        template_config = {
+                            "command": "npx",
+                            "args": [
+                                "-y",
+                                "-g",
+                                name
+                            ]
+                        }
+                    break  
+
+        # 深度合并两个配置，以用户配置为主
+        config = self._deep_merge_dicts(template_config, config)
+
+        if not config.get("args") and (name.startswith("@") or config.get("command") in ["npx", "npm"]):
+            config["args"] = ["-y", "-g", name]            
+        
+        ## 如果有模板，则无需再次安装，处理模板的时候会自动安装
+        if not template_config:    
+            # Install package if needed
+            if name.startswith("@") or config.get("command") in ["npx", "npm"]:            
+                self._install_node_package(name)
+            else:
+                self._install_python_package(name)
+            
+        return name, config
+        
+    def _parse_json_config(self, server_name_or_config: str) -> tuple[str, dict]:
+        """Parse JSON configuration into name and config"""
+        raw_config = json.loads(server_name_or_config)
+        # 用户给了一个完整的配置
+        if "mcpServers" in raw_config:
+            raw_config = raw_config["mcpServers"]
+
+        # 取第一个server 配置
+        config = list(raw_config.values())[0]
+        name = list(raw_config.keys())[0]
+        if name.startswith("@") or config["command"] in ["npx", "npm"]:
+            for item in config["args"]:
+                if name in item:
+                    self._install_node_package(item)                    
+        else:
+            self._install_python_package(name)
+            
+        return name, config
+            
+
     async def _install_server(self, request: McpInstallRequest, hub: McpHub) -> McpResponse:
         """Install an MCP server with module dependency check"""
         name = ""
         config = {}
         try:
             server_name_or_config = request.server_name_or_config
-            try:
-                raw_config = json.loads(server_name_or_config)
-                # 用户给了一个完整的配置
-                if "mcpServers" in raw_config:
-                    raw_config = raw_config["mcpServers"]
-
-                # 取第一个server 配置
-                config = list(raw_config.values())[0]
-                name = list(raw_config.keys())[0]
-                if name.startswith("@") or config["command"] in ["npx","npm"]:
-                    for item in config["args"]:
-                        if name in item:
-                            self._install_node_package(item)                    
-                else:
-                    self._install_python_package(name)
-            except json.JSONDecodeError:
-                name = server_name_or_config.strip()
-                if name not in MCP_BUILD_IN_SERVERS:
-                    # 查找外部server
-                    external_servers = get_mcp_external_servers()
-                    for s in external_servers:
-                        if s.name == name:
-                            if s.runtime == "python":
-                                self._install_python_package(name)
-                                config = {
-                                    "command": "python",
-                                    "args": [
-                                        "-m", name.replace("-", "_")
-                                    ],
-                                }
-                            elif s.runtime == "node":
-                                self._install_node_package(name)
-                                config = {
-                                    "command": "npx",
-                                    "args": [
-                                        "-y",
-                                        "-g",
-                                        name
-                                    ]
-                                }
-                            break
-                else:
-                    config = MCP_BUILD_IN_SERVERS[name]
+            
+            # Try different parsing methods
+            if server_name_or_config.strip().startswith("--"):
+                # Command-line style arguments
+                name, config = self._parse_command_line_args(server_name_or_config)
+            else:
+                try:
+                    # Try parsing as JSON
+                    name, config = self._parse_json_config(server_name_or_config)
+                except json.JSONDecodeError:
+                    logger.error(f"Failed to parse JSON config: {server_name_or_config}")
+                    pass
+            
             if not name:
                 raise ValueError(
                     "MCP server name is not available in MCP_BUILD_IN_SERVERS or external servers")
@@ -214,10 +295,13 @@ class McpServer:
             if not config:
                 raise ValueError(f"MCP server {name} config is not available")
 
-            await hub.add_server_config(name, config)
-            return McpResponse(result=f"Successfully installed MCP server: {request.server_name_or_config}")
+            is_success = await hub.add_server_config(name, config)
+            if is_success:
+                return McpResponse(result=get_message_with_format("mcp_install_success", result=request.server_name_or_config))
+            else:
+                return McpResponse(result="", error=get_message_with_format("mcp_install_error", error="Failed to establish connection"))
         except Exception as e:
-            return McpResponse(result="", error=f"Failed to install MCP server: {str(e)}")
+            return McpResponse(result="", error=get_message_with_format("mcp_install_error", error=str(e)))
 
     async def _process_request(self):
         hub = McpHub()
@@ -236,9 +320,11 @@ class McpServer:
                 elif isinstance(request, McpRemoveRequest):
                     try:
                         await hub.remove_server_config(request.server_name)
-                        await self._response_queue.put(McpResponse(result=f"Successfully removed MCP server: {request.server_name}"))
+                        await self._response_queue.put(McpResponse(
+                            result=get_message_with_format("mcp_remove_success", result=request.server_name)))
                     except Exception as e:
-                        await self._response_queue.put(McpResponse(result="", error=f"Failed to remove MCP server: {str(e)}"))
+                        await self._response_queue.put(McpResponse(
+                            result="", error=get_message_with_format("mcp_remove_error", error=str(e))))
 
                 elif isinstance(request, McpListRequest):
                     try:
@@ -253,20 +339,22 @@ class McpServer:
 
                         # Combine results
                         all_servers = builtin_servers + external_list
-                        result = "Available MCP servers:\n" + \
-                            "\n".join(all_servers)
+                        result =  "\n".join(all_servers)
 
                         await self._response_queue.put(McpResponse(result=result))
                     except Exception as e:
-                        await self._response_queue.put(McpResponse(result="", error=f"Failed to list servers: {str(e)}"))
+                        await self._response_queue.put(McpResponse(
+                            result="", error=get_message_with_format("mcp_list_builtin_error", error=str(e))))
 
                 elif isinstance(request, McpListRunningRequest):
                     try:
                         running_servers = "\n".join(
                             [f"- {server.name}" for server in hub.get_servers()])
-                        await self._response_queue.put(McpResponse(result=running_servers))
+                        result = running_servers if running_servers else ""
+                        await self._response_queue.put(McpResponse(result=result))
                     except Exception as e:
-                        await self._response_queue.put(McpResponse(result="", error=f"Failed to list running servers: {str(e)}"))
+                        await self._response_queue.put(McpResponse(
+                            result="", error=get_message_with_format("mcp_list_running_error", error=str(e))))
 
                 elif isinstance(request, McpRefreshRequest):
                     try:
@@ -274,23 +362,36 @@ class McpServer:
                             await hub.refresh_server_connection(request.name)
                         else:
                             await hub.initialize()
-                        await self._response_queue.put(McpResponse(result="Successfully refreshed MCP server connections"))
+                        await self._response_queue.put(McpResponse(
+                            result=get_message_with_format("mcp_refresh_success")))
                     except Exception as e:
-                        await self._response_queue.put(McpResponse(result="", error=f"Failed to refresh MCP servers: {str(e)}"))
+                        await self._response_queue.put(McpResponse(
+                            result="", error=get_message_with_format("mcp_refresh_error", error=str(e))))
 
                 else:
-                    llm = get_single_llm(request.model,product_mode=request.product_mode)
+                    if not request.query.strip():
+                        await self._response_queue.put(McpResponse(
+                            result="", error=get_message_with_format("mcp_query_empty")))
+                        continue
+                        
+                    llm = get_single_llm(request.model, product_mode=request.product_mode)
                     mcp_executor = McpExecutor(hub, llm)
                     conversations = [
                         {"role": "user", "content": request.query}]
                     _, results = await mcp_executor.run(conversations)
+                    
                     if not results:
-                        await self._response_queue.put(McpResponse(result="[No Result]", error="No results"))
-                    results_str = "\n\n".join(
-                        mcp_executor.format_mcp_result(result) for result in results)
-                    await self._response_queue.put(McpResponse(result=results_str))
+                        await self._response_queue.put(McpResponse(
+                            result=get_message_with_format("mcp_error_title"), 
+                            error="No results"))
+                    else:
+                        results_str = "\n\n".join(
+                            mcp_executor.format_mcp_result(result) for result in results)
+                        await self._response_queue.put(McpResponse(
+                            result=get_message_with_format("mcp_response_title") + "\n" + results_str))
             except Exception as e:
-                await self._response_queue.put(McpResponse(result="", error=str(e)))
+                await self._response_queue.put(McpResponse(
+                    result="", error=get_message_with_format("mcp_error_title") + ": " + str(e)))
 
     def send_request(self, request: McpRequest) -> McpResponse:
         async def _send():