agent/local_mcp.py

"""
Local MCP Client for Overgrowth - connects to local MCP server
"""

import subprocess
import json
import sys
import logging
import uuid
from typing import Dict, List, Any, Optional
from pathlib import Path

# Configure logging
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')

# Import API monitor for tracking
try:
    from agent.api_monitor import monitor
except ImportError:
    logger.warning("API monitor not available - tracking disabled")
    monitor = None

# Path to local MCP server
# In production (HuggingFace Space), use relative path
# In dev, use absolute path
import os
if os.path.exists("/home/gpaasch/overgrowth-mcp-server/server.py"):
    MCP_SERVER_PATH = Path("/home/gpaasch/overgrowth-mcp-server/server.py")
    VENV_PYTHON = Path("/home/gpaasch/overgrowth-mcp-server/venv/bin/python")
elif os.path.exists("mcp-server/server.py"):
    # Running from overgrowth directory
    MCP_SERVER_PATH = Path("mcp-server/server.py")
    # Try to use the overgrowth-mcp-server venv if it exists
    if os.path.exists("/home/gpaasch/overgrowth-mcp-server/venv/bin/python"):
        VENV_PYTHON = Path("/home/gpaasch/overgrowth-mcp-server/venv/bin/python")
    else:
        VENV_PYTHON = Path("/usr/local/bin/python3")  # Fallback
else:
    # HuggingFace Space - use bundled MCP server
    MCP_SERVER_PATH = Path(__file__).parent.parent / "mcp-server" / "server.py"
    VENV_PYTHON = Path("/usr/local/bin/python3")  # Use system python in Space


