264 lines
8.9 KiB
Python
264 lines
8.9 KiB
Python
|
"""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}"
|
||
|
|
)
|