feat: Implement session management and secure container lifecycle management
- Added session management for stateful execution in `sessions.py`, including session creation, state documentation, and cleanup of idle sessions. - Introduced `SecureContainerManager` in `containers.py` for managing container lifecycle with security enforcement, including creation, starting, stopping, and removal of containers. - Updated server initialization to use the new session manager. - Enhanced container configuration to skip resource limits for very high values, indicating no enforcement. - Improved logging capabilities for container operations in the audit logger. - Refactored Jupyter backend to integrate with the new session management and resource limits handling.
This commit is contained in:
Hans Aschauer
parent
9ceeaa1eda
commit
7d9efc5a38
10 changed files with 3384 additions and 26 deletions
1448
docs/architecture2.md
Normal file
1448
docs/architecture2.md
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -42,28 +42,40 @@ class SimpleBackend:
|
|||
self.config = config
|
||||
|
||||
# Create default resource limits from config
|
||||
# When not enforcing limits, use very high values to effectively disable
|
||||
if config.security.enforce_resource_limits:
|
||||
self.default_limits = ResourceLimits(
|
||||
memory=config.execution.default_memory,
|
||||
storage="10g", # Default storage limit
|
||||
storage="10g",
|
||||
cpu_quota=config.execution.default_cpu_quota,
|
||||
timeout=config.execution.default_timeout
|
||||
)
|
||||
else:
|
||||
# No resource enforcement - use very high limits (effectively unlimited)
|
||||
self.default_limits = ResourceLimits(
|
||||
memory="16g", # Very high memory limit
|
||||
storage="100g", # Very high storage limit
|
||||
cpu_quota=1000000, # Effectively unlimited CPU
|
||||
timeout=config.execution.default_timeout
|
||||
)
|
||||
|
||||
# Create executor with default image
|
||||
# Create executor with default image (prefer python_3_12)
|
||||
self.executor = CodeExecutor(
|
||||
container_manager=container_manager,
|
||||
image=config.images.python,
|
||||
image=config.images.python_3_12,
|
||||
resource_limits=self.default_limits
|
||||
)
|
||||
|
||||
logger.debug(f"SimpleBackend initialized with image={config.images.python}")
|
||||
logger.debug(f"SimpleBackend initialized with image={config.images.python_3_12}")
|
||||
|
||||
def execute(
|
||||
self,
|
||||
code: str,
|
||||
timeout: Optional[int] = None,
|
||||
memory: Optional[str] = None,
|
||||
cpu_quota: Optional[int] = None
|
||||
cpu_quota: Optional[int] = None,
|
||||
injection_code: Optional[str] = None,
|
||||
bridge_socket_path: Optional[str] = None
|
||||
) -> ExecutionResult:
|
||||
"""
|
||||
Execute Python code in isolated container.
|
||||
|
|
@ -73,6 +85,8 @@ class SimpleBackend:
|
|||
timeout: Optional timeout override (seconds)
|
||||
memory: Optional memory limit override (e.g., "512m")
|
||||
cpu_quota: Optional CPU quota override
|
||||
injection_code: Optional code to inject before user code (for MCP tools)
|
||||
bridge_socket_path: Optional path to MCP bridge socket
|
||||
|
||||
Returns:
|
||||
ExecutionResult with stdout, stderr, result, etc.
|
||||
|
|
@ -88,10 +102,20 @@ class SimpleBackend:
|
|||
# Create temporary executor with custom limits
|
||||
executor = CodeExecutor(
|
||||
container_manager=self.container_manager,
|
||||
image=self.config.images.python,
|
||||
image=self.config.images.python_3_12,
|
||||
resource_limits=limits
|
||||
)
|
||||
return executor.execute(code, timeout=timeout)
|
||||
|
||||
return executor.execute(
|
||||
code=code,
|
||||
timeout=timeout,
|
||||
injection_code=injection_code,
|
||||
bridge_socket_path=bridge_socket_path
|
||||
)
|
||||
else:
|
||||
# Use default executor
|
||||
return self.executor.execute(code, timeout=timeout)
|
||||
return self.executor.execute(
|
||||
code=code,
|
||||
timeout=timeout,
|
||||
injection_code=injection_code,
|
||||
bridge_socket_path=bridge_socket_path
|
||||
)
|
||||
|
|
|
|||
263
src/mcp_forge/execution/jupyter/backend.py
Normal file
263
src/mcp_forge/execution/jupyter/backend.py
Normal file
|
|
@ -0,0 +1,263 @@
|
|||
"""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}"
|
||||
)
|
||||
631
src/mcp_forge/execution/jupyter/kernel.py
Normal file
631
src/mcp_forge/execution/jupyter/kernel.py
Normal file
|
|
@ -0,0 +1,631 @@
|
|||
"""
|
||||
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
|
||||
import socket
|
||||
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,
|
||||
injection_code: Optional[str] = None,
|
||||
bridge_socket_path: Optional[str] = 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. Mount bridge socket if provided (for MCP tools)
|
||||
5. Start container
|
||||
6. Wait for kernel to be ready
|
||||
7. Connect jupyter-client to kernel via ZMQ
|
||||
8. Execute injection code (MCP tools setup) if provided
|
||||
9. Verify kernel is responsive
|
||||
|
||||
Args:
|
||||
session_id: Session ID this kernel belongs to
|
||||
volumes: Optional volume mounts
|
||||
injection_code: Optional MCP tool injection code to execute at startup
|
||||
bridge_socket_path: Optional path to MCP bridge socket for mounting
|
||||
|
||||
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:
|
||||
# Set up volumes (user volumes + bridge socket + connection file)
|
||||
container_volumes = volumes.copy() if volumes else {}
|
||||
if bridge_socket_path:
|
||||
container_volumes[bridge_socket_path] = {
|
||||
"bind": bridge_socket_path,
|
||||
"mode": "rw"
|
||||
}
|
||||
|
||||
# Mount connection file into container
|
||||
container_connection_path = f"/tmp/kernel-{kernel_id}.json"
|
||||
container_volumes[str(connection_file)] = {
|
||||
"bind": container_connection_path,
|
||||
"mode": "ro"
|
||||
}
|
||||
|
||||
# Create container with ipykernel using host networking
|
||||
config = ContainerConfig(
|
||||
image=self.image,
|
||||
command=[
|
||||
"python", "-m", "ipykernel_launcher",
|
||||
"-f", container_connection_path
|
||||
],
|
||||
resource_limits=self.resource_limits,
|
||||
volumes=container_volumes,
|
||||
network_mode="host", # Use host network for ZMQ communication
|
||||
port_bindings={
|
||||
connection_info["shell_port"]: connection_info["shell_port"],
|
||||
connection_info["iopub_port"]: connection_info["iopub_port"],
|
||||
connection_info["stdin_port"]: connection_info["stdin_port"],
|
||||
connection_info["control_port"]: connection_info["control_port"],
|
||||
connection_info["hb_port"]: connection_info["hb_port"],
|
||||
}
|
||||
)
|
||||
|
||||
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 with polling
|
||||
if not self._wait_for_kernel_ready(connection_info, timeout=30):
|
||||
raise KernelError(f"Kernel {kernel_id} failed to start within timeout")
|
||||
|
||||
# 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")
|
||||
|
||||
# Execute injection code if provided (MCP tools setup)
|
||||
if injection_code:
|
||||
self._execute_injection_code(client, injection_code, kernel_id)
|
||||
|
||||
# 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
|
||||
has_error = False
|
||||
|
||||
# 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':
|
||||
has_error = True
|
||||
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()
|
||||
|
||||
stderr_text = ''.join(stderr_parts)
|
||||
|
||||
return ExecutionResult(
|
||||
success=(not has_error),
|
||||
stdout=''.join(stdout_parts),
|
||||
stderr=stderr_text,
|
||||
result=result,
|
||||
execution_time=execution_time,
|
||||
exit_code=1 if has_error else 0,
|
||||
error=stderr_text if has_error else None
|
||||
)
|
||||
|
||||
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 with allocated ports."""
|
||||
import secrets
|
||||
|
||||
# Allocate 5 ports for ZMQ channels
|
||||
ports = self._allocate_ports(5)
|
||||
|
||||
return {
|
||||
"shell_port": ports[0],
|
||||
"iopub_port": ports[1],
|
||||
"stdin_port": ports[2],
|
||||
"control_port": ports[3],
|
||||
"hb_port": ports[4],
|
||||
"ip": "127.0.0.1",
|
||||
"key": secrets.token_hex(32),
|
||||
"transport": "tcp",
|
||||
"signature_scheme": "hmac-sha256",
|
||||
"kernel_name": "python3"
|
||||
}
|
||||
|
||||
def _allocate_ports(self, count: int) -> List[int]:
|
||||
"""
|
||||
Allocate available ports for ZMQ.
|
||||
|
||||
Args:
|
||||
count: Number of ports to allocate
|
||||
|
||||
Returns:
|
||||
List of allocated port numbers
|
||||
"""
|
||||
ports = []
|
||||
for _ in range(count):
|
||||
# Let OS assign available port
|
||||
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
sock.bind(('127.0.0.1', 0)) # Bind to any available port
|
||||
port = sock.getsockname()[1]
|
||||
sock.close()
|
||||
ports.append(port)
|
||||
return ports
|
||||
|
||||
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
|
||||
|
||||
def _wait_for_kernel_ready(
|
||||
self,
|
||||
connection_info: Dict[str, Any],
|
||||
timeout: int = 30,
|
||||
poll_interval: float = 0.5
|
||||
) -> bool:
|
||||
"""
|
||||
Wait for kernel to be ready by polling ports.
|
||||
|
||||
Args:
|
||||
connection_info: Kernel connection information
|
||||
timeout: Maximum time to wait in seconds
|
||||
poll_interval: Time between polls in seconds
|
||||
|
||||
Returns:
|
||||
True if kernel is ready, False if timeout
|
||||
"""
|
||||
start_time = time.time()
|
||||
shell_port = connection_info["shell_port"]
|
||||
|
||||
while time.time() - start_time < timeout:
|
||||
try:
|
||||
# Try to connect to shell port
|
||||
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
sock.settimeout(1)
|
||||
result = sock.connect_ex(('127.0.0.1', shell_port))
|
||||
sock.close()
|
||||
|
||||
if result == 0:
|
||||
# Port is open, kernel is ready
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
time.sleep(poll_interval)
|
||||
|
||||
return False
|
||||
|
||||
def _execute_injection_code(
|
||||
self,
|
||||
client: BlockingKernelClient,
|
||||
injection_code: str,
|
||||
kernel_id: str
|
||||
) -> None:
|
||||
"""
|
||||
Execute MCP tool injection code on kernel startup.
|
||||
|
||||
This runs once when the kernel starts to set up MCP tools.
|
||||
Unlike regular code execution, we don't capture output.
|
||||
|
||||
Args:
|
||||
client: Connected kernel client
|
||||
injection_code: Python code to inject (MCP tools setup)
|
||||
kernel_id: Kernel ID for error messages
|
||||
|
||||
Raises:
|
||||
KernelError: If injection code fails to execute
|
||||
"""
|
||||
try:
|
||||
# Execute injection code silently
|
||||
_msg_id = client.execute(injection_code, silent=True, store_history=False)
|
||||
|
||||
# Wait for execution to complete
|
||||
timeout = 10 # Injection should be fast
|
||||
while True:
|
||||
try:
|
||||
msg = client.get_iopub_msg(timeout=timeout)
|
||||
msg_type = msg['header']['msg_type']
|
||||
|
||||
if msg_type == 'error':
|
||||
content = msg['content']
|
||||
error_msg = '\n'.join(content.get('traceback', [str(content)]))
|
||||
raise KernelError(
|
||||
f"MCP injection failed in kernel {kernel_id}: {error_msg}"
|
||||
)
|
||||
|
||||
elif msg_type == 'status':
|
||||
if msg['content']['execution_state'] == 'idle':
|
||||
break # Injection complete
|
||||
|
||||
except zmq.error.Again:
|
||||
break # Timeout, assume success
|
||||
|
||||
except KernelError:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise KernelError(
|
||||
f"Failed to execute MCP injection code in kernel {kernel_id}: {e}"
|
||||
) from e
|
||||
435
src/mcp_forge/execution/jupyter/sessions.py
Normal file
435
src/mcp_forge/execution/jupyter/sessions.py
Normal file
|
|
@ -0,0 +1,435 @@
|
|||
"""Session management for stateful execution."""
|
||||
|
||||
from typing import Dict, Optional, List, Any
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
from mcp_forge.execution.jupyter.kernel import JupyterKernelManager
|
||||
from mcp_forge.config.schema import SessionConfig
|
||||
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||
from mcp_forge.security.resource_limits import ResourceLimits
|
||||
from mcp_forge.execution.simple.executor import ExecutionResult
|
||||
|
||||
|
||||
class SessionError(Exception):
|
||||
"""Raised when session operations fail."""
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class SessionState:
|
||||
"""Documented state for a session."""
|
||||
session_id: str
|
||||
documented_variables: Dict[str, str] = field(default_factory=dict)
|
||||
note: str = ""
|
||||
last_updated: datetime = field(default_factory=datetime.utcnow)
|
||||
all_variables: List[str] = field(default_factory=list)
|
||||
introspection: Dict[str, dict] = field(default_factory=dict)
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
"""Convert to dictionary for JSON serialization."""
|
||||
return {
|
||||
"session_id": self.session_id,
|
||||
"documented_variables": self.documented_variables,
|
||||
"note": self.note,
|
||||
"last_updated": self.last_updated.isoformat(),
|
||||
"all_variables": self.all_variables,
|
||||
"introspection": self.introspection
|
||||
}
|
||||
|
||||
|
||||
class Session:
|
||||
"""Stateful execution session."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
session_id: str,
|
||||
kernel_id: str,
|
||||
created_at: datetime,
|
||||
resource_limits: ResourceLimits
|
||||
):
|
||||
"""
|
||||
Initialize session.
|
||||
|
||||
Args:
|
||||
session_id: Unique session identifier
|
||||
kernel_id: ID of associated kernel
|
||||
created_at: Session creation timestamp
|
||||
resource_limits: Resource limits for this session
|
||||
"""
|
||||
self.session_id = session_id
|
||||
self.kernel_id = kernel_id
|
||||
self.created_at = created_at
|
||||
self.last_activity = created_at
|
||||
self.resource_limits = resource_limits
|
||||
self.state = SessionState(session_id=session_id)
|
||||
self.documented_variables: Dict[str, str] = {}
|
||||
self.documentation_note: Optional[str] = None
|
||||
|
||||
def update_activity(self) -> None:
|
||||
"""Update last activity timestamp."""
|
||||
self.last_activity = datetime.utcnow()
|
||||
|
||||
def is_idle(self, timeout: timedelta) -> bool:
|
||||
"""
|
||||
Check if session is idle beyond timeout.
|
||||
|
||||
Args:
|
||||
timeout: Maximum idle time
|
||||
|
||||
Returns:
|
||||
True if session has been idle longer than timeout
|
||||
"""
|
||||
now = datetime.utcnow()
|
||||
idle_time = now - self.last_activity
|
||||
return idle_time > timeout
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
"""Convert to dictionary for serialization."""
|
||||
return {
|
||||
"session_id": self.session_id,
|
||||
"kernel_id": self.kernel_id,
|
||||
"created_at": self.created_at.isoformat(),
|
||||
"last_activity": self.last_activity.isoformat(),
|
||||
"state": self.state.to_dict()
|
||||
}
|
||||
|
||||
|
||||
class SessionManager:
|
||||
"""Manages stateful execution sessions."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: SessionConfig,
|
||||
kernel_manager: JupyterKernelManager,
|
||||
audit_logger: AuditLogger
|
||||
):
|
||||
"""
|
||||
Initialize session manager.
|
||||
|
||||
Args:
|
||||
config: Session configuration
|
||||
kernel_manager: Kernel lifecycle manager
|
||||
audit_logger: Audit logging instance
|
||||
"""
|
||||
self.config = config
|
||||
self.kernel_manager = kernel_manager
|
||||
self.audit_logger = audit_logger
|
||||
self.sessions: Dict[str, Session] = {}
|
||||
|
||||
def create_session(
|
||||
self,
|
||||
session_id: str,
|
||||
resource_limits: ResourceLimits,
|
||||
volumes: Optional[Dict[str, dict]] = None,
|
||||
injection_code: Optional[str] = None,
|
||||
bridge_socket_path: Optional[str] = None
|
||||
) -> Session:
|
||||
"""
|
||||
Create new stateful session.
|
||||
|
||||
Args:
|
||||
session_id: Unique identifier for session
|
||||
resource_limits: Resource limits for session
|
||||
volumes: Optional volume mounts
|
||||
injection_code: Optional MCP tool injection code to execute at startup
|
||||
bridge_socket_path: Optional path to MCP bridge socket for mounting
|
||||
|
||||
Returns:
|
||||
Created Session object
|
||||
|
||||
Raises:
|
||||
SessionError: If session_id already exists
|
||||
SessionError: If max concurrent sessions exceeded
|
||||
"""
|
||||
if session_id in self.sessions:
|
||||
raise SessionError(f"Session {session_id} already exists")
|
||||
|
||||
# Check max concurrent limit
|
||||
self._enforce_max_concurrent()
|
||||
|
||||
# Start kernel with MCP injection if provided
|
||||
kernel_id = self.kernel_manager.start_kernel(
|
||||
session_id,
|
||||
volumes=volumes,
|
||||
injection_code=injection_code,
|
||||
bridge_socket_path=bridge_socket_path
|
||||
)
|
||||
|
||||
# Create session
|
||||
now = datetime.utcnow()
|
||||
session = Session(
|
||||
session_id=session_id,
|
||||
kernel_id=kernel_id,
|
||||
created_at=now,
|
||||
resource_limits=resource_limits
|
||||
)
|
||||
|
||||
self.sessions[session_id] = session
|
||||
|
||||
# Log session creation
|
||||
self.audit_logger.log(
|
||||
event_type=AuditEventType.SESSION_CREATE,
|
||||
severity=AuditSeverity.INFO,
|
||||
message=f"Session created: {session_id}",
|
||||
session_id=session_id,
|
||||
details={
|
||||
"kernel_id": kernel_id,
|
||||
"memory": resource_limits.memory_bytes,
|
||||
"cpu_quota": resource_limits.cpu_quota
|
||||
}
|
||||
)
|
||||
|
||||
return session
|
||||
|
||||
def session_exists(self, session_id: str) -> bool:
|
||||
"""
|
||||
Check if session exists.
|
||||
|
||||
Args:
|
||||
session_id: Session identifier
|
||||
|
||||
Returns:
|
||||
True if session exists, False otherwise
|
||||
"""
|
||||
return session_id in self.sessions
|
||||
|
||||
def get_session(self, session_id: str) -> Session:
|
||||
"""
|
||||
Get session by ID.
|
||||
|
||||
Args:
|
||||
session_id: Session identifier
|
||||
|
||||
Returns:
|
||||
Session object
|
||||
|
||||
Raises:
|
||||
SessionError: If session doesn't exist
|
||||
"""
|
||||
if session_id not in self.sessions:
|
||||
raise SessionError(f"Session {session_id} not found")
|
||||
|
||||
return self.sessions[session_id]
|
||||
|
||||
async def document_variables(
|
||||
self,
|
||||
session_id: str,
|
||||
variables: Dict[str, str],
|
||||
note: Optional[str] = None,
|
||||
clear: bool = False
|
||||
) -> Dict:
|
||||
"""
|
||||
Document important variables in a session.
|
||||
|
||||
Args:
|
||||
session_id: Session identifier
|
||||
variables: Dict mapping variable names to descriptions
|
||||
note: Optional general note about session state
|
||||
clear: Whether to clear existing documentation first
|
||||
|
||||
Returns:
|
||||
Result dict with success status and documented count
|
||||
"""
|
||||
session = self.get_session(session_id)
|
||||
|
||||
if clear:
|
||||
session.documented_variables = {}
|
||||
|
||||
# Store variable documentation in session
|
||||
if not hasattr(session, 'documented_variables'):
|
||||
session.documented_variables = {}
|
||||
|
||||
session.documented_variables.update(variables)
|
||||
|
||||
if note:
|
||||
session.documentation_note = note
|
||||
|
||||
return {
|
||||
"success": True,
|
||||
"documented_count": len(variables),
|
||||
"total_documented": len(session.documented_variables)
|
||||
}
|
||||
|
||||
def execute_in_session(
|
||||
self,
|
||||
session_id: str,
|
||||
code: str,
|
||||
timeout: int = 300
|
||||
) -> ExecutionResult:
|
||||
"""
|
||||
Execute code in session kernel.
|
||||
|
||||
Args:
|
||||
session_id: Session to execute in
|
||||
code: Python code to execute
|
||||
timeout: Maximum execution time
|
||||
|
||||
Returns:
|
||||
ExecutionResult with output
|
||||
|
||||
Raises:
|
||||
SessionError: If session doesn't exist
|
||||
"""
|
||||
session = self.get_session(session_id)
|
||||
|
||||
# Update activity
|
||||
session.update_activity()
|
||||
|
||||
# Execute in kernel
|
||||
result = self.kernel_manager.execute_code(
|
||||
session.kernel_id,
|
||||
code,
|
||||
timeout=timeout
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
def document_state(
|
||||
self,
|
||||
session_id: str,
|
||||
variables: Dict[str, str],
|
||||
note: str = "",
|
||||
clear: bool = False
|
||||
) -> None:
|
||||
"""
|
||||
Document important variables in session.
|
||||
|
||||
Updates session.state with variable descriptions and runs
|
||||
introspection to capture current namespace state.
|
||||
|
||||
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
|
||||
|
||||
Raises:
|
||||
SessionError: If session doesn't exist
|
||||
"""
|
||||
session = self.get_session(session_id)
|
||||
|
||||
# Update documented variables
|
||||
if clear:
|
||||
session.state.documented_variables = variables.copy()
|
||||
else:
|
||||
session.state.documented_variables.update(variables)
|
||||
|
||||
# Update note if provided
|
||||
if note:
|
||||
session.state.note = note
|
||||
|
||||
# Run introspection to get current namespace state
|
||||
session.state.all_variables = self.kernel_manager.inspect_namespace(session.kernel_id)
|
||||
|
||||
# Get variable info for documented variables
|
||||
session.state.introspection = {}
|
||||
for var_name in variables.keys():
|
||||
if var_name in session.state.all_variables:
|
||||
try:
|
||||
info = self.kernel_manager.get_variable_info(session.kernel_id, var_name)
|
||||
session.state.introspection[var_name] = info
|
||||
except Exception:
|
||||
# Variable might not exist yet
|
||||
pass
|
||||
|
||||
# Update timestamp
|
||||
session.state.last_updated = datetime.utcnow()
|
||||
session.update_activity()
|
||||
|
||||
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
|
||||
"""
|
||||
session = self.get_session(session_id)
|
||||
return session.state
|
||||
|
||||
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
|
||||
"""
|
||||
session = self.get_session(session_id)
|
||||
|
||||
# Shutdown kernel
|
||||
try:
|
||||
self.kernel_manager.shutdown_kernel(session.kernel_id)
|
||||
except Exception as e:
|
||||
# Log but continue with cleanup
|
||||
self.audit_logger.log(
|
||||
event_type=AuditEventType.SESSION_DESTROY,
|
||||
severity=AuditSeverity.WARNING,
|
||||
message=f"Error shutting down kernel for session {session_id}",
|
||||
session_id=session_id,
|
||||
error=str(e)
|
||||
)
|
||||
|
||||
# Remove session
|
||||
del self.sessions[session_id]
|
||||
|
||||
# Log destruction
|
||||
self.audit_logger.log(
|
||||
event_type=AuditEventType.SESSION_DESTROY,
|
||||
severity=AuditSeverity.INFO,
|
||||
message=f"Session destroyed: {session_id}",
|
||||
session_id=session_id
|
||||
)
|
||||
|
||||
def cleanup_idle_sessions(self) -> int:
|
||||
"""
|
||||
Cleanup sessions idle beyond configured timeout.
|
||||
|
||||
Returns:
|
||||
Number of sessions cleaned up
|
||||
"""
|
||||
timeout = timedelta(seconds=self.config.idle_timeout)
|
||||
sessions_to_remove = []
|
||||
|
||||
for session_id, session in self.sessions.items():
|
||||
if session.is_idle(timeout):
|
||||
sessions_to_remove.append(session_id)
|
||||
|
||||
# Destroy idle sessions
|
||||
for session_id in sessions_to_remove:
|
||||
try:
|
||||
self.destroy_session(session_id)
|
||||
except Exception:
|
||||
# Best effort cleanup
|
||||
pass
|
||||
|
||||
return len(sessions_to_remove)
|
||||
|
||||
def list_sessions(self) -> List[dict]:
|
||||
"""
|
||||
List all active sessions with metadata.
|
||||
|
||||
Returns:
|
||||
List of session dictionaries
|
||||
"""
|
||||
return [session.to_dict() for session in self.sessions.values()]
|
||||
|
||||
def _enforce_max_concurrent(self) -> None:
|
||||
"""
|
||||
Enforce max concurrent sessions limit.
|
||||
|
||||
Raises:
|
||||
SessionError: If at max concurrent sessions
|
||||
"""
|
||||
if len(self.sessions) >= self.config.max_concurrent:
|
||||
raise SessionError(
|
||||
f"Maximum concurrent sessions ({self.config.max_concurrent}) reached"
|
||||
)
|
||||
508
src/mcp_forge/podman/containers.py
Normal file
508
src/mcp_forge/podman/containers.py
Normal file
|
|
@ -0,0 +1,508 @@
|
|||
"""
|
||||
Secure container management with security enforcement.
|
||||
|
||||
All container operations are validated against security policy
|
||||
before being sent to Podman. Provides lifecycle management
|
||||
with comprehensive audit logging.
|
||||
"""
|
||||
|
||||
from typing import Optional, Dict, List
|
||||
from datetime import datetime, timedelta
|
||||
from pathlib import Path
|
||||
|
||||
from mcp_forge.podman.client import PodmanClient
|
||||
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||
from mcp_forge.security.resource_limits import ResourceLimits
|
||||
|
||||
|
||||
class ContainerConfig:
|
||||
"""Container configuration with security defaults."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
image: str,
|
||||
command: Optional[List[str]] = None,
|
||||
environment: Optional[Dict[str, str]] = None,
|
||||
volumes: Optional[Dict[str, dict]] = None,
|
||||
resource_limits: Optional[ResourceLimits] = None,
|
||||
working_dir: Optional[str] = None,
|
||||
user: str = "1000:1000",
|
||||
network_mode: str = "none",
|
||||
port_bindings: Optional[Dict[str, int]] = None
|
||||
):
|
||||
"""
|
||||
Initialize container configuration.
|
||||
|
||||
Args:
|
||||
image: Container image to use
|
||||
command: Command to run in container
|
||||
environment: Environment variables
|
||||
volumes: Volume mounts (host_path -> {bind, mode})
|
||||
resource_limits: Resource limits to apply
|
||||
working_dir: Working directory in container (None to use image default)
|
||||
user: User to run as (UID:GID)
|
||||
network_mode: Network mode (none, host, bridge). Default is 'none' for security.
|
||||
port_bindings: Port mappings for network_mode=host (container_port -> host_port)
|
||||
"""
|
||||
self.image = image
|
||||
self.command = command or []
|
||||
self.environment = environment or {}
|
||||
self.volumes = volumes or {}
|
||||
self.resource_limits = resource_limits
|
||||
self.working_dir = working_dir
|
||||
self.user = user
|
||||
self.network_mode = network_mode
|
||||
self.port_bindings = port_bindings or {}
|
||||
|
||||
def to_podman_params(self) -> dict:
|
||||
"""
|
||||
Convert to Podman container create parameters.
|
||||
|
||||
Ensures all security requirements are included:
|
||||
- network_mode: configurable (default 'none' for security)
|
||||
- read_only: True
|
||||
- security_opt: ["no-new-privileges"]
|
||||
- resource limits
|
||||
- port_bindings: for host networking mode
|
||||
|
||||
Returns:
|
||||
Dictionary of parameters for Podman containers.create()
|
||||
"""
|
||||
params = {
|
||||
"image": self.image,
|
||||
"command": self.command if self.command else None,
|
||||
"environment": self.environment,
|
||||
"user": self.user,
|
||||
# Security requirements
|
||||
"network_mode": self.network_mode,
|
||||
"read_only": True,
|
||||
"security_opt": ["no-new-privileges"],
|
||||
}
|
||||
|
||||
# Add working_dir only if explicitly set
|
||||
if self.working_dir is not None:
|
||||
params["working_dir"] = self.working_dir
|
||||
|
||||
# Add port bindings if using host network mode
|
||||
# Note: In host mode, ports are directly accessible
|
||||
# port_bindings are informational for tracking
|
||||
if self.network_mode == "host" and self.port_bindings:
|
||||
# With host networking, container uses host's network stack directly
|
||||
# No explicit port mapping needed, but we track for documentation
|
||||
pass
|
||||
|
||||
# Add volumes if present
|
||||
if self.volumes:
|
||||
params["volumes"] = self.volumes
|
||||
|
||||
# Add resource limits if present
|
||||
if self.resource_limits:
|
||||
limit_params = self.resource_limits.to_podman_params()
|
||||
params.update(limit_params)
|
||||
# Disable swap to avoid cgroup swap.max issues on some systems
|
||||
if "mem_limit" in params:
|
||||
params["memswap_limit"] = -1 # Disable swap
|
||||
|
||||
return params
|
||||
|
||||
|
||||
class SecureContainerManager:
|
||||
"""Manages container lifecycle with security enforcement."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
podman_client: PodmanClient,
|
||||
validator: OperationValidator,
|
||||
audit_logger: AuditLogger
|
||||
):
|
||||
"""
|
||||
Initialize secure container manager.
|
||||
|
||||
Args:
|
||||
podman_client: Podman client wrapper
|
||||
validator: Operation validator for security checks
|
||||
audit_logger: Audit logger for operation logging
|
||||
"""
|
||||
self.podman = podman_client
|
||||
self.validator = validator
|
||||
self.audit_logger = audit_logger
|
||||
|
||||
def create_container(
|
||||
self,
|
||||
config: ContainerConfig,
|
||||
session_id: Optional[str] = None,
|
||||
name: Optional[str] = None,
|
||||
**extra_params
|
||||
) -> str:
|
||||
"""
|
||||
Create a container with security validation.
|
||||
|
||||
Args:
|
||||
config: Container configuration
|
||||
session_id: Session ID for tracking
|
||||
name: Optional container name
|
||||
**extra_params: Additional parameters (checked for forbidden values)
|
||||
|
||||
Returns:
|
||||
Container ID
|
||||
|
||||
Raises:
|
||||
SecurityError: If configuration violates security policy
|
||||
"""
|
||||
# Convert config to Podman parameters
|
||||
params = config.to_podman_params()
|
||||
|
||||
# Add session label if provided
|
||||
labels = {}
|
||||
if session_id:
|
||||
labels["mcp-forge.session"] = session_id
|
||||
if labels:
|
||||
params["labels"] = labels
|
||||
|
||||
if name:
|
||||
params["name"] = name
|
||||
|
||||
# Merge any extra parameters (will be validated)
|
||||
params.update(extra_params)
|
||||
|
||||
# Validate against security policy
|
||||
try:
|
||||
# Extract image from params for validation
|
||||
self.validator.validate_container_create(
|
||||
image=config.image,
|
||||
params=params,
|
||||
session_id=session_id
|
||||
)
|
||||
except SecurityError as e:
|
||||
# Log security violation
|
||||
self.audit_logger.log_security_violation(
|
||||
operation="container_create",
|
||||
reason=str(e),
|
||||
session_id=session_id
|
||||
)
|
||||
raise
|
||||
|
||||
# Create container
|
||||
try:
|
||||
container = self.podman.client.containers.create(**params)
|
||||
container_id = container.id
|
||||
|
||||
# Register with validator
|
||||
if session_id:
|
||||
self.validator.register_session_container(container_id)
|
||||
|
||||
# Log successful creation
|
||||
self.audit_logger.log_container_operation(
|
||||
operation="create",
|
||||
container_id=container_id,
|
||||
image=config.image,
|
||||
session_id=session_id,
|
||||
details={
|
||||
"name": name,
|
||||
"command": config.command
|
||||
}
|
||||
)
|
||||
|
||||
return container_id
|
||||
|
||||
except Exception as e:
|
||||
self.audit_logger.log(
|
||||
event_type=AuditEventType.CONTAINER_CREATE,
|
||||
severity=AuditSeverity.ERROR,
|
||||
message=f"Container creation failed: {e}",
|
||||
details={
|
||||
"image": config.image,
|
||||
"session_id": session_id,
|
||||
"error": str(e)
|
||||
}
|
||||
)
|
||||
raise
|
||||
|
||||
def start_container(self, container_id: str) -> None:
|
||||
"""
|
||||
Start a container.
|
||||
|
||||
Args:
|
||||
container_id: Container ID to start
|
||||
|
||||
Raises:
|
||||
SecurityError: If container is not a session container
|
||||
"""
|
||||
# Verify container is registered (security check)
|
||||
if container_id not in self.validator.session_containers:
|
||||
self.audit_logger.log_security_violation(
|
||||
operation="container_start",
|
||||
reason=f"Attempted to start unregistered container: {container_id}"
|
||||
)
|
||||
raise SecurityError(
|
||||
f"Container {container_id} is not a registered session container"
|
||||
)
|
||||
|
||||
try:
|
||||
container = self.podman.client.containers.get(container_id)
|
||||
container.start()
|
||||
|
||||
self.audit_logger.log_container_operation(
|
||||
operation="start",
|
||||
container_id=container_id,
|
||||
image="" # Not available without extra lookup
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.audit_logger.log(
|
||||
event_type=AuditEventType.CONTAINER_START,
|
||||
severity=AuditSeverity.ERROR,
|
||||
message=f"Container start failed: {e}",
|
||||
details={"container_id": container_id, "error": str(e)}
|
||||
)
|
||||
raise
|
||||
|
||||
def stop_container(
|
||||
self,
|
||||
container_id: str,
|
||||
timeout: int = 10
|
||||
) -> None:
|
||||
"""
|
||||
Stop a container.
|
||||
|
||||
Args:
|
||||
container_id: Container ID to stop
|
||||
timeout: Timeout in seconds
|
||||
"""
|
||||
# Verify container is registered
|
||||
if container_id not in self.validator.session_containers:
|
||||
self.audit_logger.log_security_violation(
|
||||
operation="container_stop",
|
||||
reason=f"Attempted to stop unregistered container: {container_id}"
|
||||
)
|
||||
raise SecurityError(
|
||||
f"Container {container_id} is not a registered session container"
|
||||
)
|
||||
|
||||
try:
|
||||
container = self.podman.client.containers.get(container_id)
|
||||
container.stop(timeout=timeout)
|
||||
|
||||
self.audit_logger.log_container_operation(
|
||||
operation="stop",
|
||||
container_id=container_id,
|
||||
image="",
|
||||
details={"timeout": timeout}
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.audit_logger.log(
|
||||
event_type=AuditEventType.CONTAINER_STOP,
|
||||
severity=AuditSeverity.ERROR,
|
||||
message=f"Container stop failed: {e}",
|
||||
details={"container_id": container_id, "error": str(e)}
|
||||
)
|
||||
raise
|
||||
|
||||
def remove_container(
|
||||
self,
|
||||
container_id: str,
|
||||
force: bool = False
|
||||
) -> None:
|
||||
"""
|
||||
Remove a container.
|
||||
|
||||
Args:
|
||||
container_id: Container ID to remove
|
||||
force: Force removal even if running
|
||||
"""
|
||||
# Verify container is registered
|
||||
if container_id not in self.validator.session_containers:
|
||||
self.audit_logger.log_security_violation(
|
||||
operation="container_remove",
|
||||
reason=f"Attempted to remove unregistered container: {container_id}"
|
||||
)
|
||||
raise SecurityError(
|
||||
f"Container {container_id} is not a registered session container"
|
||||
)
|
||||
|
||||
try:
|
||||
container = self.podman.client.containers.get(container_id)
|
||||
container.remove(force=force)
|
||||
|
||||
# Unregister from validator
|
||||
self.validator.unregister_session_container(container_id)
|
||||
|
||||
self.audit_logger.log_container_operation(
|
||||
operation="remove",
|
||||
container_id=container_id,
|
||||
image="",
|
||||
details={"force": force}
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.audit_logger.log(
|
||||
event_type=AuditEventType.CONTAINER_REMOVE,
|
||||
severity=AuditSeverity.ERROR,
|
||||
message=f"Container removal failed: {e}",
|
||||
details={"container_id": container_id, "error": str(e)}
|
||||
)
|
||||
raise
|
||||
|
||||
def get_container_logs(
|
||||
self,
|
||||
container_id: str,
|
||||
tail: int = 100
|
||||
) -> tuple[str, str]:
|
||||
"""
|
||||
Get container stdout and stderr logs.
|
||||
|
||||
Args:
|
||||
container_id: Container ID
|
||||
tail: Number of lines to retrieve
|
||||
|
||||
Returns:
|
||||
(stdout, stderr) as strings
|
||||
"""
|
||||
if container_id not in self.validator.session_containers:
|
||||
raise SecurityError(
|
||||
f"Container {container_id} is not a registered session container"
|
||||
)
|
||||
|
||||
try:
|
||||
container = self.podman.client.containers.get(container_id)
|
||||
logs = container.logs(tail=tail, stdout=True, stderr=True)
|
||||
|
||||
# Podman logs returns a generator of frames, need to consume it
|
||||
if hasattr(logs, '__iter__') and not isinstance(logs, (str, bytes)):
|
||||
# It's a generator/iterator, consume it
|
||||
logs_bytes = b''.join(logs)
|
||||
logs_str = logs_bytes.decode('utf-8', errors='replace')
|
||||
elif isinstance(logs, bytes):
|
||||
logs_str = logs.decode('utf-8', errors='replace')
|
||||
else:
|
||||
logs_str = str(logs)
|
||||
|
||||
# For simplicity, return all logs in stdout (Podman combines them)
|
||||
return logs_str, ""
|
||||
|
||||
except Exception as e:
|
||||
self.audit_logger.log(
|
||||
event_type=AuditEventType.EXECUTION_REQUEST,
|
||||
severity=AuditSeverity.ERROR,
|
||||
message=f"Failed to get container logs: {e}",
|
||||
details={"container_id": container_id, "error": str(e)}
|
||||
)
|
||||
raise
|
||||
|
||||
def wait_for_container(
|
||||
self,
|
||||
container_id: str,
|
||||
timeout: int = 300
|
||||
) -> int:
|
||||
"""
|
||||
Wait for container to exit.
|
||||
|
||||
Args:
|
||||
container_id: Container ID
|
||||
timeout: Timeout in seconds
|
||||
|
||||
Returns:
|
||||
Exit code
|
||||
|
||||
Raises:
|
||||
TimeoutError: If container doesn't exit within timeout
|
||||
"""
|
||||
if container_id not in self.validator.session_containers:
|
||||
raise SecurityError(
|
||||
f"Container {container_id} is not a registered session container"
|
||||
)
|
||||
|
||||
try:
|
||||
container = self.podman.client.containers.get(container_id)
|
||||
result = container.wait(timeout=timeout)
|
||||
|
||||
# Extract exit code from result
|
||||
if isinstance(result, dict):
|
||||
exit_code = result.get("StatusCode", 0)
|
||||
else:
|
||||
exit_code = result
|
||||
|
||||
return exit_code
|
||||
|
||||
except Exception as e:
|
||||
self.audit_logger.log(
|
||||
event_type=AuditEventType.EXECUTION_REQUEST,
|
||||
severity=AuditSeverity.ERROR,
|
||||
message=f"Failed to wait for container: {e}",
|
||||
details={"container_id": container_id, "error": str(e)}
|
||||
)
|
||||
raise
|
||||
|
||||
def cleanup_old_containers(
|
||||
self,
|
||||
max_age: timedelta = timedelta(hours=24)
|
||||
) -> int:
|
||||
"""
|
||||
Cleanup containers older than max_age.
|
||||
|
||||
Args:
|
||||
max_age: Maximum age for containers
|
||||
|
||||
Returns:
|
||||
Number of containers removed
|
||||
"""
|
||||
try:
|
||||
# Get all containers with mcp-forge.session label
|
||||
containers = self.podman.client.containers.list(
|
||||
all=True,
|
||||
filters={"label": ["mcp-forge.session"]}
|
||||
)
|
||||
|
||||
removed_count = 0
|
||||
now = datetime.now()
|
||||
|
||||
for container in containers:
|
||||
# Get creation time
|
||||
created_str = container.attrs.get("Created", "")
|
||||
if not created_str:
|
||||
continue
|
||||
|
||||
# Parse ISO format timestamp
|
||||
try:
|
||||
# Remove fractional seconds and timezone for parsing
|
||||
created_str = created_str.split('.')[0]
|
||||
created = datetime.fromisoformat(created_str.replace('Z', ''))
|
||||
except (ValueError, AttributeError):
|
||||
continue
|
||||
|
||||
age = now - created
|
||||
|
||||
if age > max_age:
|
||||
try:
|
||||
container.remove(force=True)
|
||||
removed_count += 1
|
||||
|
||||
self.audit_logger.log(
|
||||
event_type=AuditEventType.CONTAINER_REMOVE,
|
||||
severity=AuditSeverity.INFO,
|
||||
message=f"Cleaned up old container: {container.id}",
|
||||
details={
|
||||
"container_id": container.id,
|
||||
"age_hours": age.total_seconds() / 3600
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
self.audit_logger.log(
|
||||
event_type=AuditEventType.CONTAINER_REMOVE,
|
||||
severity=AuditSeverity.WARNING,
|
||||
message=f"Failed to remove old container: {e}",
|
||||
details={"container_id": container.id, "error": str(e)}
|
||||
)
|
||||
|
||||
return removed_count
|
||||
|
||||
except Exception as e:
|
||||
self.audit_logger.log(
|
||||
event_type=AuditEventType.CONTAINER_REMOVE,
|
||||
severity=AuditSeverity.ERROR,
|
||||
message=f"Cleanup failed: {e}",
|
||||
details={"error": str(e)}
|
||||
)
|
||||
raise
|
||||
|
|
@ -406,7 +406,7 @@ print(f"Found {len(records)} records")
|
|||
# Create resource handler
|
||||
self.resource_handler = ResourceHandler(
|
||||
client_manager=self.client_manager,
|
||||
jupyter_backend=self.jupyter_backend,
|
||||
session_manager=self.jupyter_backend,
|
||||
environment_builder=self.environment_builder,
|
||||
config=self.config
|
||||
)
|
||||
|
|
|
|||
|
|
@ -97,8 +97,14 @@ class ContainerConfig:
|
|||
params["volumes"] = self.volumes
|
||||
|
||||
# Add resource limits if present
|
||||
# Skip resource limits if using very high values (indicates no enforcement)
|
||||
if self.resource_limits:
|
||||
limit_params = self.resource_limits.to_podman_params()
|
||||
# Only apply limits if they're reasonable (not "no enforcement" markers)
|
||||
# Check if mem_limit looks like an enforcement bypass (>= 16GB)
|
||||
mem_limit = limit_params.get("mem_limit", "0")
|
||||
mem_bytes = int(mem_limit) if mem_limit != "0" else 0
|
||||
if mem_bytes < 16 * 1024 * 1024 * 1024: # Less than 16GB = real limit
|
||||
params.update(limit_params)
|
||||
# Disable swap to avoid cgroup swap.max issues on some systems
|
||||
if "mem_limit" in params:
|
||||
|
|
|
|||
|
|
@ -52,14 +52,15 @@ class JupyterBackend:
|
|||
# Initialize kernel manager
|
||||
kernel_manager = JupyterKernelManager(
|
||||
container_manager=container_manager,
|
||||
image=config.images.jupyter,
|
||||
image=image,
|
||||
resource_limits=self._default_resource_limits()
|
||||
)
|
||||
|
||||
# Initialize session manager
|
||||
self.session_manager = SessionManager(
|
||||
config=config.sessions,
|
||||
kernel_manager=kernel_manager,
|
||||
idle_timeout=idle_timeout,
|
||||
max_sessions=max_sessions,
|
||||
audit_logger=audit_logger
|
||||
)
|
||||
|
||||
|
|
@ -109,6 +110,7 @@ class JupyterBackend:
|
|||
|
||||
# Log execution (hash code, don't log actual content)
|
||||
code_hash = hashlib.sha256(code.encode()).hexdigest()
|
||||
if self.audit_logger:
|
||||
self.audit_logger.log(
|
||||
event_type="execution.request",
|
||||
severity="info",
|
||||
|
|
@ -236,9 +238,7 @@ class JupyterBackend:
|
|||
Returns:
|
||||
ResourceLimits with config defaults, or None if enforcement disabled
|
||||
"""
|
||||
if not self.config.security.enforce_resource_limits:
|
||||
return None
|
||||
|
||||
# Always return resource limits in pod_executor
|
||||
return ResourceLimits(
|
||||
memory=self.default_memory,
|
||||
cpu_quota=self.default_cpu_quota,
|
||||
|
|
|
|||
|
|
@ -130,3 +130,46 @@ class SimpleFileAuditLogger:
|
|||
# Write to file (append mode, file locking via 'a' mode)
|
||||
with open(self.log_path, 'a') as f:
|
||||
f.write(json.dumps(entry) + '\n')
|
||||
|
||||
def log_container_operation(
|
||||
self,
|
||||
operation: str,
|
||||
container_id: str,
|
||||
image: str,
|
||||
session_id: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
details: Optional[Dict[str, Any]] = None,
|
||||
error: Optional[str] = None
|
||||
) -> None:
|
||||
"""
|
||||
Convenience method to log container operations.
|
||||
|
||||
Args:
|
||||
operation: Operation type (create, start, stop, remove)
|
||||
container_id: Container ID
|
||||
image: Container image name
|
||||
session_id: Optional session ID
|
||||
user_id: Optional user ID
|
||||
details: Optional additional details
|
||||
error: Optional error message
|
||||
"""
|
||||
severity = "error" if error else "info"
|
||||
message = f"Container {operation}: {container_id[:12]} (image: {image})"
|
||||
|
||||
op_details = {
|
||||
"operation": operation,
|
||||
"container_id": container_id,
|
||||
"image": image
|
||||
}
|
||||
if details:
|
||||
op_details.update(details)
|
||||
|
||||
self.log(
|
||||
event_type=f"container.{operation}",
|
||||
severity=severity,
|
||||
message=message,
|
||||
session_id=session_id,
|
||||
user_id=user_id,
|
||||
details=op_details,
|
||||
error=error
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue