mcp-forge/src/mcp_forge/execution/jupyter/kernel.py

"""
Real Jupyter kernel management for stateful execution.

This module implements proper Jupyter kernel management:
- jupyter-client runs on host (MCP-Forge server)
- ipykernel runs inside Podman containers
- Communication via ZMQ protocol
- 1:1 mapping: one container per session, one kernel per container
"""

from typing import Dict, Optional, List, Any
from dataclasses import dataclass
from datetime import datetime, timedelta
import uuid
import json
import tempfile
import time
from pathlib import Path

from jupyter_client.blocking.client import BlockingKernelClient
import zmq

from mcp_forge.podman.containers import SecureContainerManager, ContainerConfig
from mcp_forge.security.resource_limits import ResourceLimits
from mcp_forge.execution.simple.executor import ExecutionResult


class KernelError(Exception):
    """Raised when kernel operations fail."""
    pass


@dataclass
class KernelInfo:
    """Information about a running kernel."""
    kernel_id: str
    container_id: str
    session_id: str
    connection_file: Path
    connection_info: Dict[str, Any]  # ZMQ ports and keys
    started_at: datetime
    last_activity: datetime
    client: Optional[BlockingKernelClient] = None
    
    def to_dict(self) -> dict:
        """Convert to dictionary for JSON serialization."""
        return {
            "kernel_id": self.kernel_id,
            "container_id": self.container_id,
            "session_id": self.session_id,
            "started_at": self.started_at.isoformat(),
            "last_activity": self.last_activity.isoformat(),
            "connection_info": {
                "shell_port": self.connection_info.get("shell_port"),
                "iopub_port": self.connection_info.get("iopub_port"),
                "stdin_port": self.connection_info.get("stdin_port"),
                "control_port": self.connection_info.get("control_port"),
                "hb_port": self.connection_info.get("hb_port"),
            }
        }