class LocalMCPClient:
    """Client for local Overgrowth MCP server via stdio"""
    
    def __init__(self):
        self.server_path = MCP_SERVER_PATH
        self.python_path = VENV_PYTHON if VENV_PYTHON.exists() else sys.executable
        self._request_id = 0
    
    def _get_request_id(self) -> int:
        """Get next request ID"""
        self._request_id += 1
        return self._request_id
    
    def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """
        Call a tool on the MCP server
        
        Args:
            tool_name: Name of the tool to call
            arguments: Tool arguments
            
        Returns:
            Tool response data
        """
        call_id = str(uuid.uuid4())
        
        if monitor:
            monitor.start_call(call_id, "gns3", "local-mcp", tool_name, arguments=arguments)
        
        try:
            # Initialize request
            init_request = {
                "jsonrpc": "2.0",
                "id": self._get_request_id(),
                "method": "initialize",
                "params": {
                    "protocolVersion": "2024-11-05",
                    "capabilities": {},
                    "clientInfo": {
                        "name": "overgrowth-gradio-client",
                        "version": "0.1.0"
                    }
                }
            }
            
            # Tool call request
            tool_request = {
                "jsonrpc": "2.0",
                "id": self._get_request_id(),
                "method": "tools/call",
                "params": {
                    "name": tool_name,
                    "arguments": arguments
                }
            }
            
            # Prepare input (init + tool call)
            input_data = json.dumps(init_request) + "\n" + json.dumps(tool_request) + "\n"
            
            logger.info(f"Calling MCP tool: {tool_name} with server at {self.server_path}")
            logger.debug(f"Using python: {self.python_path}")
            
            # Run server process
            result = subprocess.run(
                [str(self.python_path), str(self.server_path)],
                input=input_data,
                capture_output=True,
                text=True,
                timeout=60  # Increased timeout for network operations
            )
            
            logger.debug(f"MCP server exit code: {result.returncode}")
            if result.stderr:
                logger.debug(f"MCP server stderr: {result.stderr[:500]}")
            
            # Parse responses
            lines = result.stdout.strip().split('\n')
            
            # Look for JSON-RPC responses
            for line in lines:
                if not line.strip() or line.startswith('INFO:') or line.startswith('WARNING:'):
                    continue
                    
                try:
                    response = json.loads(line)
                    
                    # Check if it's a tool response
                    if response.get('id') == tool_request['id']:
                        if 'error' in response:
                            error_msg = response['error']
                            logger.error(f"MCP tool {tool_name} failed: {error_msg}")
                            if monitor:
                                monitor.complete_call(call_id, success=False, error_message=str(error_msg))
                            raise Exception(f"MCP Error: {error_msg}")
                        
                        # Extract text content from result
                        result_data = response.get('result', {})
                        content = result_data.get('content', [])
                        
                        if content and len(content) > 0:
                            logger.info(f"MCP tool {tool_name} completed successfully")
                            if monitor:
                                monitor.complete_call(call_id, success=True)
                            return {
                                'success': True,
                                'text': content[0].get('text', ''),
                                'raw': result_data
                            }
                        
                        logger.info(f"MCP tool {tool_name} completed (no text content)")
                        if monitor:
                            monitor.complete_call(call_id, success=True)
                        return {
                            'success': True,
                            'raw': result_data
                        }
                        
                except json.JSONDecodeError as e:
                    logger.debug(f"Skipping non-JSON line: {line[:100]}")
                    continue
            
            # If we get here, no valid response found
            logger.error(f"No valid response from MCP server for tool {tool_name}")
            logger.error(f"stdout: {result.stdout[:1000]}")
            if monitor:
                monitor.complete_call(call_id, success=False, error_message="No valid response from MCP server")
            raise Exception(f"No valid response from MCP server. stdout: {result.stdout[:500]}, stderr: {result.stderr[:500]}")
            
        except subprocess.TimeoutExpired:
            logger.error(f"MCP server call to {tool_name} timed out after 60 seconds")
            if monitor:
                monitor.complete_call(call_id, success=False, error_message="Timeout after 60s")
            raise Exception(f"MCP server call to {tool_name} timed out - this may indicate GNS3 server is unreachable or slow")
        except Exception as e:
            logger.error(f"Failed to call MCP tool {tool_name}: {str(e)}")
            if monitor:
                monitor.complete_call(call_id, success=False, error_message=str(e))
            raise Exception(f"Failed to call MCP tool {tool_name}: {str(e)}")
    
    def get_topology(self, project_name: str = "overgrowth") -> Dict[str, Any]:
        """Get GNS3 topology information"""
        result = self.call_tool("get_topology", {"project_name": project_name})
        
        if result.get('success'):
            # Parse JSON from text response
            text = result.get('text', '{}')
            return json.loads(text)
        
        raise Exception("Failed to get topology")
    
    def get_projects(self) -> List[Dict[str, Any]]:
        """Get list of GNS3 projects"""
        result = self.call_tool("get_projects", {})
        
        if result.get('success'):
            text = result.get('text', '[]')
            return json.loads(text)
        
        return []
    
    def start_device(self, project_name: str, device_name: str) -> bool:
        """Start a network device"""
        result = self.call_tool("start_device", {
            "project_name": project_name,
            "device_name": device_name
        })
        
        return result.get('success', False)
    
    def stop_device(self, project_name: str, device_name: str) -> bool:
        """Stop a network device"""
        result = self.call_tool("stop_device", {
            "project_name": project_name,
            "device_name": device_name
        })
        
        return result.get('success', False)
    
    def get_device_config(self, host: str, username: str = "admin", port: int = 22) -> str:
        """Get device running configuration via SSH"""
        result = self.call_tool("get_device_config", {
            "host": host,
            "username": username,
            "port": port
        })
        
        if result.get('success'):
            return result.get('text', '')
        
        raise Exception(f"Failed to get config from {host}")
    
    def configure_device(self, host: str, commands: List[str], 
                        username: str = "admin", port: int = 22) -> str:
        """Send configuration commands to device"""
        result = self.call_tool("configure_device", {
            "host": host,
            "commands": commands,
            "username": username,
            "port": port
        })
        
        if result.get('success'):
            return result.get('text', '')
        
        raise Exception(f"Failed to configure {host}")
    
    def backup_config(self, host: str, device_name: str, username: str = "admin") -> str:
        """Backup device configuration"""
        # Prefer the updated tool name, fall back to the older alias for compatibility.
        try:
            result = self.call_tool("backup_device_config", {
                "host": host,
                "device_name": device_name,
                "username": username
            })
        except Exception:
            result = self.call_tool("backup_config", {
                "host": host,
                "device_name": device_name,
                "username": username
            })
        
        if result.get('success'):
            return result.get('text', '')
        
        raise Exception(f"Failed to backup {device_name}")


# Singleton instance
_client = None

def get_mcp_client() -> LocalMCPClient:
    """Get or create MCP client singleton"""
    global _client
    if _client is None:
        _client = LocalMCPClient()
    return _client