kollabor 0.4.9__py3-none-any.whl → 0.4.15__py3-none-any.whl

core/llm/mcp_integration.py

@@ -2,6 +2,11 @@
 
 Provides integration with MCP servers for tool discovery and execution,
 enabling the LLM to interact with external tools and services.
+
+Implements MCP JSON-RPC 2.0 protocol over stdio for:
+- Server initialization handshake
+- tools/list for tool discovery
+- tools/call for tool execution
 """
 
 import asyncio
@@ -10,32 +15,208 @@ import logging
 import re
 import shlex
 import subprocess
+import uuid
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
 
 logger = logging.getLogger(__name__)
 
 
+class MCPServerConnection:
+    """Manages a connection to an MCP server via stdio."""
+
+    def __init__(self, server_name: str, command: str):
+        self.server_name = server_name
+        self.command = command
+        self.process: Optional[asyncio.subprocess.Process] = None
+        self.initialized = False
+        self._read_buffer = ""
+
+    async def connect(self) -> bool:
+        """Start the MCP server process."""
+        try:
+            command_parts = shlex.split(self.command)
+            self.process = await asyncio.create_subprocess_exec(
+                *command_parts,
+                stdin=asyncio.subprocess.PIPE,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE
+            )
+            logger.info(f"Started MCP server process: {self.server_name}")
+            return True
+        except Exception as e:
+            logger.error(f"Failed to start MCP server {self.server_name}: {e}")
+            return False
+
+    async def initialize(self) -> bool:
+        """Send MCP initialize request and wait for response."""
+        if not self.process:
+            return False
+
+        request = {
+            "jsonrpc": "2.0",
+            "id": str(uuid.uuid4()),
+            "method": "initialize",
+            "params": {
+                "protocolVersion": "2024-11-05",
+                "capabilities": {},
+                "clientInfo": {
+                    "name": "kollabor-cli",
+                    "version": "1.0.0"
+                }
+            }
+        }
+
+        response = await self._send_request(request)
+        if response and "result" in response:
+            self.initialized = True
+            logger.info(f"MCP server {self.server_name} initialized")
+
+            # Send initialized notification
+            notification = {
+                "jsonrpc": "2.0",
+                "method": "notifications/initialized"
+            }
+            await self._send_notification(notification)
+            return True
+
+        logger.warning(f"MCP server {self.server_name} initialization failed")
+        return False
+
+    async def list_tools(self) -> List[Dict[str, Any]]:
+        """Request tools list from server."""
+        if not self.initialized:
+            return []
+
+        request = {
+            "jsonrpc": "2.0",
+            "id": str(uuid.uuid4()),
+            "method": "tools/list",
+            "params": {}
+        }
+
+        response = await self._send_request(request)
+        if response and "result" in response:
+            tools = response["result"].get("tools", [])
+            logger.info(f"Got {len(tools)} tools from {self.server_name}")
+            return tools
+
+        return []
+
+    async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
+        """Call a tool on the server."""
+        if not self.initialized:
+            return {"error": "Server not initialized"}
+
+        request = {
+            "jsonrpc": "2.0",
+            "id": str(uuid.uuid4()),
+            "method": "tools/call",
+            "params": {
+                "name": tool_name,
+                "arguments": arguments
+            }
+        }
+
+        response = await self._send_request(request)
+        if response:
+            if "result" in response:
+                return response["result"]
+            elif "error" in response:
+                return {"error": response["error"].get("message", "Unknown error")}
+
+        return {"error": "No response from server"}
+
+    async def _send_request(self, request: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        """Send JSON-RPC request and wait for response."""
+        if not self.process or not self.process.stdin:
+            return None
+
+        try:
+            message = json.dumps(request) + "\n"
+            self.process.stdin.write(message.encode())
+            await self.process.stdin.drain()
+
+            # Read response with timeout
+            response = await asyncio.wait_for(self._read_response(), timeout=30)
+            return response
+        except asyncio.TimeoutError:
+            logger.warning(f"Timeout waiting for response from {self.server_name}")
+            return None
+        except Exception as e:
+            logger.error(f"Error sending request to {self.server_name}: {e}")
+            return None
+
+    async def _send_notification(self, notification: Dict[str, Any]) -> None:
+        """Send JSON-RPC notification (no response expected)."""
+        if not self.process or not self.process.stdin:
+            return
+
+        try:
+            message = json.dumps(notification) + "\n"
+            self.process.stdin.write(message.encode())
+            await self.process.stdin.drain()
+        except Exception as e:
+            logger.error(f"Error sending notification to {self.server_name}: {e}")
+
+    async def _read_response(self) -> Optional[Dict[str, Any]]:
+        """Read a JSON-RPC response from stdout."""
+        if not self.process or not self.process.stdout:
+            return None
+
+        while True:
+            # Check if we have a complete message in buffer
+            if "\n" in self._read_buffer:
+                line, self._read_buffer = self._read_buffer.split("\n", 1)
+                if line.strip():
+                    try:
+                        return json.loads(line)
+                    except json.JSONDecodeError:
+                        continue
+
+            # Read more data
+            chunk = await self.process.stdout.read(4096)
+            if not chunk:
+                return None
+            self._read_buffer += chunk.decode()
+
+    async def close(self) -> None:
+        """Close the server connection."""
+        if self.process:
+            try:
+                self.process.terminate()
+                await asyncio.wait_for(self.process.wait(), timeout=5)
+            except asyncio.TimeoutError:
+                self.process.kill()
+            except Exception as e:
+                logger.warning(f"Error closing MCP server {self.server_name}: {e}")
+            self.process = None
+            self.initialized = False
+            logger.info(f"Closed MCP server connection: {self.server_name}")
+
+
 class MCPIntegration:
     """MCP server and tool integration.
