server.py

#!/usr/bin/env python3
"""
FORUS Cortex Agent - Working MCP Server
Using proper MCP JSON-RPC protocol with streamable HTTP transport
"""

import os
import sys
import json
import asyncio
import logging
import pickle
from typing import Dict, List, Any, Optional
from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.parse import urlparse
import subprocess
import tempfile

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

class MCPJSONRPCHandler(BaseHTTPRequestHandler):
    """HTTP handler implementing MCP JSON-RPC protocol"""
    
    def do_POST(self):
        """Handle MCP JSON-RPC requests"""
        try:
            # Parse request
            content_length = int(self.headers.get('Content-Length', 0))
            if content_length == 0:
                self._send_error(-32600, "Invalid Request: Empty body")
                return
            
            body = self.rfile.read(content_length)
            try:
                request = json.loads(body.decode('utf-8'))
            except json.JSONDecodeError as e:
                self._send_error(-32700, f"Parse error: {str(e)}")
                return
            
            # Validate JSON-RPC format
            if not isinstance(request, dict):
                self._send_error(-32600, "Invalid Request: Not a JSON object")
                return
            
            if request.get('jsonrpc') != '2.0':
                self._send_error(-32600, "Invalid Request: Missing or invalid jsonrpc version")
                return
            
            # Handle the request
            response = self._handle_jsonrpc_request(request)
            
            # Send response
            self.send_response(200)
            self.send_header('Content-Type', 'application/json')
            self.send_header('Access-Control-Allow-Origin', '*')
            self.send_header('Access-Control-Allow-Methods', 'POST, OPTIONS')
            self.send_header('Access-Control-Allow-Headers', 'Content-Type, Authorization')
            self.end_headers()
            
            self.wfile.write(json.dumps(response).encode('utf-8'))
            
        except Exception as e:
            logger.error(f"Request handling error: {e}", exc_info=True)
            self._send_error(-32603, f"Internal error: {str(e)}")
    
    def do_OPTIONS(self):
        """Handle CORS preflight requests"""
        self.send_response(200)
        self.send_header('Access-Control-Allow-Origin', '*')
        self.send_header('Access-Control-Allow-Methods', 'POST, OPTIONS')
        self.send_header('Access-Control-Allow-Headers', 'Content-Type, Authorization')
        self.end_headers()
    
    def do_GET(self):
        """Handle GET requests for server info"""
        if self.path == '/' or self.path == '/mcp':
            response = {
                "server": "FORUS Cortex Agent MCP Server",
                "version": "2.0.0",
                "protocol": "MCP (Model Context Protocol)",
                "transport": "HTTP JSON-RPC",
                "endpoints": {
                    "/mcp": "Main MCP JSON-RPC endpoint (POST)",
                    "/": "Server information (GET)"
                },
                "tools_count": 6,
                "documentation": "https://modelcontextprotocol.io/"
            }
            
            self.send_response(200)
            self.send_header('Content-Type', 'application/json')
            self.send_header('Access-Control-Allow-Origin', '*')
            self.end_headers()
            self.wfile.write(json.dumps(response, indent=2).encode('utf-8'))
        else:
            self.send_response(404)
            self.send_header('Content-Type', 'application/json')
            self.end_headers()
            self.wfile.write(json.dumps({"error": "Not found"}).encode('utf-8'))
    
    def _send_error(self, code: int, message: str, request_id=None):
        """Send JSON-RPC error response"""
        error_response = {
            "jsonrpc": "2.0",
            "id": request_id,
            "error": {
                "code": code,
                "message": message
            }
        }
        
        self.send_response(400 if code == -32600 else 500)
        self.send_header('Content-Type', 'application/json')
        self.send_header('Access-Control-Allow-Origin', '*')
        self.end_headers()
        self.wfile.write(json.dumps(error_response).encode('utf-8'))
    
    def _handle_jsonrpc_request(self, request: Dict) -> Dict:
        """Handle JSON-RPC request and return response"""
        request_id = request.get('id')
        method = request.get('method', '')
        params = request.get('params', {})
        
        try:
            if method == 'initialize':
                return self._handle_initialize(request_id, params)
            elif method == 'tools/list':
                return self._handle_tools_list(request_id)
            elif method == 'tools/call':
                return self._handle_tools_call(request_id, params)
            elif method == 'resources/list':
                return self._handle_resources_list(request_id)
            elif method == 'prompts/list':
                return self._handle_prompts_list(request_id)
            else:
                return {
                    "jsonrpc": "2.0",
                    "id": request_id,
                    "error": {
                        "code": -32601,
                        "message": f"Method not found: {method}"
                    }
                }
                
        except Exception as e:
            logger.error(f"Method {method} error: {e}", exc_info=True)
            return {
                "jsonrpc": "2.0",
                "id": request_id,
                "error": {
                    "code": -32603,
                    "message": f"Internal error: {str(e)}"
                }
            }
    
    def _handle_initialize(self, request_id, params):
        """Handle MCP initialize request"""
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "result": {
                "protocolVersion": "2024-11-05",
                "capabilities": {
                    "tools": {},
                    "resources": {},
                    "prompts": {}
                },
                "serverInfo": {
                    "name": "forus-cortex-agent",
                    "version": "2.0.0"
                }
            }
        }
    
    def _handle_tools_list(self, request_id):
        """List available MCP tools"""
        tools = [
            {
                "name": "memory_search",
                "description": "Search the agent's memory using semantic similarity",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "query": {"type": "string", "description": "Search query"},
                        "limit": {"type": "integer", "default": 10, "description": "Max results"},
                        "area": {"type": "string", "enum": ["main", "solutions", "errors"], "description": "Memory area filter"}
                    },
                    "required": ["query"]
                }
            },
            {
                "name": "query_msa",
                "description": "Query Master Services Agreement documents",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "query": {"type": "string", "description": "Search query for MSA documents"}
                    },
                    "required": ["query"]
                }
            },
            {
                "name": "execute_python",
                "description": "Execute Python code in a sandboxed environment",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "code": {"type": "string", "description": "Python code to execute"},
                        "timeout": {"type": "integer", "default": 30, "description": "Timeout in seconds"}
                    },
                    "required": ["code"]
                }
            },
            {
                "name": "execute_shell",
                "description": "Execute shell commands safely",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "command": {"type": "string", "description": "Shell command to execute"},
                        "timeout": {"type": "integer", "default": 30, "description": "Timeout in seconds"}
                    },
                    "required": ["command"]
                }
            },
            {
                "name": "get_system_status",
                "description": "Get current status of the agent system",
                "inputSchema": {
                    "type": "object",
                    "properties": {}
                }
            },
            {
                "name": "list_mcp_tools",
                "description": "List all available MCP tools with descriptions",
                "inputSchema": {
                    "type": "object",
                    "properties": {}
                }
            }
        ]
        
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "result": {
                "tools": tools
            }
        }
    
    def _handle_tools_call(self, request_id, params):
        """Execute a tool"""
        tool_name = params.get('name', '')
        arguments = params.get('arguments', {})
        
        # Execute the tool
        if tool_name == 'memory_search':
            result = self._tool_memory_search(arguments)
        elif tool_name == 'query_msa':
            result = self._tool_query_msa(arguments)
        elif tool_name == 'execute_python':
            result = self._tool_execute_python(arguments)
        elif tool_name == 'execute_shell':
            result = self._tool_execute_shell(arguments)
        elif tool_name == 'get_system_status':
            result = self._tool_get_system_status(arguments)
        elif tool_name == 'list_mcp_tools':
            result = self._tool_list_mcp_tools(arguments)
        else:
            return {
                "jsonrpc": "2.0",
                "id": request_id,
                "error": {
                    "code": -32601,
                    "message": f"Tool not found: {tool_name}"
                }
            }
        
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "result": {
                "content": [
                    {
                        "type": "text",
                        "text": json.dumps(result, indent=2)
                    }
                ]
            }
        }
    
    def _handle_resources_list(self, request_id):
        """List available resources"""
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "result": {
                "resources": []  # No resources implemented yet
            }
        }
    
    def _handle_prompts_list(self, request_id):
        """List available prompts"""
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "result": {
                "prompts": []  # No prompts implemented yet
            }
        }
    
    # Tool implementations
    def _tool_memory_search(self, args):
        """Search agent memory"""
        try:
            query = args.get('query', '')
            limit = args.get('limit', 10)
            area = args.get('area')
            
            memory_path = "/root/forus_cortex_agent/memory/default/index.pkl"
            if not os.path.exists(memory_path):
                return {"success": False, "error": "Memory index not found"}
            
            with open(memory_path, "rb") as f:
                docstore, index_to_id = pickle.load(f)
            
            results = []
            query_lower = query.lower()
            
            for doc_id, doc in list(docstore._dict.items())[:limit * 2]:
                if hasattr(doc, 'page_content'):
                    content = doc.page_content
                    metadata = getattr(doc, 'metadata', {})
                    
                    if area and metadata.get('area') != area:
                        continue
                    
                    if query_lower in content.lower():
                        score = 1.0 if query_lower in content.lower()[:100] else 0.5
                        results.append({
                            "id": doc_id,
                            "text": content[:500] + ("..." if len(content) > 500 else ""),
                            "score": score,
                            "area": metadata.get("area", "main"),
                            "metadata": metadata
                        })
                        
                        if len(results) >= limit:
                            break
            
            return {
                "success": True,
                "query": query,
                "count": len(results),
                "results": results
            }
            
        except Exception as e:
            logger.error(f"Memory search error: {e}")
            return {"success": False, "error": f"Memory search failed: {str(e)}"}
    
    def _tool_query_msa(self, args):
        """Query MSA documents"""
        try:
            query = args.get('query', '')
            
            memory_path = "/root/forus_cortex_agent/memory/default/index.pkl"
            if not os.path.exists(memory_path):
                return {"success": False, "error": "MSA documents not found"}
            
            with open(memory_path, "rb") as f:
                docstore, index_to_id = pickle.load(f)
            
            results = []
            query_lower = query.lower()
            
            for doc_id, doc in docstore._dict.items():
                if hasattr(doc, 'page_content') and hasattr(doc, 'metadata'):
                    content = doc.page_content.lower()
                    source = doc.metadata.get('source', '')
                    
                    if 'master_services_agreement' in source.lower() and query_lower in content:
                        relevance = "high" if query_lower in content[:200] else "medium"
                        results.append({
                            "content": doc.page_content[:1000] + ("..." if len(doc.page_content) > 1000 else ""),
                            "source": source.split('/')[-1] if source else "Unknown",
                            "relevance": relevance,
                            "doc_id": doc_id
                        })
                        
                        if len(results) >= 5:
                            break
            
            return {
                "success": True,
                "query": query,
                "count": len(results),
                "results": results
            }
            
        except Exception as e:
            logger.error(f"MSA query error: {e}")
            return {"success": False, "error": f"MSA query failed: {str(e)}"}
    
    def _tool_execute_python(self, args):
        """Execute Python code"""
        try:
            code = args.get('code', '')
            timeout = min(args.get('timeout', 30), 60)
            
            with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
                f.write(code)
                temp_file = f.name
            
            try:
                result = subprocess.run(
                    [sys.executable, temp_file],
                    capture_output=True,
                    text=True,
                    timeout=timeout,
                    cwd=tempfile.gettempdir()
                )
                
                return {
                    "success": True,
                    "stdout": result.stdout,
                    "stderr": result.stderr,
                    "return_code": result.returncode,
                    "execution_time": f"< {timeout}s"
                }
                
            finally:
                try:
                    os.unlink(temp_file)
                except OSError:
                    pass
                    
        except subprocess.TimeoutExpired:
            return {
                "success": False,
                "error": f"Code execution timed out after {timeout} seconds"
            }
        except Exception as e:
            logger.error(f"Python execution error: {e}")
            return {"success": False, "error": f"Execution failed: {str(e)}"}
    
    def _tool_execute_shell(self, args):
        """Execute shell command"""
        try:
            command = args.get('command', '')
            timeout = min(args.get('timeout', 30), 30)
            
            ALLOWED_COMMANDS = {
                'ls', 'pwd', 'echo', 'cat', 'head', 'tail', 'grep', 'find', 'wc', 'sort',
                'date', 'whoami', 'id', 'hostname', 'uname', 'df', 'free', 'ps'
            }
            
            cmd_parts = command.strip().split()
            if not cmd_parts:
                return {"success": False, "error": "Empty command"}
            
            base_command = cmd_parts[0]
            if base_command not in ALLOWED_COMMANDS:
                return {
                    "success": False,
                    "error": f"Command '{base_command}' not allowed. Allowed: {', '.join(sorted(ALLOWED_COMMANDS))}"
                }
            
            result = subprocess.run(
                command,
                shell=True,
                capture_output=True,
                text=True,
                timeout=timeout
            )
            
            return {
                "success": True,
                "stdout": result.stdout,
                "stderr": result.stderr,
                "return_code": result.returncode,
                "command": command
            }
            
        except subprocess.TimeoutExpired:
            return {"success": False, "error": f"Command timed out after {timeout} seconds"}
        except Exception as e:
            logger.error(f"Shell execution error: {e}")
            return {"success": False, "error": f"Command failed: {str(e)}"}
    
    def _tool_get_system_status(self, args):
        """Get system status"""
        try:
            import platform
            
            memory_available = os.path.exists("/root/forus_cortex_agent/memory/default/index.pkl")
            
            system_info = {
                "platform": platform.system(),
                "python_version": platform.python_version(),
                "hostname": platform.node()
            }
            
            try:
                import psutil
                cpu_percent = psutil.cpu_percent(interval=1)
                memory = psutil.virtual_memory()
                disk = psutil.disk_usage('/')
                
                resources = {
                    "cpu_percent": cpu_percent,
                    "memory_percent": memory.percent,
                    "memory_available_gb": round(memory.available / (1024**3), 2),
                    "disk_free_gb": round(disk.free / (1024**3), 2),
                    "disk_percent": round((disk.used / disk.total) * 100, 2)
                }
            except ImportError:
                resources = {"error": "psutil not available"}
            
            return {
                "success": True,
                "server": "FORUS Cortex Agent MCP Server",
                "version": "2.0.0",
                "protocol": "MCP (Model Context Protocol)",
                "transport": "HTTP JSON-RPC",
                "memory_available": memory_available,
                "system": system_info,
                "resources": resources,
                "tools_count": 6
            }
            
        except Exception as e:
            logger.error(f"System status error: {e}")
            return {"success": False, "error": f"Failed to get system status: {str(e)}"}
    
    def _tool_list_mcp_tools(self, args):
        """List MCP tools"""
        tools = [
            {"name": "memory_search", "category": "memory", "description": "Search agent memory"},
            {"name": "query_msa", "category": "knowledge", "description": "Query MSA documents"},
            {"name": "execute_python", "category": "execution", "description": "Execute Python code"},
            {"name": "execute_shell", "category": "execution", "description": "Execute shell commands"},
            {"name": "get_system_status", "category": "system", "description": "Get system status"},
            {"name": "list_mcp_tools", "category": "system", "description": "List all tools"}
        ]
        
        return {
            "success": True,
            "server": "FORUS Cortex Agent MCP Server",
            "total_tools": len(tools),
            "tools": tools,
            "categories": {
                "memory": 1,
                "knowledge": 1,
                "execution": 2,
                "system": 2
            }
        }
    
    def log_message(self, format, *args):
        """Override to use proper logging"""
        logger.info(f"HTTP: {format % args}")

def main():
    """Main entry point"""
    try:
        logger.info("=" * 60)
        logger.info("FORUS Cortex Agent - Working MCP Server")
        logger.info("=" * 60)
        
        # Check memory availability
        memory_path = "/root/forus_cortex_agent/memory/default/index.pkl"
        if os.path.exists(memory_path):
            logger.info(f"✓ Memory index found: {memory_path}")
        else:
            logger.warning(f"⚠ Memory index not found: {memory_path}")
        
        # Start HTTP server
        host = '0.0.0.0'
        port = 8083
        
        server = HTTPServer((host, port), MCPJSONRPCHandler)
        
        logger.info(f"MCP Server running on http://{host}:{port}")
        logger.info(f"MCP Endpoint: http://93.127.162.206:{port}/mcp")
        logger.info(f"Server Info: http://93.127.162.206:{port}/")
        logger.info("Protocol: MCP JSON-RPC 2.0")
        logger.info("Transport: HTTP")
        logger.info("=" * 60)
        
        server.serve_forever()
        
    except KeyboardInterrupt:
        logger.info("Server shutdown requested")
    except Exception as e:
        logger.error(f"Server error: {e}", exc_info=True)
        raise

if __name__ == "__main__":
    main()