emdash-core 0.1.7__py3-none-any.whl

emdash_core/agent/mcp/manager.py

@@ -0,0 +1,496 @@
+"""MCP Server Manager for lifecycle and tool registry.
+
+This module provides the MCPServerManager class that handles:
+- Loading MCP config from .emdash/mcp.json
+- Starting/stopping MCP servers
+- Registering tools from all servers
+- Handling tool name collisions
+"""
+
+import atexit
+from pathlib import Path
+from typing import Optional
+
+from .client import GenericMCPClient, MCPResponse, MCPToolInfo, MCPError
+from .config import MCPConfigFile, MCPServerConfig, get_default_mcp_config_path, ensure_mcp_config
+from ...utils.logger import log
+
+
+class MCPServerManager:
+    """Manages lifecycle of multiple MCP servers.
+
+    This class is responsible for:
+    - Loading MCP server configurations from file
+    - Starting servers on demand (lazy initialization)
+    - Maintaining a unified tool registry from all servers
+    - Handling tool name collisions with prefixing
+    - Graceful shutdown of all servers
+
+    Example:
+        manager = MCPServerManager(config_path=Path(".emdash/mcp.json"))
+
+        # List all available tools from all servers
+        tools = manager.get_all_tools()
+
+        # Call a tool (server starts automatically if needed)
+        result = manager.call_tool("read_file", {"path": "/tmp/test.txt"})
+
+        # Cleanup
+        manager.shutdown_all()
+    """
+
+    def __init__(self, config_path: Optional[Path] = None, repo_root: Optional[Path] = None):
+        """Initialize the MCP server manager.
+
+        Args:
+            config_path: Path to mcp.json config file. If None, uses default.
+            repo_root: Repository root for default config path resolution.
+        """
+        self.repo_root = repo_root or Path.cwd()
+        self.config_path = config_path or get_default_mcp_config_path(self.repo_root)
+
+        self._config: Optional[MCPConfigFile] = None
+        self._clients: dict[str, GenericMCPClient] = {}
+        self._tool_registry: dict[str, tuple[str, MCPToolInfo]] = {}  # tool_name -> (server_name, tool_info)
+        self._started = False
+
+        # Register atexit handler for cleanup
+        atexit.register(self._cleanup)
+
+    def _cleanup(self) -> None:
+        """Cleanup handler for atexit."""
+        try:
+            self.shutdown_all()
+        except Exception as e:
+            log.warning(f"Error during MCP cleanup: {e}")
+
+    def load_config(self) -> MCPConfigFile:
+        """Load MCP configuration from file, creating default if needed.
+
+        Returns:
+            MCPConfigFile with loaded server configurations
+        """
+        if self._config is None:
+            self._config = ensure_mcp_config(self.config_path)
+        return self._config
+
+    def reload_config(self) -> MCPConfigFile:
+        """Reload configuration from file.
+
+        Returns:
+            MCPConfigFile with fresh configuration
+        """
+        self._config = MCPConfigFile.load(self.config_path)
+        return self._config
+
+    def save_config(self) -> None:
+        """Save current configuration to file."""
+        if self._config:
+            self._config.save(self.config_path)
+
+    def get_enabled_servers(self) -> list[MCPServerConfig]:
+        """Get list of enabled server configurations.
+
+        Returns:
+            List of enabled MCPServerConfig instances
+        """
+        config = self.load_config()
+        return config.get_enabled_servers()
+
+    def _resolve_args(self, args: list[str], resolved_env: dict[str, str]) -> list[str]:
+        """Resolve ${VAR} placeholders in args using resolved env.
+
+        Args:
+            args: List of argument strings
+            resolved_env: Already-resolved environment variables
+
+        Returns:
+            List of args with placeholders resolved
+        """
+        import re
+        pattern = re.compile(r"\$\{([^}]+)\}")
+        resolved = []
+        for arg in args:
+            def replace_var(m):
+                var_name = m.group(1)
+                return resolved_env.get(var_name, "")
+            resolved.append(pattern.sub(replace_var, arg))
+        return resolved
+
+    def start_server(self, name: str) -> GenericMCPClient:
+        """Start an MCP server by name.
+
+        Args:
+            name: Server name from configuration
+
+        Returns:
+            Started GenericMCPClient instance
+
+        Raises:
+            MCPError: If server not found in config or fails to start
+        """
+        # Check if already running
+        if name in self._clients and self._clients[name].is_running:
+            return self._clients[name]
+
+        config = self.load_config()
+        server_config = config.get_server(name)
+
+        if not server_config:
+            raise MCPError(f"MCP server '{name}' not found in configuration")
+
+        if not server_config.enabled:
+            raise MCPError(f"MCP server '{name}' is disabled")
+
+        # Resolve env vars first, then args
+        resolved_env = server_config.get_resolved_env()
+        resolved_args = self._resolve_args(server_config.args, resolved_env)
+
+        # Create and start client
+        client = GenericMCPClient(
+            name=server_config.name,
+            command=server_config.command,
+            args=resolved_args,
+            env=resolved_env,
+            timeout=server_config.timeout,
+        )
+
+        try:
+            client.start()
+            self._clients[name] = client
+
+            # Register tools from this server
+            self._register_server_tools(name, client)
+
+            return client
+        except Exception as e:
+            log.error(f"Failed to start MCP server '{name}': {e}")
+            raise
+
+    def stop_server(self, name: str) -> bool:
+        """Stop an MCP server by name.
+
+        Args:
+            name: Server name
+
+        Returns:
+            True if server was stopped, False if not running
+        """
+        if name not in self._clients:
+            return False
+
+        client = self._clients[name]
+        client.stop()
+        del self._clients[name]
+
+        # Unregister tools from this server
+        self._unregister_server_tools(name)
+
+        return True
+
+    def start_all_enabled(self) -> list[str]:
+        """Start all enabled MCP servers.
+
+        Returns:
+            List of server names that were started successfully
+        """
+        started = []
+        for server_config in self.get_enabled_servers():
+            try:
+                self.start_server(server_config.name)
+                started.append(server_config.name)
+            except MCPError as e:
+                log.warning(f"Failed to start MCP server '{server_config.name}': {e}")
+
+        self._started = True
+        return started
+
+    def shutdown_all(self) -> None:
+        """Stop all running MCP servers."""
+        for name in list(self._clients.keys()):
+            try:
+                self.stop_server(name)
+            except Exception as e:
+                log.warning(f"Error stopping MCP server '{name}': {e}")
+
+        self._clients.clear()
+        self._tool_registry.clear()
+        self._started = False
+
+    def _register_server_tools(self, server_name: str, client: GenericMCPClient) -> None:
+        """Register tools from a server into the unified registry.
+
+        Handles name collisions by prefixing with server name.
+
+        Args:
+            server_name: Name of the server
+            client: Started client to get tools from
+        """
+        try:
+            tools = client.list_tools()
+        except MCPError as e:
+            log.warning(f"Failed to list tools from '{server_name}': {e}")
+            return
+
+        for tool in tools:
+            # Check for name collision
+            if tool.name in self._tool_registry:
+                existing_server, _ = self._tool_registry[tool.name]
+                # Use prefixed name for collision
+                prefixed_name = f"{server_name}_{tool.name}"
+                log.warning(
+                    f"Tool name collision: '{tool.name}' exists in '{existing_server}'. "
+                    f"Registering as '{prefixed_name}' for server '{server_name}'."
+                )
+                self._tool_registry[prefixed_name] = (server_name, tool)
+            else:
+                self._tool_registry[tool.name] = (server_name, tool)
+
+        log.info(f"Registered {len(tools)} tools from MCP server '{server_name}'")
+
+    def _unregister_server_tools(self, server_name: str) -> None:
+        """Remove tools from a server from the registry.
+
+        Args:
+            server_name: Name of the server
+        """
+        to_remove = [
+            name for name, (srv, _) in self._tool_registry.items()
+            if srv == server_name
+        ]
+        for name in to_remove:
+            del self._tool_registry[name]
+
+    def get_all_tools(self) -> list[tuple[str, str, MCPToolInfo]]:
+        """Get all registered tools from all servers.
+
+        Starts all enabled servers if not already started.
+
+        Returns:
+            List of (tool_name, server_name, MCPToolInfo) tuples
+        """
+        if not self._started:
+            self.start_all_enabled()
+
+        return [
+            (name, server_name, tool_info)
+            for name, (server_name, tool_info) in self._tool_registry.items()
+        ]
+
+    def get_tool(self, name: str) -> Optional[tuple[str, MCPToolInfo]]:
+        """Get a tool by name.
+
+        Args:
+            name: Tool name (may be prefixed for collisions)
+
+        Returns:
+            (server_name, MCPToolInfo) tuple or None if not found
+        """
+        if not self._started:
+            self.start_all_enabled()
+
+        return self._tool_registry.get(name)
+
+    def call_tool(self, tool_name: str, arguments: dict) -> MCPResponse:
+        """Call a tool by name.
+
+        Automatically routes to the correct server.
+
+        Args:
+            tool_name: Tool name (may be prefixed)
+            arguments: Tool arguments
+
+        Returns:
+            MCPResponse from the tool call
+
+        Raises:
+            MCPError: If tool not found or call fails
+        """
+        tool_info = self.get_tool(tool_name)
+        if not tool_info:
+            raise MCPError(f"Tool '{tool_name}' not found in any MCP server")
+
+        server_name, mcp_tool = tool_info
+
+        # Get the client, start if needed
+        if server_name not in self._clients:
+            self.start_server(server_name)
+
+        client = self._clients[server_name]
+
+        # Use original tool name (strip prefix if present)
+        original_name = mcp_tool.name
+
+        return client.call_tool(original_name, arguments)
+
+    def add_server(self, config: MCPServerConfig) -> None:
+        """Add a new server to configuration.
+
+        Args:
+            config: Server configuration to add
+        """
+        self.load_config()
+        self._config.add_server(config)
+        self.save_config()
+        log.info(f"Added MCP server '{config.name}' to configuration")
+
+    def remove_server(self, name: str) -> bool:
+        """Remove a server from configuration.
+
+        Also stops the server if running.
+
+        Args:
+            name: Server name to remove
+
+        Returns:
+            True if removed, False if not found
+        """
+        # Stop if running
+        if name in self._clients:
+            self.stop_server(name)
+
+        self.load_config()
+        if self._config.remove_server(name):
+            self.save_config()
+            log.info(f"Removed MCP server '{name}' from configuration")
+            return True
+        return False
+
+    def enable_server(self, name: str) -> bool:
+        """Enable a server.
+
+        Args:
+            name: Server name
+
+        Returns:
+            True if enabled, False if not found
+        """
+        self.load_config()
+        server = self._config.get_server(name)
+        if server:
+            server.enabled = True
+            self.save_config()
+            return True
+        return False
+
+    def disable_server(self, name: str) -> bool:
+        """Disable a server.
+
+        Also stops the server if running.
+
+        Args:
+            name: Server name
+
+        Returns:
+            True if disabled, False if not found
+        """
+        # Stop if running
+        if name in self._clients:
+            self.stop_server(name)
+
+        self.load_config()
+        server = self._config.get_server(name)
+        if server:
+            server.enabled = False
+            self.save_config()
+            return True
+        return False
+
+    def list_servers(self) -> list[dict]:
+        """List all configured servers with their status.
+
+        Returns:
+            List of server info dicts with keys:
+            - name: Server name
+            - command: Server command
+            - enabled: Whether enabled
+            - running: Whether currently running
+            - tool_count: Number of tools (if running)
+        """
+        self.load_config()
+        servers = []
+
+        for name, server_config in self._config.servers.items():
+            is_running = name in self._clients and self._clients[name].is_running
+            tool_count = len([
+                1 for _, (srv, _) in self._tool_registry.items()
+                if srv == name
+            ]) if is_running else 0
+
+            servers.append({
+                "name": name,
+                "command": server_config.command,
+                "args": server_config.args,
+                "enabled": server_config.enabled,
+                "running": is_running,
+                "tool_count": tool_count,
+            })
+
+        return servers
+
+    def describe_server(self, name: str) -> Optional[dict]:
+        """Get detailed information about a server including its tools.
+
+        Args:
+            name: Server name
+
+        Returns:
+            Dict with server details and tools, or None if not found
+        """
+        self.load_config()
+        server_config = self._config.get_server(name)
+
+        if not server_config:
+            return None
+
+        is_running = name in self._clients and self._clients[name].is_running
+
+        # Get tools for this server
+        tools = []
+        if is_running:
+            for tool_name, (srv, tool_info) in self._tool_registry.items():
+                if srv == name:
+                    tools.append({
+                        "name": tool_name,
+                        "original_name": tool_info.name,
+                        "description": tool_info.description,
+                        "input_schema": tool_info.input_schema,
+                    })
+
+        return {
+            "name": name,
+            "command": server_config.command,
+            "args": server_config.args,
+            "env": server_config.env,
+            "enabled": server_config.enabled,
+            "timeout": server_config.timeout,
+            "running": is_running,
+            "tools": tools,
+        }
+
+
+# Global manager instance (lazy initialization)
+_manager: Optional[MCPServerManager] = None
+
+
+def get_mcp_manager(config_path: Optional[Path] = None) -> MCPServerManager:
+    """Get or create the global MCP manager instance.
+
+    Args:
+        config_path: Optional path to override default config location
+
+    Returns:
+        MCPServerManager instance
+    """
+    global _manager
+    if _manager is None:
+        _manager = MCPServerManager(config_path=config_path)
+    return _manager
+
+
+def reset_mcp_manager() -> None:
+    """Reset the global MCP manager (for testing)."""
+    global _manager
+    if _manager:
+        _manager.shutdown_all()
+    _manager = None