-    
+
     Manages discovery, registration, and execution of MCP tools,
     bridging external services with the LLM core service.
+
+    Uses proper MCP JSON-RPC protocol for server communication.
     """
-    
+
     def __init__(self):
         """Initialize MCP integration."""
-        self.mcp_servers = {}
-        self.tool_registry = {}
-        self.active_connections = {}
-        
+        self.mcp_servers: Dict[str, Dict[str, Any]] = {}
+        self.tool_registry: Dict[str, Dict[str, Any]] = {}
+        self.server_connections: Dict[str, MCPServerConnection] = {}
+
         # MCP configuration directories (local project first, then global)
         self.local_mcp_dir = Path.cwd() / ".kollabor-cli" / "mcp"
         self.global_mcp_dir = Path.home() / ".kollabor-cli" / "mcp"
-        
+
         # Load from both local and global configs
         self._load_mcp_config()
-        
+
         logger.info("MCP Integration initialized")
     
     def _load_mcp_config(self):
@@ -50,7 +231,7 @@ class MCPIntegration:
     
     def _load_config_from_dir(self, config_dir: Path, config_type: str):
         """Load MCP config from a specific directory.
-        
+
         Args:
             config_dir: Directory to load config from
             config_type: Type of config (local/global) for logging
@@ -67,28 +248,97 @@ class MCPIntegration:
             logger.warning(f"Failed to load {config_type} MCP config: {e}")
     
     async def discover_mcp_servers(self) -> Dict[str, Any]:
