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 2026-03-04 22:54:56 +01:00
parent 9ceeaa1eda
commit 7d9efc5a38
10 changed files with 3384 additions and 26 deletions

1448
docs/architecture2.md Normal file

File diff suppressed because it is too large Load diff

View file

@ -42,28 +42,40 @@ class SimpleBackend:
self.config = config
# Create default resource limits from config
self.default_limits = ResourceLimits(
memory=config.execution.default_memory,
storage="10g", # Default storage limit
cpu_quota=config.execution.default_cpu_quota,
timeout=config.execution.default_timeout
)
# 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",
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)
# Use default executor
return self.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=code,
timeout=timeout,
injection_code=injection_code,
bridge_socket_path=bridge_socket_path
)

View 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}"
)

View 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

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

View 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

View file

@ -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
)

View file

@ -97,12 +97,18 @@ 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()
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
# 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:
params["memswap_limit"] = -1 # Disable swap
return params

View file

@ -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,7 +110,8 @@ class JupyterBackend:
# Log execution (hash code, don't log actual content)
code_hash = hashlib.sha256(code.encode()).hexdigest()
self.audit_logger.log(
if self.audit_logger:
self.audit_logger.log(
event_type="execution.request",
severity="info",
message="Stateful code execution requested",
@ -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,

View file

@ -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
)