emdash_core/agent/mcp/tool_factory.py

@@ -0,0 +1,213 @@
+"""Tool factory for creating BaseTool instances from MCP servers.
+
+This module provides functions to dynamically create agent tools
+from MCP server tool definitions.
+"""
+
+import json
+from typing import Any, Optional
+
+from ..tools.base import BaseTool, ToolResult, ToolCategory
+from .manager import MCPServerManager
+from .client import MCPToolInfo, MCPError
+from ...utils.logger import log
+
+
+class MCPDynamicTool(BaseTool):
+    """A tool dynamically created from an MCP server tool definition.
+
+    This allows MCP tools to be used seamlessly alongside native tools
+    in the agent toolkit.
+    """
+
+    def __init__(
+        self,
+        mcp_manager: MCPServerManager,
+        tool_name: str,
+        tool_info: MCPToolInfo,
+        server_name: str,
+        connection: Any = None,
+    ):
+        """Initialize the dynamic tool.
+
+        Args:
+            mcp_manager: MCP server manager for calling tools
+            tool_name: Tool name (may be prefixed)
+            tool_info: Tool info from MCP server
+            server_name: Name of the source server
+            connection: Optional Kuzu connection (for compatibility)
+        """
+        super().__init__(connection)
+        self._mcp_manager = mcp_manager
+        self._tool_name = tool_name
+        self._tool_info = tool_info
+        self._server_name = server_name
+
+        # Set BaseTool attributes
+        self.name = tool_name
+        self.description = tool_info.description or f"MCP tool: {tool_name}"
+        self.category = self._infer_category(tool_name, tool_info)
+
+    def _infer_category(self, name: str, info: MCPToolInfo) -> ToolCategory:
+        """Infer tool category from name and description.
+
+        Args:
+            name: Tool name
+            info: Tool info
+
+        Returns:
+            Inferred ToolCategory
+        """
+        name_lower = name.lower()
+        desc_lower = (info.description or "").lower()
+
+        # Search-related
+        if any(kw in name_lower for kw in ["search", "find", "query", "grep"]):
+            return ToolCategory.SEARCH
+
+        # Traversal-related
+        if any(kw in name_lower for kw in ["expand", "caller", "callee", "dependency", "neighbor"]):
+            return ToolCategory.TRAVERSAL
+
+        # Analytics-related
+        if any(kw in name_lower for kw in ["pagerank", "community", "importance", "metric"]):
+            return ToolCategory.ANALYTICS
+
+        # History-related (PRs, commits)
+        if any(kw in name_lower for kw in ["pr", "commit", "history", "github"]):
+            return ToolCategory.HISTORY
+
+        # Default to search
+        return ToolCategory.SEARCH
+
+    def execute(self, **kwargs) -> ToolResult:
+        """Execute the MCP tool.
+
+        Args:
+            **kwargs: Tool arguments
+
+        Returns:
+            ToolResult with tool output
+        """
+        try:
+            response = self._mcp_manager.call_tool(self._tool_name, kwargs)
+
+            if response.is_error:
+                return ToolResult.error_result(
+                    f"MCP tool error: {response.get_text()}",
+                )
+
+            # Parse response content
+            result_data = self._parse_response(response)
+
+            return ToolResult.success_result(
+                data=result_data,
+                metadata={
+                    "server": self._server_name,
+                    "tool": self._tool_name,
+                },
+            )
+
+        except MCPError as e:
+            return ToolResult.error_result(
+                f"MCP error: {str(e)}",
+                suggestions=["Check if the MCP server is running"],
+            )
+        except Exception as e:
+            log.exception(f"MCP tool execution error: {self._tool_name}")
+            return ToolResult.error_result(
+                f"Tool execution failed: {str(e)}",
+            )
+
+    def _parse_response(self, response) -> dict:
+        """Parse MCP response into result data.
+
+        Args:
+            response: MCPResponse
+
+        Returns:
+            Parsed result dict
+        """
+        result = {"content": [], "text": ""}
+
+        for item in response.content:
+            if item.get("type") == "text":
+                text = item.get("text", "")
+                result["text"] += text + "\n"
+
+                # Try to parse as JSON
+                try:
+                    parsed = json.loads(text)
+                    if isinstance(parsed, dict):
+                        result.update(parsed)
+                    elif isinstance(parsed, list):
+                        result["results"] = parsed
+                except json.JSONDecodeError:
+                    pass
+
+            result["content"].append(item)
+
+        result["text"] = result["text"].strip()
+        return result
+
+    def get_schema(self) -> dict:
+        """Get OpenAI function calling schema.
+
+        Returns:
+            OpenAI function schema dict
+        """
+        input_schema = self._tool_info.input_schema or {}
+
+        # Convert MCP schema to OpenAI format
+        parameters = {
+            "type": "object",
+            "properties": input_schema.get("properties", {}),
+            "required": input_schema.get("required", []),
+        }
+
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": parameters,
+            },
+        }
+
+
+def create_tools_from_mcp(
+    manager: MCPServerManager,
+    connection: Any = None,
+) -> list[BaseTool]:
+    """Create BaseTool instances from all MCP server tools.
+
+    This function queries all enabled MCP servers and creates
+    dynamic tool wrappers for each tool they provide.
+
+    Args:
+        manager: MCP server manager
+        connection: Optional Kuzu connection
+
+    Returns:
+        List of MCPDynamicTool instances
+    """
+    tools = []
+
+    try:
+        all_tools = manager.get_all_tools()
+
+        for tool_name, server_name, tool_info in all_tools:
+            tool = MCPDynamicTool(
+                mcp_manager=manager,
+                tool_name=tool_name,
+                tool_info=tool_info,
+                server_name=server_name,
+                connection=connection,
+            )
+            tools.append(tool)
+            log.debug(f"Created MCP tool: {tool_name} from {server_name}")
+
+    except Exception as e:
+        log.warning(f"Failed to create tools from MCP: {e}")
+
+    return tools