-        """Auto-discover available MCP servers.
-        
+        """Auto-discover available MCP servers and their tools.
+
+        Connects to each configured stdio server using MCP protocol,
+        initializes it, and queries available tools.
+
         Returns:
             Dictionary of discovered MCP servers and their capabilities
         """
         discovered = {}
-        
-        # Check for local MCP servers
+
+        # Check for local MCP servers (manifest-based)
         await self._discover_local_servers(discovered)
-        
-        # Check for configured servers
+
+        # Connect to configured stdio servers using MCP protocol
         for server_name, server_config in self.mcp_servers.items():
-            if await self._validate_server(server_config):
-                discovered[server_name] = {
-                    "name": server_name,
-                    "type": server_config.get("type", "unknown"),
-                    "capabilities": await self._get_server_capabilities(server_config),
-                    "status": "available"
-                }
-                logger.info(f"Discovered MCP server: {server_name}")
-        
+            if not server_config.get("enabled", True):
+                logger.debug(f"Skipping disabled MCP server: {server_name}")
+                continue
+
+            if server_config.get("type") == "stdio":
+                command = server_config.get("command")
+                if command:
+                    tools = await self._connect_and_list_tools(server_name, command)
+                    discovered[server_name] = {
+                        "name": server_name,
+                        "type": "stdio",
+                        "tools": [t.get("name") for t in tools],
+                        "tool_count": len(tools),
+                        "status": "connected" if tools else "no_tools"
+                    }
+
         return discovered
+
+    async def _connect_and_list_tools(self, server_name: str, command: str) -> List[Dict[str, Any]]:
+        """Connect to an MCP server and list its tools.
+
+        Uses proper MCP JSON-RPC protocol:
+        1. Start server process
+        2. Send initialize request
+        3. Send tools/list request
+        4. Register discovered tools
+
+        Args:
+            server_name: Name of the server
+            command: Command to start the server
+
+        Returns:
+            List of tool definitions from the server
+        """
+        # Close existing connection if any
+        if server_name in self.server_connections:
+            await self.server_connections[server_name].close()
+
+        # Create new connection
+        connection = MCPServerConnection(server_name, command)
+
+        if not await connection.connect():
+            logger.warning(f"Failed to connect to MCP server: {server_name}")
+            return []
+
+        if not await connection.initialize():
+            logger.warning(f"Failed to initialize MCP server: {server_name}")
+            await connection.close()
+            return []
+
+        # List tools
+        tools = await connection.list_tools()
+
+        # Register tools
+        for tool in tools:
+            tool_name = tool.get("name")
+            if tool_name:
+                self.tool_registry[tool_name] = {
+                    "server": server_name,
+                    "definition": {
+                        "name": tool_name,
+                        "description": tool.get("description", ""),
+                        "parameters": tool.get("inputSchema", {
+                            "type": "object",
+                            "properties": {},
+                            "required": []
+                        })
+                    },
+                    "enabled": True
+                }
+                logger.info(f"Registered MCP tool: {tool_name} from {server_name}")
+
+        # Keep connection open for tool calls
+        self.server_connections[server_name] = connection
+
+        return tools
     
     async def _discover_local_servers(self, discovered: Dict):
         """Discover locally running MCP servers."""
@@ -184,12 +434,12 @@ class MCPIntegration:
             return False
     
     async def call_mcp_tool(self, tool_name: str, params: Dict[str, Any]) -> Dict[str, Any]:
-        """Execute an MCP tool call.
-        
+        """Execute an MCP tool call using proper MCP protocol.
+
         Args:
             tool_name: Name of the tool to execute
             params: Parameters for the tool
-            
+
         Returns:
             Tool execution result
         """
@@ -198,28 +448,30 @@ class MCPIntegration:
                 "error": f"Tool '{tool_name}' not found",
                 "available_tools": list(self.tool_registry.keys())
             }
-        
+
         tool_info = self.tool_registry[tool_name]
         server_name = tool_info["server"]
-        
+
         if not tool_info["enabled"]:
             return {"error": f"Tool '{tool_name}' is disabled"}
-        
-        try:
-            # Get server configuration
+
+        # Get active connection
+        connection = self.server_connections.get(server_name)
+        if not connection or not connection.initialized:
+            # Try to reconnect
             server_config = self.mcp_servers.get(server_name, {})
-            
-            # Execute tool based on server type
-            if server_config.get("type") == "stdio":
-                result = await self._execute_stdio_tool(server_config, tool_name, params)
-            elif server_config.get("type") == "http":
-                result = await self._execute_http_tool(server_config, tool_name, params)
-            else:
-                result = {"error": f"Unknown server type for {server_name}"}
-            
+            command = server_config.get("command")
+            if command:
+                await self._connect_and_list_tools(server_name, command)
+                connection = self.server_connections.get(server_name)
+
+            if not connection or not connection.initialized:
+                return {"error": f"No active connection to server '{server_name}'"}
+
+        try:
+            result = await connection.call_tool(tool_name, params)
             logger.info(f"Executed MCP tool: {tool_name}")
             return result
-            
         except Exception as e:
             logger.error(f"Failed to execute MCP tool {tool_name}: {e}")
             return {"error": str(e)}
@@ -257,8 +509,9 @@ class MCPIntegration:
         # Add parameters as JSON input
         input_json = json.dumps(params)
 