class JupyterKernelManager:
    """
    Manages IPython kernels in containers via jupyter-client.
    
    Architecture:
    - This class runs on host (MCP-Forge server process)
    - Creates one container per session with ipykernel running inside
    - Connects to kernel via ZMQ protocol (jupyter-client)
    - Communicates using Jupyter message protocol
    
    Each session gets:
    - Dedicated container
    - Dedicated kernel process
    - Isolated Python namespace
    - Independent resource limits
    """
    
    def __init__(
        self,
        container_manager: SecureContainerManager,
        image: str,
        resource_limits: Optional[ResourceLimits] = None
    ):
        """
        Initialize kernel manager.
        
        Args:
            container_manager: Container lifecycle manager
            image: Docker/Podman image with ipykernel installed
            resource_limits: Default resource limits for kernels
        """
        self.container_manager = container_manager
        self.image = image
        self.resource_limits = resource_limits
        self.kernels: Dict[str, KernelInfo] = {}
    
    def start_kernel(
        self,
        session_id: str,
        volumes: Optional[Dict[str, dict]] = None
    ) -> str:
        """
        Start IPython kernel in dedicated container.
        
        Process:
        1. Generate ZMQ connection info (ports, keys)
        2. Create connection file
        3. Create container with ipykernel command
        4. Start container
        5. Wait for kernel to be ready
        6. Connect jupyter-client to kernel via ZMQ
        7. Verify kernel is responsive
        
        Args:
            session_id: Session ID this kernel belongs to
            volumes: Optional volume mounts
        
        Returns:
            kernel_id: Unique identifier for this kernel
        
        Raises:
            KernelError: If kernel startup fails
        """
        kernel_id = f"kernel-{uuid.uuid4().hex[:16]}"
        
        # Generate connection info
        connection_info = self._generate_connection_info()
        
        # Create connection file
        connection_file = self._create_connection_file(kernel_id, connection_info)
        
        try:
            # Create container with ipykernel
            config = ContainerConfig(
                image=self.image,
                command=[
                    "python", "-m", "ipykernel_launcher",
                    "-f", f"/tmp/kernel-{kernel_id}.json"
                ],
                resource_limits=self.resource_limits,
                volumes=volumes or {},
                # TODO: Port mappings for ZMQ
                # TODO: Mount connection file into container
            )
            
            container_id = self.container_manager.create_container(
                config,
                session_id=session_id,
                name=f"jupyter-{kernel_id}"
            )
            
            # Start container
            self.container_manager.start_container(container_id)
            
            # Wait for kernel to be ready
            time.sleep(2)  # TODO: Better readiness check
            
            # Connect client
            client = self._connect_client(connection_info)
            
            # Verify kernel is responsive
            if not self._verify_kernel(client):
                raise KernelError(f"Kernel {kernel_id} not responsive")
            
            # Register kernel
            now = datetime.utcnow()
            kernel_info = KernelInfo(
                kernel_id=kernel_id,
                container_id=container_id,
                session_id=session_id,
                connection_file=connection_file,
                connection_info=connection_info,
                started_at=now,
                last_activity=now,
                client=client
            )
            self.kernels[kernel_id] = kernel_info
            
            return kernel_id
            
        except Exception as e:
            # Cleanup on failure
            connection_file.unlink(missing_ok=True)
            raise KernelError(f"Failed to start kernel: {e}") from e
    
    def execute_code(
        self,
        kernel_id: str,
        code: str,
        timeout: int = 300
    ) -> ExecutionResult:
        """
        Execute code in kernel via ZMQ.
        
        Uses jupyter-client to:
        1. Send execute_request message
        2. Receive stream (stdout/stderr) messages
        3. Receive execute_result/display_data messages
        4. Collect and parse all output
        
        Args:
            kernel_id: Kernel to execute in
            code: Python code to execute
            timeout: Maximum execution time in seconds
        
        Returns:
            ExecutionResult with stdout, stderr, result
        
        Raises:
            KernelError: If kernel not found or execution fails
        """
        kernel_info = self._get_kernel(kernel_id)
        client = kernel_info.client
        
        if not client:
            raise KernelError(f"Kernel {kernel_id} has no connected client")
        
        start_time = time.time()
        
        try:
            # Execute code
            _msg_id = client.execute(code, silent=False, store_history=True)
            
            # Collect output
            stdout_parts = []
            stderr_parts = []
            result = None
            
            # Wait for execution to complete
            while True:
                try:
                    msg = client.get_iopub_msg(timeout=timeout)
                    msg_type = msg['header']['msg_type']
                    content = msg['content']
                    
                    if msg_type == 'stream':
                        if content['name'] == 'stdout':
                            stdout_parts.append(content['text'])
                        elif content['name'] == 'stderr':
                            stderr_parts.append(content['text'])
                    
                    elif msg_type == 'execute_result':
                        result = content.get('data', {}).get('text/plain', '')
                    
                    elif msg_type == 'error':
                        stderr_parts.append('\n'.join(content['traceback']))
                    
                    elif msg_type == 'status':
                        if content['execution_state'] == 'idle':
                            break
                
                except zmq.error.Again:
                    break
            
            execution_time = time.time() - start_time
            
            # Update last activity
            kernel_info.last_activity = datetime.utcnow()
            
            return ExecutionResult(
                success=True,
                stdout=''.join(stdout_parts),
                stderr=''.join(stderr_parts),
                result=result,
                execution_time=execution_time,
                exit_code=0
            )
            
        except Exception as e:
            execution_time = time.time() - start_time
            return ExecutionResult(
                success=False,
                stdout='',
                stderr=str(e),
                result=None,
                execution_time=execution_time,
                exit_code=1,
                error=str(e)
            )
    
    def shutdown_kernel(self, kernel_id: str) -> None:
        """
        Shutdown kernel and cleanup container.
        
        1. Send shutdown_request via ZMQ
        2. Wait for kernel shutdown
        3. Stop and remove container
        4. Cleanup connection file
        
        Args:
            kernel_id: Kernel to shutdown
        """
        kernel_info = self._get_kernel(kernel_id)
        
        try:
            # Shutdown kernel
            if kernel_info.client:
                kernel_info.client.shutdown()
                kernel_info.client.stop_channels()
            
            # Stop and remove container
            self.container_manager.stop_container(kernel_info.container_id)
            self.container_manager.remove_container(kernel_info.container_id)
            
            # Cleanup connection file
            kernel_info.connection_file.unlink(missing_ok=True)
            
        finally:
            # Remove from registry
            del self.kernels[kernel_id]
    
    def inspect_namespace(self, kernel_id: str) -> List[str]:
        """
        Get list of variables in kernel namespace.
        
        Executes introspection code:
        [var for var in dir() if not var.startswith('_')]
        
        Args:
            kernel_id: Kernel to inspect
        
        Returns:
            List of variable names
        """
        code = "[var for var in dir() if not var.startswith('_')]"
        result = self.execute_code(kernel_id, code, timeout=5)
        
        if result.success and result.result:
            # Parse result (it's a string representation of a list)
            try:
                return eval(result.result)  # nosec - controlled code
            except Exception:
                return []
        return []
    
    def get_variable_info(
        self,
        kernel_id: str,
        variable_name: str
    ) -> Dict[str, Any]:
        """
        Get detailed information about a variable.
        
        Executes introspection code to get:
        - type(var).__name__
        - sys.getsizeof(var) if available
        - var.shape if hasattr(var, 'shape')
        - repr(var)[:100]
        
        Args:
            kernel_id: Kernel to inspect
            variable_name: Name of variable to inspect
        
        Returns:
            Dict with type, size, shape, repr
        """
        code = f"""
import sys
_var = {variable_name}
_info = {{
    'type': type(_var).__name__,
    'repr': repr(_var)[:100],
}}
try:
    _info['size_bytes'] = sys.getsizeof(_var)
except:
    pass
if hasattr(_var, 'shape'):
    _info['shape'] = _var.shape
_info
"""
        result = self.execute_code(kernel_id, code, timeout=5)
        
        if result.success and result.result:
            try:
                return eval(result.result)  # nosec - controlled code
            except Exception:
                return {}
        return {}
    
    def restart_kernel(self, kernel_id: str) -> None:
        """
        Restart kernel (namespace reset, container kept).
        
        Strategy: shutdown current kernel and start new one in same container.
        Note: In a full implementation, we'd use KernelManager.restart_kernel().
        
        Args:
            kernel_id: Kernel to restart
        """
        kernel_info = self._get_kernel(kernel_id)
        
        # For now, just record activity - full restart implementation requires
        # KernelManager integration (not just BlockingKernelClient)
        # TODO: Implement proper kernel restart via KernelManager
        kernel_info.last_activity = datetime.utcnow()
    
    def cleanup_idle_kernels(
        self,
        idle_timeout: timedelta
    ) -> int:
        """
        Cleanup kernels idle longer than timeout.
        
        Args:
            idle_timeout: Maximum idle time before cleanup
        
        Returns:
            Number of kernels cleaned up
        """
        now = datetime.utcnow()
        cleaned_up = 0
        
        for kernel_id in list(self.kernels.keys()):
            kernel_info = self.kernels[kernel_id]
            idle_time = now - kernel_info.last_activity
            
            if idle_time > idle_timeout:
                try:
                    self.shutdown_kernel(kernel_id)
                    cleaned_up += 1
                except Exception:
                    pass  # Continue cleanup even if one fails
        
        return cleaned_up
    
    def _get_kernel(self, kernel_id: str) -> KernelInfo:
        """Get kernel info or raise error."""
        if kernel_id not in self.kernels:
            raise KernelError(f"Kernel {kernel_id} not found")
        return self.kernels[kernel_id]
    
    def _generate_connection_info(self) -> Dict[str, Any]:
        """Generate ZMQ connection information."""
        import secrets
        
        return {
            "shell_port": 0,  # Let ZMQ assign
            "iopub_port": 0,
            "stdin_port": 0,
            "control_port": 0,
            "hb_port": 0,
            "ip": "127.0.0.1",
            "key": secrets.token_hex(32),
            "transport": "tcp",
            "signature_scheme": "hmac-sha256",
            "kernel_name": "python3"
        }
    
    def _create_connection_file(
        self,
        kernel_id: str,
        connection_info: Dict[str, Any]
    ) -> Path:
        """Create connection file for kernel."""
        # Create temp file
        fd, path = tempfile.mkstemp(suffix=f"-kernel-{kernel_id}.json")
        
        # Write connection info
        with open(fd, 'w') as f:
            json.dump(connection_info, f)
        
        return Path(path)
    
    def _connect_client(self, connection_info: Dict[str, Any]) -> BlockingKernelClient:
        """Connect jupyter-client to kernel."""
        client = BlockingKernelClient()
        client.load_connection_info(connection_info)
        client.start_channels()
        return client
    
    def _verify_kernel(self, client: BlockingKernelClient, timeout: int = 10) -> bool:
        """Verify kernel is responsive."""
        try:
            client.wait_for_ready(timeout=timeout)
            return True
        except Exception:
            return False