"""Jupyter backend for stateful code execution.""" from typing import Optional, Dict, List import hashlib from mcp_forge.config.schema import ForgeConfig from mcp_forge.podman.containers import SecureContainerManager from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity from mcp_forge.security.resource_limits import ResourceLimits, parse_memory_string from mcp_forge.execution.simple.executor import ExecutionResult from mcp_forge.execution.jupyter.kernel import JupyterKernelManager from mcp_forge.execution.jupyter.sessions import SessionManager, SessionState, SessionError class JupyterBackend: """Stateful code execution backend using Jupyter kernels.""" def __init__( self, config: ForgeConfig, container_manager: SecureContainerManager, audit_logger: AuditLogger ): """ Initialize Jupyter backend. Args: config: Forge configuration container_manager: Container lifecycle manager audit_logger: Audit logging instance """ self.config = config self.container_manager = container_manager self.audit_logger = audit_logger # Initialize kernel manager kernel_manager = JupyterKernelManager( container_manager=container_manager, image=config.images.jupyter, resource_limits=self._default_resource_limits() ) # Initialize session manager self.session_manager = SessionManager( config=config.sessions, kernel_manager=kernel_manager, audit_logger=audit_logger ) def execute( self, code: str, session_id: str, timeout: Optional[int] = None, memory: Optional[str] = None, cpu_quota: Optional[int] = None, custom_image: Optional[str] = None, volumes: Optional[Dict[str, dict]] = None, injection_code: Optional[str] = None, bridge_socket_path: Optional[str] = None ) -> ExecutionResult: """ Execute code in stateful session. Creates session if it doesn't exist, reuses existing session otherwise. Session maintains namespace state across multiple executions. Args: code: Python code to execute session_id: Unique session identifier timeout: Max execution time in seconds (uses config default if None) memory: Memory limit string (uses config default if None) cpu_quota: CPU quota (uses config default if None) custom_image: Custom image name (uses config default if None) volumes: Volume mounts dict injection_code: Optional MCP tool injection code (executed once at session start) bridge_socket_path: Optional path to MCP bridge socket for mounting Returns: ExecutionResult with execution output and metadata Raises: ValueError: If limits exceed configured maximums SessionError: If session operation fails """ # Use defaults from config if not specified timeout = timeout if timeout is not None else self.config.execution.default_timeout memory = memory if memory is not None else self.config.execution.default_memory cpu_quota = cpu_quota if cpu_quota is not None else self.config.execution.default_cpu_quota # Validate limits against maximums self._validate_limits(timeout, memory, cpu_quota) # Log execution (hash code, don't log actual content) code_hash = hashlib.sha256(code.encode()).hexdigest() self.audit_logger.log( event_type=AuditEventType.EXECUTION_REQUEST, severity=AuditSeverity.INFO, message="Stateful code execution requested", session_id=session_id, details={ "code_hash": code_hash, "timeout": timeout, "memory": memory, "cpu_quota": cpu_quota } ) # Check if session exists, create if needed try: self.session_manager.get_session(session_id) except SessionError: # Session doesn't exist, create it with MCP injection resource_limits = ResourceLimits( memory=memory, cpu_quota=cpu_quota, storage="1g", # Default storage quota timeout=timeout ) self.session_manager.create_session( session_id=session_id, resource_limits=resource_limits, volumes=volumes, injection_code=injection_code, bridge_socket_path=bridge_socket_path ) # Execute in session result = self.session_manager.execute_in_session( session_id=session_id, code=code, timeout=timeout ) return result def document_state( self, session_id: str, variables: Dict[str, str], note: str = "", clear: bool = False ) -> dict: """ Document important variables in session. Args: session_id: Session to document variables: Dictionary of variable_name -> description note: Optional note about session state clear: If True, replace all documented variables; if False, merge Returns: Dictionary with updated state info Raises: SessionError: If session doesn't exist """ self.session_manager.document_state( session_id=session_id, variables=variables, note=note, clear=clear ) # Return updated state state = self.session_manager.get_session_state(session_id) return state.to_dict() def get_session_state(self, session_id: str) -> SessionState: """ Get documented state for session. Args: session_id: Session identifier Returns: SessionState object Raises: SessionError: If session doesn't exist """ return self.session_manager.get_session_state(session_id) def destroy_session(self, session_id: str) -> None: """ Destroy session and cleanup kernel. Args: session_id: Session to destroy Raises: SessionError: If session doesn't exist """ self.session_manager.destroy_session(session_id) def list_sessions(self) -> List[dict]: """ List all active sessions with metadata. Returns: List of session dictionaries """ return self.session_manager.list_sessions() def cleanup_idle_sessions(self) -> int: """ Cleanup sessions idle beyond configured timeout. Returns: Number of sessions cleaned up """ return self.session_manager.cleanup_idle_sessions() def _default_resource_limits(self) -> Optional[ResourceLimits]: """ Get default resource limits from config. Returns: ResourceLimits with config defaults, or None if enforcement disabled """ if not self.config.security.enforce_resource_limits: return None return ResourceLimits( memory=self.config.execution.default_memory, cpu_quota=self.config.execution.default_cpu_quota, storage="1g", timeout=self.config.execution.default_timeout ) def _validate_limits(self, timeout: int, memory: str, cpu_quota: int) -> None: """ Validate resource limits against configured maximums. Args: timeout: Timeout in seconds memory: Memory limit string cpu_quota: CPU quota value Raises: ValueError: If any limit exceeds maximum """ # Validate timeout if timeout > self.config.execution.max_timeout: raise ValueError( f"Timeout {timeout} exceeds maximum {self.config.execution.max_timeout}" ) # Validate memory memory_bytes = parse_memory_string(memory) max_memory_bytes = parse_memory_string(self.config.execution.max_memory) if memory_bytes > max_memory_bytes: raise ValueError( f"Memory {memory} exceeds maximum {self.config.execution.max_memory}" ) # Validate CPU quota if cpu_quota > self.config.execution.max_cpu_quota: raise ValueError( f"CPU quota {cpu_quota} exceeds maximum {self.config.execution.max_cpu_quota}" )