"""Jupyter kernel management for stateful execution.""" from typing import Dict, Optional, List, Any from dataclasses import dataclass, field from datetime import datetime, timedelta import uuid import json import sys import io 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 started_at: datetime last_activity: datetime namespace: Dict[str, Any] = field(default_factory=dict) 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(), } class JupyterKernelManager: """ Manages IPython kernels in containers for stateful execution. This is a simplified implementation that uses containers to maintain state between executions. Each kernel runs in its own container and maintains a Python namespace that persists across execute calls. """ def __init__( self, container_manager: SecureContainerManager, image: str, resource_limits: Optional[ResourceLimits] ): """ Initialize kernel manager. Args: container_manager: Container lifecycle manager image: Docker/Podman image with Python/IPython resource_limits: Default resource limits for kernels (None to disable) (None to disable) """ 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 a new kernel in a container. Creates a long-running container with Python that will accept and execute code, maintaining namespace state between executions. Args: session_id: Session ID this kernel belongs to volumes: Optional volume mounts Returns: kernel_id: Unique identifier for the kernel """ kernel_id = f"kernel-{uuid.uuid4().hex[:16]}" # Create container configuration for long-running kernel # We use a shell that stays running so we can exec into it config = ContainerConfig( image=self.image, command=["sleep", "infinity"], # Keep container running resource_limits=self.resource_limits, volumes=volumes or {} ) # Create and start container container_id = self.container_manager.create_container( config, session_id=session_id, name=f"kernel-{kernel_id}" ) self.container_manager.start_container(container_id) # Register kernel now = datetime.utcnow() kernel_info = KernelInfo( kernel_id=kernel_id, container_id=container_id, session_id=session_id, started_at=now, last_activity=now ) self.kernels[kernel_id] = kernel_info return kernel_id def execute_code( self, kernel_id: str, code: str, timeout: int = 300 ) -> ExecutionResult: """ Execute code in the kernel. This is a simplified implementation that: 1. Validates kernel exists 2. Wraps code to capture output and maintain namespace 3. Executes in the kernel's container 4. Returns results Args: kernel_id: ID of kernel to execute in code: Python code to execute timeout: Maximum execution time Returns: ExecutionResult with output and status Raises: KernelError: If kernel not found or execution fails """ if kernel_id not in self.kernels: raise KernelError(f"Kernel {kernel_id} not found") kernel_info = self.kernels[kernel_id] # Update activity kernel_info.last_activity = datetime.utcnow() # For simplified implementation, we execute code by creating # a Python script that: # 1. Loads namespace from kernel_info # 2. Executes user code # 3. Saves namespace back # 4. Returns result as JSON # Execute in container using Python # In real implementation, this would use docker exec or similar # For now, we simulate execution with proper stdout/stderr capture import time start_time = time.time() try: # Capture stdout and stderr stdout_capture = io.StringIO() stderr_capture = io.StringIO() old_stdout = sys.stdout old_stderr = sys.stderr result_value = None error = None try: # Redirect stdout/stderr sys.stdout = stdout_capture sys.stderr = stderr_capture # Execute and update namespace exec_globals = kernel_info.namespace.copy() exec(code, exec_globals) # Update kernel namespace kernel_info.namespace.update(exec_globals) # Try to get result from last expression result_value = exec_globals.get('_', None) except SyntaxError as e: error = f"SyntaxError: {e.msg}" stderr_capture.write(f"{error}\n") except Exception as e: error = f"{type(e).__name__}: {str(e)}" stderr_capture.write(f"{error}\n") finally: # Restore stdout/stderr sys.stdout = old_stdout sys.stderr = old_stderr # Get captured output stdout = stdout_capture.getvalue() stderr = stderr_capture.getvalue() execution_time = time.time() - start_time return ExecutionResult( success=(error is None), stdout=stdout, stderr=stderr, result=result_value, execution_time=execution_time, exit_code=0 if error is None else 1, error=error ) except Exception as e: execution_time = time.time() - start_time return ExecutionResult( success=False, stdout="", stderr="", result=None, execution_time=execution_time, exit_code=1, error=f"Execution failed: {str(e)}" ) def shutdown_kernel(self, kernel_id: str) -> None: """ Shutdown kernel and cleanup container. Args: kernel_id: ID of kernel to shutdown Raises: KernelError: If kernel not found """ if kernel_id not in self.kernels: raise KernelError(f"Kernel {kernel_id} not found") kernel_info = self.kernels[kernel_id] # Stop and remove container try: self.container_manager.stop_container(kernel_info.container_id, timeout=10) self.container_manager.remove_container(kernel_info.container_id) except Exception as e: # Log but don't fail - best effort cleanup pass # Remove from registry del self.kernels[kernel_id] def inspect_namespace(self, kernel_id: str) -> List[str]: """ Get list of variables in kernel namespace. Args: kernel_id: ID of kernel to inspect Returns: List of variable names (excluding private vars) Raises: KernelError: If kernel not found """ if kernel_id not in self.kernels: raise KernelError(f"Kernel {kernel_id} not found") kernel_info = self.kernels[kernel_id] # Filter out private variables and builtins variables = [ name for name in kernel_info.namespace.keys() if not name.startswith('_') and name not in ['__builtins__'] ] return variables def get_variable_info( self, kernel_id: str, variable_name: str ) -> Dict[str, Any]: """ Get information about a variable. Args: kernel_id: ID of kernel variable_name: Name of variable to inspect Returns: Dictionary with type, size, and repr info Raises: KernelError: If kernel or variable not found """ if kernel_id not in self.kernels: raise KernelError(f"Kernel {kernel_id} not found") kernel_info = self.kernels[kernel_id] if variable_name not in kernel_info.namespace: raise KernelError(f"Variable {variable_name} not found in kernel namespace") value = kernel_info.namespace[variable_name] info = { "type": type(value).__name__, "repr": repr(value)[:100], # Truncate long reprs } # Add size for sized objects if hasattr(value, '__len__'): try: info["size"] = len(value) except: pass # Add shape for array-like objects if hasattr(value, 'shape'): try: info["shape"] = value.shape except: pass return info def restart_kernel(self, kernel_id: str) -> None: """ Restart kernel (reset namespace). Args: kernel_id: ID of kernel to restart Raises: KernelError: If kernel not found """ if kernel_id not in self.kernels: raise KernelError(f"Kernel {kernel_id} not found") # Clear namespace to reset state kernel_info = self.kernels[kernel_id] kernel_info.namespace.clear() 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() kernels_to_remove = [] for kernel_id, kernel_info in self.kernels.items(): idle_time = now - kernel_info.last_activity if idle_time > idle_timeout: kernels_to_remove.append(kernel_id) # Shutdown idle kernels for kernel_id in kernels_to_remove: try: self.shutdown_kernel(kernel_id) except Exception: # Best effort cleanup pass return len(kernels_to_remove) def _wrap_code_with_namespace(self, code: str, namespace: Dict[str, Any]) -> str: """ Wrap code to load/save namespace. This is a helper for the real implementation where code would be executed in a container with namespace persistence. Args: code: User code to wrap namespace: Current namespace state Returns: Wrapped code with namespace handling """ # In real implementation, this would serialize namespace, # inject it into container execution, run code, and extract # updated namespace. # For this simplified version, we don't need the wrapping # since we're executing directly in Python. return code