""" 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