-        try:
-            result = subprocess.run(
+        def run_subprocess():
+            """Run subprocess in thread to avoid blocking event loop."""
+            return subprocess.run(
                 command_parts,
                 shell=False,
                 input=input_json,
@@ -266,7 +519,12 @@ class MCPIntegration:
                 text=True,
                 timeout=30
             )
-            
+
+        try:
+            # Run blocking subprocess in executor to not freeze event loop
+            loop = asyncio.get_event_loop()
+            result = await loop.run_in_executor(None, run_subprocess)
+
             if result.returncode == 0:
                 # Try to parse JSON output
                 try:
@@ -275,7 +533,7 @@ class MCPIntegration:
                     return {"output": result.stdout}
             else:
                 return {"error": result.stderr or f"Tool exited with code {result.returncode}"}
-                
+
         except subprocess.TimeoutExpired:
             return {"error": "Tool execution timed out"}
         except Exception as e:
@@ -327,7 +585,7 @@ class MCPIntegration:
     
     def list_available_tools(self) -> List[Dict[str, Any]]:
         """List all available MCP tools.
-        
+
         Returns:
             List of available tools with their information
         """
@@ -340,6 +598,233 @@ class MCPIntegration:
                 "definition": tool_info.get("definition", {})
             })
         return tools
+
+    def get_tool_definitions_for_api(self) -> List[Dict[str, Any]]:
+        """Convert registered MCP tools to API tool schema format.
+
+        Returns generic format that adapters (OpenAI/Anthropic) auto-convert:
+        - OpenAI wraps in: {type: "function", function: {...}}
+        - Anthropic uses: {name, description, input_schema}
+
+        Returns:
+            List of tool definitions in generic API format
+        """
+        tools = []
+
+        # Add MCP tools from registry
+        for tool_name, tool_info in self.tool_registry.items():
+            if not tool_info.get("enabled", True):
+                continue
+
+            definition = tool_info.get("definition", {})
+            tools.append({
+                "name": tool_name,
+                "description": definition.get("description", f"MCP tool: {tool_name}"),
+                "parameters": definition.get("parameters", definition.get("inputSchema", {
+                    "type": "object",
+                    "properties": {},
+                    "required": []
+                }))
+            })
+
+        # Add built-in file operation tools
+        tools.extend(self._get_file_operation_tools())
+
+        logger.debug(f"Prepared {len(tools)} tools for API ({len(self.tool_registry)} MCP + file ops)")
+        return tools
+
+    def _get_file_operation_tools(self) -> List[Dict[str, Any]]:
+        """Get built-in file operation tool definitions.
+
+        These tools provide file manipulation capabilities that work alongside
+        MCP tools. They match the XML-based file operations but as native functions.
+
+        Returns:
+            List of file operation tool definitions
+        """
+        return [
+            {
+                "name": "file_read",
+                "description": "Read content from a file. Use this to examine existing files before editing.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "file": {"type": "string", "description": "Relative file path to read"},
+                        "offset": {"type": "integer", "description": "Line offset to start reading from (0-indexed, optional)"},
+                        "limit": {"type": "integer", "description": "Number of lines to read (optional)"}
+                    },
+                    "required": ["file"]
+                }
+            },
+            {
+                "name": "file_create",
+                "description": "Create a new file with content. Fails if file already exists.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "file": {"type": "string", "description": "Relative file path to create"},
+                        "content": {"type": "string", "description": "Content to write to the file"}
+                    },
+                    "required": ["file", "content"]
+                }
+            },
+            {
+                "name": "file_create_overwrite",
+                "description": "Create or overwrite a file with content. Creates backup if file exists.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "file": {"type": "string", "description": "Relative file path to create/overwrite"},
+                        "content": {"type": "string", "description": "Content to write to the file"}
+                    },
+                    "required": ["file", "content"]
+                }
+            },
+            {
+                "name": "file_edit",
+                "description": "Find and replace text in a file. Replaces ALL occurrences of the pattern.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "file": {"type": "string", "description": "Relative file path to edit"},
+                        "find": {"type": "string", "description": "Text pattern to find (exact match)"},
+                        "replace": {"type": "string", "description": "Text to replace with"}
+                    },
+                    "required": ["file", "find", "replace"]
+                }
+            },
+            {
+                "name": "file_append",
+                "description": "Append content to the end of a file.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "file": {"type": "string", "description": "Relative file path to append to"},
+                        "content": {"type": "string", "description": "Content to append"}
+                    },
+                    "required": ["file", "content"]
+                }
+            },
+            {
+                "name": "file_insert_after",
+                "description": "Insert content after a pattern in a file. Pattern must match exactly.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "file": {"type": "string", "description": "Relative file path"},
+                        "pattern": {"type": "string", "description": "Pattern to find (exact match)"},
+                        "content": {"type": "string", "description": "Content to insert after pattern"}
+                    },
+                    "required": ["file", "pattern", "content"]
+                }
+            },
+            {
+                "name": "file_insert_before",
+                "description": "Insert content before a pattern in a file. Pattern must match exactly.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "file": {"type": "string", "description": "Relative file path"},
+                        "pattern": {"type": "string", "description": "Pattern to find (exact match)"},
+                        "content": {"type": "string", "description": "Content to insert before pattern"}
+                    },
+                    "required": ["file", "pattern", "content"]
+                }
+            },
+            {
+                "name": "file_delete",
+                "description": "Delete a file. Creates backup before deletion.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "file": {"type": "string", "description": "Relative file path to delete"}
+                    },
+                    "required": ["file"]
+                }
+            },
+            {
+                "name": "file_move",
+                "description": "Move or rename a file.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "from": {"type": "string", "description": "Source file path"},
+                        "to": {"type": "string", "description": "Destination file path"}
+                    },
+                    "required": ["from", "to"]
+                }
+            },
+            {
+                "name": "file_copy",
+                "description": "Copy a file. Fails if destination exists.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "from": {"type": "string", "description": "Source file path"},
+                        "to": {"type": "string", "description": "Destination file path"}
+                    },
+                    "required": ["from", "to"]
+                }
+            },
+            {
+                "name": "file_copy_overwrite",
+                "description": "Copy a file, overwriting destination if it exists.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "from": {"type": "string", "description": "Source file path"},
+                        "to": {"type": "string", "description": "Destination file path"}
+                    },
+                    "required": ["from", "to"]
+                }
+            },
+            {
+                "name": "file_mkdir",
+                "description": "Create a directory (including parent directories).",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "path": {"type": "string", "description": "Directory path to create"}
+                    },
+                    "required": ["path"]
+                }
+            },
+            {
+                "name": "file_rmdir",
+                "description": "Remove an empty directory.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "path": {"type": "string", "description": "Directory path to remove"}
+                    },
+                    "required": ["path"]
+                }
+            },
+            {
+                "name": "file_grep",
+                "description": "Search for a pattern in a file and return matching lines.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "file": {"type": "string", "description": "Relative file path to search"},
+                        "pattern": {"type": "string", "description": "Text pattern to search for"},
+                        "case_insensitive": {"type": "boolean", "description": "Whether to ignore case (default: false)"}
+                    },
+                    "required": ["file", "pattern"]
+                }
+            },
+            {
+                "name": "terminal",
+                "description": "Execute a terminal/shell command and return the output.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "command": {"type": "string", "description": "Shell command to execute"}
+                    },
+                    "required": ["command"]
+                }
+            }
+        ]
     
     def enable_tool(self, tool_name: str) -> bool:
         """Enable an MCP tool.
@@ -372,15 +857,14 @@ class MCPIntegration:
         return False
     
     async def shutdown(self):
-        """Shutdown MCP integration and close connections."""
-        # Close any active connections
-        for connection_id, connection in self.active_connections.items():
+        """Shutdown MCP integration and close all server connections."""
+        for server_name, connection in self.server_connections.items():
             try:
-                if hasattr(connection, 'close'):
-                    await connection.close()
-                logger.debug(f"Closed MCP connection: {connection_id}")
+                await connection.close()
+                logger.debug(f"Closed MCP server connection: {server_name}")
             except Exception as e:
-                logger.warning(f"Error closing MCP connection {connection_id}: {e}")
-        
-        self.active_connections.clear()
+                logger.warning(f"Error closing MCP connection {server_name}: {e}")
+
+        self.server_connections.clear()
+        self.tool_registry.clear()
         logger.info("MCP Integration shutdown complete")