auto-coder 1.0.0__py3-none-any.whl → 2.0.1__py3-none-any.whl

autocoder/common/{mcp_hub.py → mcp_tools/hub.py}

@@ -1,19 +1,27 @@
+"""
+MCP Tools Hub Module
+
+This module provides the McpHub class for managing MCP server connections
+and interactions. It handles server configuration, connection management,
+and marketplace operations.
+"""
+
 import os
 import json
 import asyncio
-import aiohttp
 import importlib
 import pkgutil
-import re
-import inspect
 from datetime import datetime, timedelta
-from typing import Dict, List, Optional, Any, Set, Optional, Tuple
+from typing import Dict, List, Optional, Any, Set, Tuple
 from pathlib import Path
 from pydantic import BaseModel, Field
 from loguru import logger
 from contextlib import AsyncExitStack
-from datetime import timedelta
-from autocoder.common.mcp_server_types import MarketplaceMCPServerItem
+
+from .types import (
+    McpTool, McpResource, McpResourceTemplate, McpServer,
+    MarketplaceMCPServerItem
+)
 
 try:
     from mcp import ClientSession
@@ -27,49 +35,13 @@ except ImportError:
     ClientSession = None
     logger.error("mcp is not installed(which requires python>=3.11), please install it by `pip install mcp`")
 
-class McpTool(BaseModel):
-    """Represents an MCP tool configuration"""
-
-    name: str
-    description: Optional[str] = None
-    input_schema: dict = Field(default_factory=dict)
-
-
-class McpResource(BaseModel):
-    """Represents an MCP resource configuration"""
-
-    uri: str
-    name: str
-    description: Optional[str] = None
-    mime_type: Optional[str] = None
-
-class McpResourceTemplate(BaseModel):
-    """Represents an MCP resource template"""
-
-    uri_template: str
-    name: str
-    description: Optional[str] = None
-    mime_type: Optional[str] = None
-
-
-class McpServer(BaseModel):
-    """Represents an MCP server configuration and status"""
-
-    name: str
-    config: str  # JSON string of server config
-    status: str = "disconnected"  # connected, disconnected, connecting
-    error: Optional[str] = None
-    tools: List[McpTool] = Field(default_factory=list)
-    resources: List[McpResource] = Field(default_factory=list)
-    resource_templates: List[McpResourceTemplate] = Field(default_factory=list)
-
 
 class McpConnection:
     """Represents an active MCP server connection"""
 
-    def __init__(self, server: McpServer, session: ClientSession):
+    def __init__(self, server: McpServer, session):
         self.server = server
-        self.session = session        
+        self.session = session
 
 
 def _generate_server_configs() -> Tuple[Dict[str, Any], Dict[str, str]]:
@@ -127,7 +99,13 @@ class McpHub:
 
     _instance = None
 
-    def __new__(cls, settings_path: Optional[str] = None):
+    def __new__(cls, settings_path: Optional[str] = None, marketplace_path: Optional[str] = None):
+        # For testing purposes, allow multiple instances with custom paths
+        if settings_path is not None or marketplace_path is not None:
+            instance = super(McpHub, cls).__new__(cls)
+            instance._initialized = False
+            return instance
+        
         if cls._instance is None:
             cls._instance = super(McpHub, cls).__new__(cls)
             cls._instance._initialized = False
@@ -178,7 +156,7 @@ class McpHub:
     def _read_marketplace(self) -> Dict[str, List[Dict[str, Any]]]:
         """Read marketplace file"""
         try:
-            with open(self.marketplace_path,"r", encoding="utf-8") as f:
+            with open(self.marketplace_path, "r", encoding="utf-8") as f:
                 return json.load(f)
         except Exception as e:
             logger.error(f"Failed to read marketplace: {e}")
@@ -216,7 +194,7 @@ class McpHub:
                 
             # Add the new item
             data = self._read_marketplace()
-            data["mcpServers"].append(item.dict())
+            data["mcpServers"].append(item.model_dump())
             
             # Write back to file
             with open(self.marketplace_path, "w", encoding="utf-8") as f:
@@ -341,12 +319,16 @@ class McpHub:
             logger.error(f"Failed to apply marketplace item: {e}")
             return False
 
-    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]) -> bool:
         """
         Add or update MCP server configuration in settings file.
 
         Args:
-            server_name_or_config: Name of the server or server configuration dictionary
+            name: Name of the server
+            config: Server configuration dictionary
+            
+        Returns:
+            bool: True if successful, False otherwise
         """
         settings = self._read_settings()
         settings["mcpServers"][name] = config
@@ -410,12 +392,16 @@ class McpHub:
         )
         transport_config = config.get("transport", {
             "type": "stdio",
-            "endpoint":""
+            "endpoint": ""
         })
 
         if transport_config["type"] not in ["stdio", "sse"]:
             raise ValueError(f"Invalid transport type: {transport_config['type']}")
         
+        # Check if MCP is available
+        if stdio_client is None or sse_client is None or ClientSession is None:
+            raise RuntimeError("MCP is not installed or not available")
+        
         if transport_config["type"] == "stdio":            
             # Setup transport parameters
             server_params = StdioServerParameters(
@@ -432,7 +418,6 @@ class McpHub:
             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(
@@ -455,11 +440,6 @@ class McpHub:
                 self.connections[name].server.error = error_msg
             return server
 
-        
-
-
-
-
     async def delete_connection(self, name: str) -> None:
         """
         Close and remove a server connection with proper cleanup
@@ -477,7 +457,6 @@ class McpHub:
                 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:
         """
@@ -577,14 +556,9 @@ class McpHub:
             if not connection:
                 return []
 
-            #       return await self.send_request(
-            #     types.ClientRequest(
-            #         types.PingRequest(
-            #             method="ping",
-            #         )
-            #     ),
-            #     types.EmptyResult,
-            # )
+            if mcp_types is None:
+                logger.warning("MCP types not available")
+                return []
 
             response = await connection.session.send_request(
                 mcp_types.ClientRequest(mcp_types.ListResourceTemplatesRequest(
@@ -617,7 +591,7 @@ class McpHub:
 
     async def call_tool(
         self, server_name: str, tool_name: str, tool_arguments: Optional[Dict] = None
-    ) -> mcp_types.CallToolResult:
+    ):
         """
         Call an MCP tool with arguments
         """
@@ -627,7 +601,7 @@ class McpHub:
 
         return await connection.session.call_tool(tool_name, tool_arguments or {})
 
-    async def read_resource(self, server_name: str, uri: str) -> mcp_types.ReadResourceResult:
+    async def read_resource(self, server_name: str, uri: str):
         """
         Read an MCP resource
         """
@@ -649,4 +623,4 @@ class McpHub:
         """
         Get all available server templates as JSON strings
         """
-        return MCP_SERVER_TEMPLATES
+        return MCP_SERVER_TEMPLATES 

autocoder/common/{mcp_server_install.py → mcp_tools/installer.py}

@@ -1,13 +1,22 @@
+"""
+MCP Tools Installer Module
+
+This module provides the McpServerInstaller class for installing and
+configuring MCP servers. It handles both built-in and external servers.
+"""
+
 import json
 import sys
 import os
 import time
 from typing import Dict, Any, Tuple, List
 from loguru import logger
-import asyncio
 
-from autocoder.common.mcp_hub import McpHub, MCP_BUILD_IN_SERVERS
-from autocoder.common.mcp_server_types import McpResponse, McpInstallRequest, InstallResult, ServerConfig, McpExternalServer
+from .hub import McpHub, MCP_BUILD_IN_SERVERS
+from .types import (
+    McpResponse, McpInstallRequest, InstallResult, ServerConfig, 
+    McpExternalServer, MarketplaceMCPServerItem
+)
 from autocoder.chat_auto_coder_lang import get_message_with_format
 
 
@@ -130,7 +139,6 @@ class McpServerInstaller:
             for s in external_servers:
                 if s.name == name:
                     if s.runtime == "python":
-                        # self.install_python_package(name)
                         template_config = {
                             "command": "python",
                             "args": [
@@ -138,7 +146,6 @@ class McpServerInstaller:
                             ],
                         }
                     elif s.runtime == "node":
-                        # self.install_node_package(name)
                         template_config = {
                             "command": "npx",
                             "args": [
@@ -155,14 +162,6 @@ class McpServerInstaller:
         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[str, Any]]:
@@ -175,16 +174,10 @@ class McpServerInstaller:
         # 取第一个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
     
-    def process_market_install_item(self, market_item) -> Tuple[str, Dict[str, Any]]:
+    def process_market_install_item(self, market_item: MarketplaceMCPServerItem) -> Tuple[str, Dict[str, Any]]:
         """Process a MarketplaceMCPServerItem into name and config"""
         name = market_item.name
         config = {
@@ -192,14 +185,6 @@ class McpServerInstaller:
             "args": market_item.args,
             "env": market_item.env
         }
-        
-        # Install package if needed
-        # if name.startswith("@") or market_item.command in ["npx", "npm"]:
-        #     for item in market_item.args:
-        #         if name in item:
-        #             self.install_node_package(item)
-        # elif market_item.command not in ["python", "node"]:
-        #     self.install_python_package(name)
             
         return name, config
 
@@ -215,6 +200,9 @@ class McpServerInstaller:
                 display_result = request.market_install_item.name
             else:
                 server_name_or_config = request.server_name_or_config
+                if server_name_or_config is None:
+                    raise ValueError("server_name_or_config is required")
+                    
                 display_result = server_name_or_config
 
                 # Try different parsing methods

autocoder/common/{mcp_server.py → mcp_tools/server.py}

@@ -1,64 +1,145 @@
+"""
+MCP Tools Server Module
+
+This module provides the McpServer class for handling MCP server operations
+in a thread-safe manner. It manages requests and responses through async queues.
+"""
+
 import asyncio
 from asyncio import Queue as AsyncQueue
 import threading
-from typing import List
-import os
-import time
+import signal
+import atexit
+from typing import Union
 from loguru import logger
-import json
 
-from autocoder.common.mcp_hub import McpHub, MCP_BUILD_IN_SERVERS, MarketplaceMCPServerItem
-from autocoder.common.mcp_tools import McpExecutor
+from .hub import McpHub, MCP_BUILD_IN_SERVERS
+from .executor import McpExecutor
+from .installer import McpServerInstaller
+from .types import (
+    McpRequest, McpInstallRequest, McpRemoveRequest, McpListRequest,
+    McpListRunningRequest, McpRefreshRequest, McpServerInfoRequest,
+    McpResponse, ServerInfo, InstallResult, RemoveResult, ListResult,
+    ListRunningResult, RefreshResult, QueryResult, ErrorResult,
+    StringResult, MarketplaceAddRequest, MarketplaceAddResult,
+    MarketplaceUpdateRequest, MarketplaceUpdateResult, MarketplaceMCPServerItem
+)
 from autocoder.utils.llms import get_single_llm
 from autocoder.chat_auto_coder_lang import get_message_with_format
-from autocoder.common.mcp_server_types import (
-    McpRequest, McpInstallRequest, McpRemoveRequest, McpListRequest, 
-    McpListRunningRequest, McpRefreshRequest, McpServerInfoRequest, 
-    McpResponse, ServerInfo, InstallResult, RemoveResult, ListResult, 
-    ListRunningResult, RefreshResult, QueryResult, ErrorResult, ServerConfig,StringResult,
-    ExternalServerInfo, McpExternalServer, MarketplaceAddRequest, MarketplaceAddResult,
-    MarketplaceUpdateRequest, MarketplaceUpdateResult
-)
-from autocoder.common.mcp_server_install import McpServerInstaller
+
 
 class McpServer:
+    """
+    MCP Server handling class that manages MCP operations in a thread-safe manner.
+    
+    This class provides an async interface for MCP operations while maintaining
+    thread safety through proper queue management.
+    """
+    
     def __init__(self):
         self._request_queue = AsyncQueue()
         self._response_queue = AsyncQueue()
         self._running = False
         self._task = None
         self._loop = None
+        self._thread = None
         self._installer = McpServerInstaller()
+        self._shutdown_event = threading.Event()
 
     def start(self):
+        """Start the MCP server in a separate thread"""
         if self._running:
             return
 
         self._running = True
+        self._shutdown_event.clear()
         self._loop = asyncio.new_event_loop()
-        threading.Thread(target=self._run_event_loop, daemon=True).start()
+        self._thread = threading.Thread(target=self._run_event_loop, daemon=True)
+        self._thread.start()
 
     def stop(self):
-        if self._running:
-            self._running = False
-            if self._loop:
-                self._loop.stop()
-                self._loop.close()
+        """Stop the MCP server"""
+        if not self._running:
+            return
+            
+        self._running = False
+        self._shutdown_event.set()
+        
+        if self._loop and self._loop.is_running():
+            # Cancel the task gracefully
+            if self._task and not self._task.done():
+                self._loop.call_soon_threadsafe(self._task.cancel)
+            
+            # Send None to break the request processing loop
+            try:
+                future = asyncio.run_coroutine_threadsafe(
+                    self._request_queue.put(None), self._loop
+                )
+                future.result(timeout=1.0)
+            except Exception:
+                pass
+            
+            # Stop the event loop
+            self._loop.call_soon_threadsafe(self._loop.stop)
+        
+        # Wait for thread to finish
+        if self._thread and self._thread.is_alive():
+            self._thread.join(timeout=2.0)
+            if self._thread.is_alive():
+                logger.warning("MCP server thread did not stop gracefully")
 
     def _run_event_loop(self):
+        """Run the event loop in the dedicated thread"""
+        if self._loop is None:
+            return
+            
         asyncio.set_event_loop(self._loop)
-        self._task = self._loop.create_task(self._process_request())
-        self._loop.run_forever()
+        
+        try:
+            self._task = self._loop.create_task(self._process_request())
+            self._loop.run_forever()
+        except KeyboardInterrupt:
+            logger.info("KeyboardInterrupt received in MCP server event loop")
+        except Exception as e:
+            logger.error(f"Error in MCP server event loop: {e}")
+        finally:
+            # Clean up
+            if self._task and not self._task.done():
+                self._task.cancel()
+                try:
+                    self._loop.run_until_complete(self._task)
+                except asyncio.CancelledError:
+                    pass
+            
+            # Close the loop
+            try:
+                self._loop.close()
+            except Exception as e:
+                logger.error(f"Error closing event loop: {e}")
 
     async def _process_request(self):
-        hub = McpHub()    
-
-        ## 可能会阻塞，当mcp settings.json 里有异常内容时    
-        await hub.initialize()
+        """
+        Process incoming requests and generate responses
+        """
+        hub = McpHub()
+        
+        try:
+            # May block if there are abnormal contents in mcp settings.json
+            await hub.initialize()
+        except Exception as e:
+            logger.error(f"Failed to initialize MCP hub: {e}")
+            return
 
         while self._running:
-            try:                
-                request = await self._request_queue.get()
+            try:
+                # Use timeout to allow checking _running flag
+                try:
+                    request = await asyncio.wait_for(
+                        self._request_queue.get(), timeout=1.0
+                    )
+                except asyncio.TimeoutError:
+                    continue
+                    
                 if request is None:
                     break
 
@@ -89,14 +170,14 @@ class McpServer:
 
                 elif isinstance(request, McpListRequest):
                     try:
-                        # Get built-in servers                        
+                        # Get built-in servers
                         builtin_servers = []
                         for name, config in MCP_BUILD_IN_SERVERS.items():
                             marketplace_item = MarketplaceMCPServerItem(
                                 name=name,
                                 description=f"Built-in: {name}",
                                 mcp_type="command",
-                                command=config.get("command", ""),
+                                command=config.get("command") or "",
                                 args=config.get("args", []),
                                 env=config.get("env", {})
                             )
@@ -156,10 +237,15 @@ class McpServer:
 
                 elif isinstance(request, McpServerInfoRequest):
                     try:
-                        llm = get_single_llm(request.model, product_mode=request.product_mode)
+                        llm = get_single_llm(request.model, product_mode=request.product_mode or "lite")
+                        if llm is None:
+                            raise ValueError("Failed to get LLM instance")
                         mcp_executor = McpExecutor(hub, llm)
                         result = mcp_executor.get_connected_servers_info()
-                        await self._response_queue.put(McpResponse(result=result, raw_result=StringResult(result=result)))
+                        await self._response_queue.put(McpResponse(
+                            result=result, 
+                            raw_result=StringResult(result=result)
+                        ))
                     except Exception as e:
                         import traceback
                         traceback.print_exc()
@@ -172,8 +258,7 @@ class McpServer:
                 elif isinstance(request, McpListRunningRequest):
                     try:
                         servers = hub.get_servers()
-                        running_servers = "\n".join(
-                            [f"- {server.name}" for server in servers])
+                        running_servers = "\n".join([f"- {server.name}" for server in servers])
                         result = running_servers if running_servers else ""
                         await self._response_queue.put(McpResponse(
                             result=result,
@@ -219,7 +304,7 @@ class McpServer:
                             name=request.name,
                             description=request.description,
                             mcp_type=request.mcp_type,
-                            command=request.command,
+                            command=request.command or "",
                             args=request.args or [],
                             env=request.env or {},
                             url=request.url or ""
@@ -264,7 +349,7 @@ class McpServer:
                             name=request.name,
                             description=request.description,
                             mcp_type=request.mcp_type,
-                            command=request.command,
+                            command=request.command or "",
                             args=request.args or [],
                             env=request.env or {},
                             url=request.url or ""
@@ -303,6 +388,7 @@ class McpServer:
                         ))
 
                 else:
+                    # Handle general MCP requests
                     if not request.query.strip():
                         await self._response_queue.put(McpResponse(
                             result="", 
@@ -316,8 +402,7 @@ class McpServer:
 
                     llm = get_single_llm(request.model, product_mode=request.product_mode)
                     mcp_executor = McpExecutor(hub, llm)
-                    conversations = [
-                        {"role": "user", "content": request.query}]
+                    conversations = [{"role": "user", "content": request.query}]
                     _, results = await mcp_executor.run(conversations)
 
                     if not results:
@@ -339,20 +424,53 @@ class McpServer:
                                 results=results
                             )
                         ))
+            except asyncio.CancelledError:
+                logger.info("MCP request processing was cancelled")
+                break
             except Exception as e:
-                await self._response_queue.put(McpResponse(
-                    result="", 
-                    error=get_message_with_format("mcp_error_title") + ": " + str(e),
-                    raw_result=ErrorResult(error=str(e))
-                ))
+                logger.error(f"Error in MCP request processing: {e}")
+                try:
+                    await self._response_queue.put(McpResponse(
+                        result="", 
+                        error=get_message_with_format("mcp_error_title") + ": " + str(e),
+                        raw_result=ErrorResult(error=str(e))
+                    ))
+                except Exception:
+                    pass
+        
+        # Clean up hub connections
+        try:
+            await hub.shutdown()
+        except Exception as e:
+            logger.error(f"Error shutting down MCP hub: {e}")
 
-    def send_request(self, request: McpRequest) -> McpResponse:
+    def send_request(self, request: Union[McpRequest, McpServerInfoRequest]) -> McpResponse:
+        """
+        Send a request to the MCP server and wait for response
+        
+        Args:
+            request: The request to send
+            
+        Returns:
+            McpResponse: The response from the server
+        """
+        if not self._running:
+            self.start()
+        
         async def _send():
             await self._request_queue.put(request)
             return await self._response_queue.get()
-
-        future = asyncio.run_coroutine_threadsafe(_send(), self._loop)
-        return future.result()
+        
+        if self._loop is None:
+            return McpResponse(result="", error="Event loop not available")
+        
+        try:
+            future = asyncio.run_coroutine_threadsafe(_send(), self._loop)
+            return future.result(timeout=30.0)  # Add timeout to prevent hanging
+        except asyncio.TimeoutError:
+            return McpResponse(result="", error="Request timeout")
+        except Exception as e:
+            return McpResponse(result="", error=f"Request failed: {str(e)}")
 
 
 # Global MCP server instance
@@ -360,10 +478,20 @@ _mcp_server = None
 
 
 def get_mcp_server():
+    """Get the global MCP server instance"""
     global _mcp_server
     if _mcp_server is None:
         _mcp_server = McpServer()
         _mcp_server.start()
+        # Register cleanup function
+        atexit.register(_cleanup_mcp_server)
     return _mcp_server
 
 
+def _cleanup_mcp_server():
+    """Clean up the global MCP server instance"""
+    global _mcp_server
+    if _mcp_server is not None:
+        logger.info("Cleaning up MCP server...")
+        _mcp_server.stop()
+        _mcp_server = None