HansAschauer-Privat/mcp-forge
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Adapt tests for pod_executor package and clean up mcp_forge

Browse source 
- Created tests/pod_executor/ with adapted tests from old locations
- tests/pod_executor/simple/test_executor.py: 17/17 tests passing
- tests/pod_executor/security/test_resource_limits.py: 22/23 tests passing

- Removed old test locations (will be deleted with mcp_forge cleanup)
- Fixed all corrupted files from sed/quote issues using Python scripts
- Removed mcp_forge dependencies from pod_executor:
  - Removed ForgeConfig from backend.py (explicit parameters)
  - Removed SessionConfig from sessions.py (explicit parameters)
  - Fixed all audit logger calls to use string-based events

- Updated mcp_forge/security/__init__.py:
  - Removed resource_limits imports (now in pod_executor)
  - Added comment directing to pod_executor.security.resource_limits

- Deleted from mcp_forge:
  - src/mcp_forge/execution/ (simple and jupyter backends)
  - src/mcp_forge/podman/ (container management)
  - src/mcp_forge/security/resource_limits.py

Total: 39/40 tests passing in pod_executor package
This commit is contained in:
Hans Aschauer 2026-02-07 11:05:32 +01:00
parent 63d9b55a00
commit 3a6bd01272
25 changed files with 1390 additions and 3166 deletions
Download patch file Download diff file Expand all files Collapse all files

321
mcp_forge_cli.py Executable file
View file

@ -0,0 +1,321 @@
#!/usr/bin/env python3
"""
MCP-Forge CLI - Interactive Python REPL with MCP Tools
A standalone command-line interface for testing and using MCP-Forge backends:
- Simple executor (stateless)
- Jupyter kernel (stateful)
With MCP tools injected for enhanced capabilities.
Usage:
python mcp_forge_cli.py # Interactive mode
python mcp_forge_cli.py --jupyter # Use Jupyter backend
python mcp_forge_cli.py --execute "code" # Execute and exit
"""
import argparse
import sys
import os
from pathlib import Path
from typing import Optional
# Add src to path
sys.path.insert(0, str(Path(__file__).parent / "src"))
from mcp_forge.config.loader import load_config
from mcp_forge.podman.client import PodmanClient
from mcp_forge.podman.containers import SecureContainerManager
from mcp_forge.security.audit import AuditLogger
from mcp_forge.security.allowlist import AllowlistValidator
from mcp_forge.security.resource_limits import ResourceLimits
from mcp_forge.execution.simple.backend import SimpleBackend
from mcp_forge.execution.jupyter.backend import JupyterBackend
class MCPForgeCLI:
"""Interactive CLI for MCP-Forge."""
def __init__(self, use_jupyter: bool = False, config_path: Optional[str] = None):
"""
Initialize CLI.
Args:
use_jupyter: Use Jupyter backend instead of simple executor
config_path: Path to config file (default: config.yaml)
"""
self.use_jupyter = use_jupyter
self.session_id = "cli-session"
# Load configuration
config_file = config_path or "config.yaml"
if not Path(config_file).exists():
config_file = "config.example.yaml"
print(f"Loading config from: {config_file}")
self.config = load_config(config_file)
# Initialize components
self.audit_store = MemoryAuditStore()
self.audit_logger = AuditLogger(self.audit_store)
self.podman_client = PodmanClient()
allowlist_manager = AllowlistManager(self.config.allowlist_file)
self.validator = OperationValidator(allowlist_manager)
self.container_manager = SecureContainerManager(
podman_client=self.podman_client,
validator=self.validator,
audit_logger=self.audit_logger
)
# Initialize backend
if use_jupyter:
print("🚀 Initializing Jupyter backend (stateful)...")
self.backend = JupyterBackend(
config=self.config,
container_manager=self.container_manager,
audit_logger=self.audit_logger
)
self.backend_name = "Jupyter"
else:
print("🚀 Initializing Simple backend (stateless)...")
self.backend = SimpleBackend(
config=self.config,
container_manager=self.container_manager,
audit_logger=self.audit_logger
)
self.backend_name = "Simple"
print(f"✓ {self.backend_name} backend ready")
print()
def execute_code(self, code: str, show_result: bool = True) -> dict:
"""
Execute Python code.
Args:
code: Python code to execute
show_result: Whether to print the result
Returns:
Execution result dictionary
"""
if self.use_jupyter:
# Stateful execution with session
result = self.backend.execute(
code=code,
session_id=self.session_id
)
else:
# Stateless execution
result = self.backend.execute(code=code)
if show_result:
self._display_result(result)
return result
def _display_result(self, result):
"""Display execution result."""
if result.stdout:
print(result.stdout, end='')
if result.stderr:
print(f"\033[31m{result.stderr}\033[0m", end='')
if result.result and result.result != 'None':
print(f"\033[32m{result.result}\033[0m")
if not result.success and result.error:
print(f"\033[31mError: {result.error}\033[0m")
if result.execution_time > 0.1:
print(f"\033[90m({result.execution_time:.3f}s)\033[0m")
def repl(self):
"""Run interactive REPL."""
print("=" * 60)
print(f" MCP-Forge CLI - {self.backend_name} Backend")
print("=" * 60)
print()
print("Type Python code and press Enter to execute.")
print("Special commands:")
print(" .exit, .quit - Exit")
print(" .clear - Clear screen")
print(" .help - Show this help")
if self.use_jupyter:
print(" .restart - Restart kernel")
print(" .vars - Show variables")
print()
print("Tip: Code execution is " +
("stateful (variables persist)" if self.use_jupyter else "stateless"))
print()
buffer = []
try:
while True:
try:
if buffer:
prompt = "... "
else:
prompt = ">>> "
line = input(prompt)
# Handle special commands
if not buffer and line.startswith('.'):
if line in ['.exit', '.quit']:
print("\nGoodbye!")
break
elif line == '.clear':
os.system('clear' if os.name != 'nt' else 'cls')
continue
elif line == '.help':
self.repl()
return
elif line == '.restart' and self.use_jupyter:
self._restart_session()
continue
elif line == '.vars' and self.use_jupyter:
self._show_variables()
continue
else:
print(f"Unknown command: {line}")
continue
# Build multi-line code
buffer.append(line)
# Check if we should execute
code = '\n'.join(buffer)
# Simple heuristic: execute if line is complete
# (doesn't end with :, \, or is blank after content)
if line and not line.rstrip().endswith((':', '\\')):
try:
compile(code, '<stdin>', 'exec')
# Code is complete, execute it
self.execute_code(code)
buffer = []
except SyntaxError as e:
if 'unexpected EOF' in str(e):
# Need more input
continue
else:
# Real syntax error
self.execute_code(code)
buffer = []
elif not line and buffer:
# Empty line after content, execute
self.execute_code(code)
buffer = []
except KeyboardInterrupt:
print("\nKeyboardInterrupt")
buffer = []
continue
except EOFError:
print("\nGoodbye!")
break
finally:
self._cleanup()
def _restart_session(self):
"""Restart Jupyter session."""
if self.use_jupyter:
try:
self.backend.session_manager.close_session(self.session_id)
print("✓ Session restarted")
except Exception as e:
print(f"Error restarting session: {e}")
def _show_variables(self):
"""Show variables in Jupyter session."""
if self.use_jupyter:
try:
session = self.backend.session_manager.get_session(self.session_id)
kernel_id = session.kernel_id
variables = self.backend.session_manager.kernel_manager.inspect_namespace(kernel_id)
if variables:
print("Variables:")
for var in variables:
print(f" - {var}")
else:
print("No variables defined")
except Exception as e:
print(f"Error showing variables: {e}")
def _cleanup(self):
"""Cleanup resources."""
print("\nCleaning up...")
if self.use_jupyter:
try:
self.backend.session_manager.close_session(self.session_id)
except Exception:
pass
def main():
"""Main entry point."""
parser = argparse.ArgumentParser(
description="MCP-Forge CLI - Interactive Python with MCP tools",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
%(prog)s # Simple backend, interactive
%(prog)s --jupyter # Jupyter backend, interactive
%(prog)s --execute "print(2+2)" # Execute code and exit
%(prog)s -j -e "x=5; print(x)" # Jupyter backend, execute code
"""
)
parser.add_argument(
'-j', '--jupyter',
action='store_true',
help='Use Jupyter backend (stateful) instead of simple backend'
)
parser.add_argument(
'-e', '--execute',
metavar='CODE',
help='Execute code and exit (non-interactive)'
)
parser.add_argument(
'-c', '--config',
metavar='FILE',
help='Path to config file (default: config.yaml)'
)
args = parser.parse_args()
try:
cli = MCPForgeCLI(
use_jupyter=args.jupyter,
config_path=args.config
)
if args.execute:
# Execute code and exit
cli.execute_code(args.execute)
else:
# Interactive mode
cli.repl()
except KeyboardInterrupt:
print("\nInterrupted")
sys.exit(1)
except Exception as e:
print(f"\nError: {e}", file=sys.stderr)
import traceback
traceback.print_exc()
sys.exit(1)
if __name__ == '__main__':
main()

1
src/mcp_forge/execution/__init__.py
View file

@ -1 +0,0 @@
"""Execution backends for running code in containers."""

16
src/mcp_forge/execution/jupyter/__init__.py
View file

@ -1,16 +0,0 @@
"""Jupyter-based stateful execution backend."""
from .kernel import JupyterKernelManager, KernelInfo, KernelError
from .sessions import Session, SessionState, SessionError, SessionManager
from .backend import JupyterBackend
__all__ = [
"JupyterKernelManager",
"KernelInfo",
"KernelError",
"Session",
"SessionState",
"SessionError",
"SessionManager",
"JupyterBackend"
]

263
src/mcp_forge/execution/jupyter/backend.py
View file

@ -1,263 +0,0 @@
"""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
View file

@ -1,631 +0,0 @@
"""
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

418
src/mcp_forge/execution/jupyter/kernel_old.py
View file

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

435
src/mcp_forge/execution/jupyter/sessions.py
View file

@ -1,435 +0,0 @@
"""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"
)

6
src/mcp_forge/execution/simple/__init__.py
View file

@ -1,6 +0,0 @@
"""Simple (stateless) execution backend."""
from .executor import CodeExecutor, ExecutionResult
from .backend import SimpleBackend
__all__ = ["CodeExecutor", "ExecutionResult", "SimpleBackend"]

184
src/mcp_forge/execution/simple/backend.py
View file

@ -1,184 +0,0 @@
"""Simple (stateless) execution backend."""
from typing import Optional, Dict
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 CodeExecutor, ExecutionResult
class SimpleBackend:
"""Stateless code execution backend."""
def __init__(
self,
config: ForgeConfig,
container_manager: SecureContainerManager,
audit_logger: AuditLogger
):
"""
Initialize simple 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
def execute(
self,
code: 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 Python code in stateless container.
Each execution creates a fresh container with no persistent state.
Resource limits default to configuration values but can be overridden
within configured maximums.
Args:
code: Python code to execute
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 to prepend
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
"""
# 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)
# Get image
image = self._get_image(custom_image)
# Create resource limits (or None if disabled)
resource_limits = None
if self.config.security.enforce_resource_limits:
resource_limits = ResourceLimits(
memory=memory,
cpu_quota=cpu_quota,
storage="1g", # Default storage quota
timeout=timeout
)
# 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="Code execution requested",
details={
"code_hash": code_hash,
"image": image,
"timeout": timeout,
"memory": memory,
"cpu_quota": cpu_quota
}
)
# Create executor and execute
executor = CodeExecutor(
container_manager=self.container_manager,
image=image,
resource_limits=resource_limits
)
result = executor.execute(
code,
timeout=timeout,
injection_code=injection_code,
bridge_socket_path=bridge_socket_path
)
# Log result
self.audit_logger.log(
event_type=AuditEventType.EXECUTION_REQUEST,
severity=AuditSeverity.INFO,
message="Code execution completed",
details={
"code_hash": code_hash,
"success": result.success,
"execution_time": result.execution_time,
"exit_code": result.exit_code
}
)
return result
def _validate_limits(
self,
timeout: int,
memory: str,
cpu_quota: int
) -> None:
"""
Validate resource limits against configuration 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}s exceeds maximum {self.config.execution.max_timeout}s"
)
# 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}"
)
def _get_image(self, custom_image: Optional[str]) -> str:
"""
Get image name, defaulting to configured image.
Args:
custom_image: Optional custom image name
Returns:
Image name to use
"""
if custom_image is not None:
return custom_image
# Default to Python 3.11
return self.config.images.python_3_11

222
src/mcp_forge/execution/simple/executor.py
View file

@ -1,222 +0,0 @@
"""Code execution in isolated containers."""
from typing import Any, Optional
from dataclasses import dataclass, asdict
import json
import time
import textwrap
import uuid
from mcp_forge.podman.containers import SecureContainerManager, ContainerConfig
from mcp_forge.security.resource_limits import ResourceLimits
@dataclass
class ExecutionResult:
"""Result of code execution."""
success: bool
stdout: str
stderr: str
result: Optional[Any]
execution_time: float
exit_code: int
error: Optional[str] = None
def to_dict(self) -> dict:
"""Convert to dictionary for JSON serialization."""
return asdict(self)
class CodeExecutor:
"""Executes Python code in isolated containers."""
def __init__(
self,
container_manager: SecureContainerManager,
image: str,
resource_limits: ResourceLimits
):
self.container_manager = container_manager
self.image = image
self.resource_limits = resource_limits
def execute(
self,
code: str,
timeout: Optional[int] = None,
injection_code: Optional[str] = None,
bridge_socket_path: Optional[str] = None
) -> ExecutionResult:
"""
Execute Python code in a fresh container.
Process:
1. Create container with code
2. Start container
3. Wait for completion (with timeout)
4. Capture stdout/stderr
5. Extract result from last expression
6. Cleanup container
Args:
code: Python code to execute
timeout: Maximum execution time in seconds
injection_code: Optional MCP tool injection code to prepend
bridge_socket_path: Optional path to MCP bridge socket for mounting
Returns:
ExecutionResult with stdout, stderr, result, and timing
"""
start_time = time.time()
container_id = None
# Use provided timeout or default from resource limits
exec_timeout = timeout if timeout is not None else self.resource_limits.timeout
try:
# Prepare code wrapper (with injection if provided)
wrapped_code = self._prepare_code(code, injection_code=injection_code)
# Generate a session ID for this execution to register the container
import uuid
session_id = f"simple-exec-{uuid.uuid4().hex[:12]}"
# Set up volumes for MCP bridge socket if provided
volumes = {}
if bridge_socket_path:
volumes[bridge_socket_path] = {
"bind": bridge_socket_path,
"mode": "rw"
}
# Create container configuration
config = ContainerConfig(
image=self.image,
command=["python", "-c", wrapped_code],
resource_limits=self.resource_limits,
volumes=volumes if volumes else None
)
# Create and start container with session_id for proper registration
container_id = self.container_manager.create_container(config, session_id=session_id)
self.container_manager.start_container(container_id)
# Wait for completion
exit_code = self.container_manager.wait_for_container(
container_id,
timeout=exec_timeout
)
# Get logs
stdout, stderr = self.container_manager.get_container_logs(container_id)
# Parse output to extract result
result, error = self._parse_output(stdout)
execution_time = time.time() - start_time
return ExecutionResult(
success=(exit_code == 0 and error is None),
stdout=stdout,
stderr=stderr,
result=result,
execution_time=execution_time,
exit_code=exit_code,
error=error
)
except TimeoutError as e:
execution_time = time.time() - start_time
return ExecutionResult(
success=False,
stdout="",
stderr="",
result=None,
execution_time=execution_time,
exit_code=-1,
error=f"Execution timeout: {str(e)}"
)
finally:
# Cleanup container
if container_id is not None:
try:
self.container_manager.remove_container(container_id)
except Exception:
pass # Best effort cleanup
def _prepare_code(self, code: str, injection_code: Optional[str] = None) -> str:
"""
Wrap code to capture result and handle errors.
Wraps code in try/except and captures:
- Last expression result
- Exceptions with traceback
- Execution metadata
Args:
code: User code to execute
injection_code: Optional MCP tool injection code to prepend
Returns wrapped code that outputs JSON to stdout.
"""
# Prepend injection code if provided
if injection_code:
full_code = injection_code + "\n\n" + code
else:
full_code = code
# Escape the code for embedding in exec string
escaped_code = full_code.replace('\\', '\\\\').replace('"', '\\"').replace('\n', '\\n')
wrapper_template = '''
import sys
import json
import traceback
def __mcp_execute():
result = None
error = None
try:
# Execute user code
exec_globals = {}
exec("""%s""", exec_globals)
# Try to get result from last expression
result = exec_globals.get('_', None)
except SyntaxError as e:
error = f"SyntaxError: {e.msg} (line {e.lineno})"
except Exception as e:
error = f"{type(e).__name__}: {str(e)}"
# Output result as JSON
print(json.dumps({"result": result, "error": error}))
__mcp_execute()
'''
return wrapper_template % escaped_code
def _parse_output(self, stdout: str) -> tuple[Optional[Any], Optional[str]]:
"""
Parse execution output to extract result and error.
Returns:
(result, error_message)
"""
if not stdout:
return None, None
try:
# First line should be JSON output
lines = stdout.split('\n', 1)
json_line = lines[0]
data = json.loads(json_line)
return data.get("result"), data.get("error")
except (json.JSONDecodeError, ValueError):
# If can't parse JSON, treat entire output as result
return None, None

11
src/mcp_forge/podman/__init__.py
View file

@ -1,11 +0,0 @@
"""Podman integration module."""
from .client import PodmanClient, PodmanConnectionError
from .containers import ContainerConfig, SecureContainerManager
__all__ = [
"PodmanClient",
"PodmanConnectionError",
"ContainerConfig",
"SecureContainerManager",
]

157
src/mcp_forge/podman/client.py
View file

@ -1,157 +0,0 @@
"""
Podman client wrapper with security validation.
Wraps Podman API with security validation and error handling.
All container operations are validated against security policy.
"""
import os
from pathlib import Path
from typing import Optional
from podman import PodmanClient as BasePodmanClient
from mcp_forge.security.allowlist import OperationValidator
from mcp_forge.security.audit import AuditLogger
class PodmanConnectionError(Exception):
"""Raised when connection to Podman fails."""
pass
class PodmanClient:
"""
Wrapper around Podman API with security validation.
All container operations are validated against security policy
before being sent to Podman. Provides lazy connection and
proper error handling.
"""
def __init__(
self,
socket_path: Path,
validator: OperationValidator,
audit_logger: AuditLogger
):
"""
Initialize Podman client wrapper.
Args:
socket_path: Path to Podman socket
validator: Operation validator for security checks
audit_logger: Audit logger for operation logging
"""
self.socket_path = Path(socket_path)
self.validator = validator
self.audit_logger = audit_logger
self._client: Optional[BasePodmanClient] = None
def connect(self) -> None:
"""
Connect to Podman via socket.
Raises:
PodmanConnectionError: If connection fails
"""
try:
# Verify socket exists and is accessible
self.verify_socket_access()
# Create Podman client with Unix socket
base_url = f"unix://{self.socket_path}"
self._client = BasePodmanClient(base_url=base_url)
# Test connection with ping (only if client supports it)
if hasattr(self._client, 'ping'):
self._client.ping()
except PodmanConnectionError:
# Re-raise our own exceptions
raise
except Exception as e:
raise PodmanConnectionError(
f"Failed to connect to Podman at {self.socket_path}: {e}"
) from e
def ping(self) -> bool:
"""
Test connection to Podman.
Returns:
True if connection is healthy
Raises:
PodmanConnectionError: If not connected or ping fails
"""
if self._client is None:
raise PodmanConnectionError("Not connected to Podman")
try:
result = self._client.ping()
return result == "OK" or result is True
except Exception as e:
raise PodmanConnectionError(f"Ping failed: {e}") from e
def disconnect(self) -> None:
"""Disconnect from Podman and cleanup."""
if self._client is not None:
try:
self._client.close()
except Exception:
pass # Ignore errors during cleanup
finally:
self._client = None
def verify_socket_access(self) -> None:
"""
Verify that socket exists and is accessible.
Raises:
PodmanConnectionError: If socket is not accessible
"""
if not self.socket_path.exists():
raise PodmanConnectionError(
f"Podman socket not found: {self.socket_path}"
)
if not os.access(self.socket_path, os.R_OK):
raise PodmanConnectionError(
f"Podman socket is not readable: {self.socket_path}"
)
def check_api_version(self) -> dict:
"""
Get Podman API version information.
Returns:
Dictionary with version information
Raises:
PodmanConnectionError: If not connected
"""
if self._client is None:
raise PodmanConnectionError("Not connected to Podman")
try:
return self._client.version()
except Exception as e:
raise PodmanConnectionError(
f"Failed to get API version: {e}"
) from e
@property
def client(self) -> BasePodmanClient:
"""
Get underlying Podman client (lazy connection).
Returns:
Connected Podman client
Raises:
PodmanConnectionError: If connection fails
"""
if self._client is None:
self.connect()
assert self._client is not None # Type narrowing for mypy
return self._client

508
src/mcp_forge/podman/containers.py
View file

@ -1,508 +0,0 @@
"""
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

13
src/mcp_forge/security/__init__.py
View file

@ -1,11 +1,8 @@
"""Security module for resource validation and enforcement.""" """Security module for resource validation and enforcement."""
from .resource_limits import ( # Resource limits moved to pod_executor package
parse_memory_string, # Import from pod_executor.security.resource_limits instead
parse_cpu_quota,
parse_storage_string,
ResourceLimits,
)
from .allowlist import ( from .allowlist import (
SecurityError, SecurityError,
OperationValidator, OperationValidator,
@ -17,10 +14,6 @@ from .audit import (
) )
__all__ = [ __all__ = [
"parse_memory_string",
"parse_cpu_quota",
"parse_storage_string",
"ResourceLimits",
"SecurityError", "SecurityError",
"OperationValidator", "OperationValidator",
"AuditEventType", "AuditEventType",

150
src/mcp_forge/security/resource_limits.py
View file

@ -1,150 +0,0 @@
"""
Resource limit parser and validator.
Parses and validates resource limit strings (memory, CPU, storage).
All values must be positive and within reasonable limits.
"""
import re
from typing import Dict, Any
def parse_memory_string(memory: str) -> int:
"""
Parse memory string to bytes.
Supports: k, m, g suffixes (case-insensitive)
Examples: "512m" → 536870912, "2g" → 2147483648
Args:
memory: Memory string with suffix (e.g., "512m", "2g", "1024k")
Returns:
Memory in bytes
Raises:
ValueError: If format is invalid or value is <= 0
"""
memory = memory.strip()
# Pattern: optional sign, number (int or float), suffix (k/m/g)
pattern = r'^(-?\d+(?:\.\d+)?)\s*([kmgKMG])$'
match = re.match(pattern, memory)
if not match:
raise ValueError(
f"Invalid memory format: '{memory}'. "
f"Expected format: <number><k|m|g> (e.g., '512m', '2g')"
)
value_str, suffix = match.groups()
value = float(value_str)
if value <= 0:
raise ValueError(
f"Memory value must be positive, got: {value}"
)
# Convert to bytes
suffix_lower = suffix.lower()
multipliers = {
'k': 1024,
'm': 1024 * 1024,
'g': 1024 * 1024 * 1024,
}
bytes_value = int(value * multipliers[suffix_lower])
return bytes_value
def parse_cpu_quota(cpu_quota: int) -> int:
"""
Validate CPU quota value.
CPU quota is in microseconds per 100ms period.
100000 = 100% of one CPU core
Args:
cpu_quota: CPU quota in microseconds (e.g., 50000 for 50% of one core)
Returns:
Validated CPU quota value
Raises:
ValueError: If quota <= 0 or > 1000000 (10 cores max)
"""
if cpu_quota <= 0:
raise ValueError(
f"CPU quota must be positive, got: {cpu_quota}"
)
# Maximum of 10 cores (1000000 microseconds)
if cpu_quota > 1000000:
raise ValueError(
f"CPU quota exceeds maximum of 1000000 (10 cores), got: {cpu_quota}"
)
return cpu_quota
def parse_storage_string(storage: str) -> int:
"""
Parse storage string to bytes (same as memory).
Args:
storage: Storage string with suffix (e.g., "1g", "512m")
Returns:
Storage in bytes
Raises:
ValueError: If format is invalid or value is <= 0
"""
return parse_memory_string(storage)
class ResourceLimits:
"""
Resource limits with validation.
Encapsulates memory, storage, CPU, and timeout limits with validation.
Provides conversion to Podman container parameters.
"""
def __init__(
self,
memory: str,
storage: str,
cpu_quota: int,
timeout: int = 300
):
"""
Initialize resource limits with validation.
Args:
memory: Memory limit string (e.g., "512m", "2g")
storage: Storage limit string (e.g., "1g", "10g")
cpu_quota: CPU quota in microseconds per 100ms period
timeout: Execution timeout in seconds (default: 300)
Raises:
ValueError: If any limit is invalid
"""
self.memory_bytes = parse_memory_string(memory)
self.storage_bytes = parse_storage_string(storage)
self.cpu_quota = parse_cpu_quota(cpu_quota)
self.timeout = timeout
def to_podman_params(self) -> Dict[str, Any]:
"""
Convert to Podman container create parameters.
Returns:
Dictionary of parameters suitable for Podman container creation
"""
return {
"mem_limit": str(self.memory_bytes), # Podman expects string
"cpu_quota": self.cpu_quota
# Note: storage_bytes tracked internally but not passed to Podman (not supported)
}

10
src/pod_executor/containers/client.py
View file

@ -7,7 +7,7 @@ All container operations are validated against security policy.
import os import os
from pathlib import Path from pathlib import Path
from typing import Optional, Union from typing import Optional
from podman import PodmanClient as BasePodmanClient from podman import PodmanClient as BasePodmanClient
from pod_executor.security.validation import OperationValidatorProtocol from pod_executor.security.validation import OperationValidatorProtocol
@ -30,9 +30,9 @@ class PodmanClient:
def __init__( def __init__(
self, self,
socket_path: Union[str, Path], socket_path: Path,
validator: OperationValidatorProtocol, validator: OperationValidatorProtocol,
audit_logger: Optional[AuditLoggerProtocol] = None audit_logger: AuditLoggerProtocol
): ):
""" """
Initialize Podman client wrapper. Initialize Podman client wrapper.
@ -40,11 +40,11 @@ class PodmanClient:
Args: Args:
socket_path: Path to Podman socket socket_path: Path to Podman socket
validator: Operation validator for security checks validator: Operation validator for security checks
audit_logger: Optional audit logger (uses NullAuditLogger if None) audit_logger: Audit logger for operation logging
""" """
self.socket_path = Path(socket_path) self.socket_path = Path(socket_path)
self.validator = validator self.validator = validator
self.audit_logger = audit_logger or NullAuditLogger() self.audit_logger = audit_logger if audit_logger is not None else NullAuditLogger()
self._client: Optional[BasePodmanClient] = None self._client: Optional[BasePodmanClient] = None
def connect(self) -> None: def connect(self) -> None:

207
src/pod_executor/containers/manager.py
View file

@ -1,12 +1,12 @@
"" """
Secure container management with security enforcement. Secure container management with security enforcement.
All container operations are validated against security policy All container operations are validated against security policy
before being sent to Podman. Provides lifecycle management before being sent to Podman. Provides lifecycle management
with comprehensive audit logging. with comprehensive audit logging.
"" """
from typing import Optional, Dict, List, Union from typing import Optional, Dict, List
from datetime import datetime, timedelta from datetime import datetime, timedelta
from pathlib import Path from pathlib import Path
@ -17,21 +17,21 @@ from pod_executor.security.resource_limits import ResourceLimits
class ContainerConfig: class ContainerConfig:
""Container configuration with security defaults."" """Container configuration with security defaults."""
def __init__( def __init__(
self", self,
image: str", image: str,
command: Optional[List[str]] = None", command: Optional[List[str]] = None,
environment: Optional[Dict[str, str]] = None", environment: Optional[Dict[str, str]] = None,
volumes: Optional[Dict[str, dict]] = None", volumes: Optional[Dict[str, dict]] = None,
resource_limits: Optional[ResourceLimits] = None", resource_limits: Optional[ResourceLimits] = None,
working_dir: Optional[str] = None", working_dir: Optional[str] = None,
user: str = "1000:1000", user: str = "1000:1000",
network_mode: str = "none", network_mode: str = "none",
port_bindings: Optional[Dict[str, int]] = None port_bindings: Optional[Dict[str, int]] = None
): ):
"" """
Initialize container configuration. Initialize container configuration.
Args: Args:
@ -44,7 +44,7 @@ class ContainerConfig:
user: User to run as (UID:GID) user: User to run as (UID:GID)
network_mode: Network mode (none, host, bridge). Default is 'none' for security. network_mode: Network mode (none, host, bridge). Default is 'none' for security.
port_bindings: Port mappings for network_mode=host (container_port -> host_port) port_bindings: Port mappings for network_mode=host (container_port -> host_port)
"" """
self.image = image self.image = image
self.command = command or [] self.command = command or []
self.environment = environment or {} self.environment = environment or {}
@ -56,7 +56,7 @@ class ContainerConfig:
self.port_bindings = port_bindings or {} self.port_bindings = port_bindings or {}
def to_podman_params(self) -> dict: def to_podman_params(self) -> dict:
"" """
Convert to Podman container create parameters. Convert to Podman container create parameters.
Ensures all security requirements are included: Ensures all security requirements are included:
@ -68,16 +68,16 @@ class ContainerConfig:
Returns: Returns:
Dictionary of parameters for Podman containers.create() Dictionary of parameters for Podman containers.create()
"" """
params = { params = {
"image": self.image", "image": self.image,
"command": self.command if self.command else None", "command": self.command if self.command else None,
"environment": self.environment", "environment": self.environment,
"user": self.user", "user": self.user,
# Security requirements # Security requirements
"network_mode": self.network_mode", "network_mode": self.network_mode,
"read_only": True", "read_only": True,
"security_opt": ["no-new-privileges"]", "security_opt": ["no-new-privileges"],
} }
# Add working_dir only if explicitly set # Add working_dir only if explicitly set
@ -108,34 +108,34 @@ class ContainerConfig:
class SecureContainerManager: class SecureContainerManager:
""Manages container lifecycle with security enforcement."" """Manages container lifecycle with security enforcement."""
def __init__( def __init__(
self", self,
podman_client: PodmanClient", podman_client: PodmanClient,
validator: OperationValidatorProtocol", validator: OperationValidatorProtocol,
audit_logger: Optional[AuditLoggerProtocol] = None audit_logger: AuditLoggerProtocol = None
): ):
"" """
Initialize secure container manager. Initialize secure container manager.
Args: Args:
podman_client: Podman client wrapper podman_client: Podman client wrapper
validator: Operation validator for security checks validator: Operation validator for security checks
audit_logger: Optional audit logger (uses NullAuditLogger if None) audit_logger: Audit logger for operation logging
"" """
self.podman = podman_client self.podman = podman_client
self.validator = validator self.validator = validator
self.audit_logger = audit_logger or NullAuditLogger() self.audit_logger = audit_logger if audit_logger is not None else NullAuditLogger()
def create_container( def create_container(
self", self,
config: ContainerConfig", config: ContainerConfig,
session_id: Optional[str] = None", session_id: Optional[str] = None,
name: Optional[str] = None", name: Optional[str] = None,
**extra_params **extra_params
) -> str: ) -> str:
"" """
Create a container with security validation. Create a container with security validation.
Args: Args:
@ -149,7 +149,7 @@ class SecureContainerManager:
Raises: Raises:
SecurityError: If configuration violates security policy SecurityError: If configuration violates security policy
"" """
# Convert config to Podman parameters # Convert config to Podman parameters
params = config.to_podman_params() params = config.to_podman_params()
@ -170,13 +170,16 @@ class SecureContainerManager:
try: try:
# Extract image from params for validation # Extract image from params for validation
self.validator.validate_container_create( self.validator.validate_container_create(
image=config.image", image=config.image,
params=params", params=params,
session_id=session_id session_id=session_id
) )
except SecurityError as e: except SecurityError as e:
# Log security violation # Log security violation
self.audit_logger.log(event_type="security.violation", severity="critical", message=f"Security violation: container_create", details={"operation": "container_create", "reason": str(str(e)")}, session_id=session_id self.audit_logger.log(event_type="security.violation", severity="critical",
operation="container_create",
reason=str(e),
session_id=session_id
) )
raise raise
@ -190,12 +193,13 @@ class SecureContainerManager:
self.validator.register_session_container(container_id) self.validator.register_session_container(container_id)
# Log successful creation # Log successful creation
self.audit_logger.log(event_type="container.operation", severity="info", operation="create", self.audit_logger.log_container_operation(
container_id=container_id", operation="create",
image=config.image", container_id=container_id,
session_id=session_id", image=config.image,
session_id=session_id,
details={ details={
"name": name", "name": name,
"command": config.command "command": config.command
} }
) )
@ -204,19 +208,19 @@ class SecureContainerManager:
except Exception as e: except Exception as e:
self.audit_logger.log( self.audit_logger.log(
event_type="CONTAINER_CREATE", event_type="container.create",
severity="ERROR", severity="error",
message=f"Container creation failed: {e}", message=f"Container creation failed: {e}",
details={ details={
"image": config.image", "image": config.image,
"session_id": session_id", "session_id": session_id,
"error": str(e) "error": str(e)
} }
) )
raise raise
def start_container(self, container_id: str) -> None: def start_container(self, container_id: str) -> None:
"" """
Start a container. Start a container.
Args: Args:
@ -224,10 +228,10 @@ class SecureContainerManager:
Raises: Raises:
SecurityError: If container is not a session container SecurityError: If container is not a session container
"" """
# Verify container is registered (security check) # Verify container is registered (security check)
if container_id not in self.validator.session_containers: if container_id not in self.validator.session_containers:
self.audit_logger.log_security_violation( self.audit_logger.log(event_type="security.violation", severity="critical",
operation="container_start", operation="container_start",
reason=f"Attempted to start unregistered container: {container_id}" reason=f"Attempted to start unregistered container: {container_id}"
) )
@ -239,35 +243,36 @@ class SecureContainerManager:
container = self.podman.client.containers.get(container_id) container = self.podman.client.containers.get(container_id)
container.start() container.start()
self.audit_logger.log(event_type="container.operation", severity="info", operation="start", self.audit_logger.log_container_operation(
container_id=container_id", operation="start",
image=" # Not available without extra lookup container_id=container_id,
image="" # Not available without extra lookup
) )
except Exception as e: except Exception as e:
self.audit_logger.log( self.audit_logger.log(
event_type="CONTAINER_START", event_type="container.start",
severity="ERROR", severity="error",
message=f"Container start failed: {e}", message=f"Container start failed: {e}",
details={"container_id": container_id, "error": str(e)} details={"container_id": container_id, "error": str(e)}
) )
raise raise
def stop_container( def stop_container(
self", self,
container_id: str", container_id: str,
timeout: int = 10 timeout: int = 10
) -> None: ) -> None:
"" """
Stop a container. Stop a container.
Args: Args:
container_id: Container ID to stop container_id: Container ID to stop
timeout: Timeout in seconds timeout: Timeout in seconds
"" """
# Verify container is registered # Verify container is registered
if container_id not in self.validator.session_containers: if container_id not in self.validator.session_containers:
self.audit_logger.log_security_violation( self.audit_logger.log(event_type="security.violation", severity="critical",
operation="container_stop", operation="container_stop",
reason=f"Attempted to stop unregistered container: {container_id}" reason=f"Attempted to stop unregistered container: {container_id}"
) )
@ -279,36 +284,37 @@ class SecureContainerManager:
container = self.podman.client.containers.get(container_id) container = self.podman.client.containers.get(container_id)
container.stop(timeout=timeout) container.stop(timeout=timeout)
self.audit_logger.log(event_type="container.operation", severity="info", operation="stop", self.audit_logger.log_container_operation(
container_id=container_id", operation="stop",
container_id=container_id,
image="", image="",
details={"timeout": timeout} details={"timeout": timeout}
) )
except Exception as e: except Exception as e:
self.audit_logger.log( self.audit_logger.log(
event_type="CONTAINER_STOP", event_type="container.stop",
severity="ERROR", severity="error",
message=f"Container stop failed: {e}", message=f"Container stop failed: {e}",
details={"container_id": container_id, "error": str(e)} details={"container_id": container_id, "error": str(e)}
) )
raise raise
def remove_container( def remove_container(
self", self,
container_id: str", container_id: str,
force: bool = False force: bool = False
) -> None: ) -> None:
"" """
Remove a container. Remove a container.
Args: Args:
container_id: Container ID to remove container_id: Container ID to remove
force: Force removal even if running force: Force removal even if running
"" """
# Verify container is registered # Verify container is registered
if container_id not in self.validator.session_containers: if container_id not in self.validator.session_containers:
self.audit_logger.log_security_violation( self.audit_logger.log(event_type="security.violation", severity="critical",
operation="container_remove", operation="container_remove",
reason=f"Attempted to remove unregistered container: {container_id}" reason=f"Attempted to remove unregistered container: {container_id}"
) )
@ -323,27 +329,28 @@ class SecureContainerManager:
# Unregister from validator # Unregister from validator
self.validator.unregister_session_container(container_id) self.validator.unregister_session_container(container_id)
self.audit_logger.log(event_type="container.operation", severity="info", operation="remove", self.audit_logger.log_container_operation(
container_id=container_id", operation="remove",
container_id=container_id,
image="", image="",
details={"force": force} details={"force": force}
) )
except Exception as e: except Exception as e:
self.audit_logger.log( self.audit_logger.log(
event_type="CONTAINER_REMOVE", event_type="container.remove",
severity="ERROR", severity="error",
message=f"Container removal failed: {e}", message=f"Container removal failed: {e}",
details={"container_id": container_id, "error": str(e)} details={"container_id": container_id, "error": str(e)}
) )
raise raise
def get_container_logs( def get_container_logs(
self", self,
container_id: str", container_id: str,
tail: int = 100 tail: int = 100
) -> tuple[str, str]: ) -> tuple[str, str]:
"" """
Get container stdout and stderr logs. Get container stdout and stderr logs.
Args: Args:
@ -352,7 +359,7 @@ class SecureContainerManager:
Returns: Returns:
(stdout, stderr) as strings (stdout, stderr) as strings
"" """
if container_id not in self.validator.session_containers: if container_id not in self.validator.session_containers:
raise SecurityError( raise SecurityError(
f"Container {container_id} is not a registered session container" f"Container {container_id} is not a registered session container"
@ -373,23 +380,23 @@ class SecureContainerManager:
logs_str = str(logs) logs_str = str(logs)
# For simplicity, return all logs in stdout (Podman combines them) # For simplicity, return all logs in stdout (Podman combines them)
return logs_str, " return logs_str, ""
except Exception as e: except Exception as e:
self.audit_logger.log( self.audit_logger.log(
event_type="EXECUTION_REQUEST", event_type="execution.request",
severity="ERROR", severity="error",
message=f"Failed to get container logs: {e}", message=f"Failed to get container logs: {e}",
details={"container_id": container_id, "error": str(e)} details={"container_id": container_id, "error": str(e)}
) )
raise raise
def wait_for_container( def wait_for_container(
self", self,
container_id: str", container_id: str,
timeout: int = 300 timeout: int = 300
) -> int: ) -> int:
"" """
Wait for container to exit. Wait for container to exit.
Args: Args:
@ -401,7 +408,7 @@ class SecureContainerManager:
Raises: Raises:
TimeoutError: If container doesn't exit within timeout TimeoutError: If container doesn't exit within timeout
"" """
if container_id not in self.validator.session_containers: if container_id not in self.validator.session_containers:
raise SecurityError( raise SecurityError(
f"Container {container_id} is not a registered session container" f"Container {container_id} is not a registered session container"
@ -421,18 +428,18 @@ class SecureContainerManager:
except Exception as e: except Exception as e:
self.audit_logger.log( self.audit_logger.log(
event_type="EXECUTION_REQUEST", event_type="execution.request",
severity="ERROR", severity="error",
message=f"Failed to wait for container: {e}", message=f"Failed to wait for container: {e}",
details={"container_id": container_id, "error": str(e)} details={"container_id": container_id, "error": str(e)}
) )
raise raise
def cleanup_old_containers( def cleanup_old_containers(
self", self,
max_age: timedelta = timedelta(hours=24) max_age: timedelta = timedelta(hours=24)
) -> int: ) -> int:
"" """
Cleanup containers older than max_age. Cleanup containers older than max_age.
Args: Args:
@ -440,11 +447,11 @@ class SecureContainerManager:
Returns: Returns:
Number of containers removed Number of containers removed
"" """
try: try:
# Get all containers with mcp-forge.session label # Get all containers with mcp-forge.session label
containers = self.podman.client.containers.list( containers = self.podman.client.containers.list(
all=True", all=True,
filters={"label": ["mcp-forge.session"]} filters={"label": ["mcp-forge.session"]}
) )
@ -453,7 +460,7 @@ class SecureContainerManager:
for container in containers: for container in containers:
# Get creation time # Get creation time
created_str = container.attrs.get("Created", ") created_str = container.attrs.get("Created", "")
if not created_str: if not created_str:
continue continue
@ -473,18 +480,18 @@ class SecureContainerManager:
removed_count += 1 removed_count += 1
self.audit_logger.log( self.audit_logger.log(
event_type="CONTAINER_REMOVE", event_type="container.remove",
severity="INFO", severity="info",
message=f"Cleaned up old container: {container.id}", message=f"Cleaned up old container: {container.id}",
details={ details={
"container_id": container.id", "container_id": container.id,
"age_hours": age.total_seconds() / 3600 "age_hours": age.total_seconds() / 3600
} }
) )
except Exception as e: except Exception as e:
self.audit_logger.log( self.audit_logger.log(
event_type="CONTAINER_REMOVE", event_type="container.remove",
severity="WARNING", severity="warning",
message=f"Failed to remove old container: {e}", message=f"Failed to remove old container: {e}",
details={"container_id": container.id, "error": str(e)} details={"container_id": container.id, "error": str(e)}
) )
@ -493,8 +500,8 @@ class SecureContainerManager:
except Exception as e: except Exception as e:
self.audit_logger.log( self.audit_logger.log(
event_type="CONTAINER_REMOVE", event_type="container.remove",
severity="ERROR", severity="error",
message=f"Cleanup failed: {e}", message=f"Cleanup failed: {e}",
details={"error": str(e)} details={"error": str(e)}
) )

126
src/pod_executor/jupyter/backend.py
View file

@ -3,9 +3,10 @@
from typing import Optional, Dict, List from typing import Optional, Dict, List
import hashlib import hashlib
from pod_executor.containers.manager import SecureContainerManager from pod_executor.containers.manager import SecureContainerManager
from pod_executor.security.resource_limits import ResourceLimits, parse_memory_string
from pod_executor.security.audit import AuditLoggerProtocol, NullAuditLogger from pod_executor.security.audit import AuditLoggerProtocol, NullAuditLogger
from pod_executor.security.resource_limits import ResourceLimits, parse_memory_string
from pod_executor.simple.executor import ExecutionResult from pod_executor.simple.executor import ExecutionResult
from pod_executor.jupyter.kernel import JupyterKernelManager from pod_executor.jupyter.kernel import JupyterKernelManager
from pod_executor.jupyter.sessions import SessionManager, SessionState, SessionError from pod_executor.jupyter.sessions import SessionManager, SessionState, SessionError
@ -17,36 +18,25 @@ class JupyterBackend:
def __init__( def __init__(
self, self,
container_manager: SecureContainerManager, container_manager: SecureContainerManager,
image: str, image: str = "mcp-forge/jupyter:latest",
default_timeout: int = 300, default_timeout: int = 300,
default_memory: str = "512m", default_memory: str = "512m",
default_cpu_quota: int = 100000, default_cpu_quota: int = 50000,
max_timeout: int = 3600, max_timeout: int = 1800,
max_memory: str = "2g", max_memory: str = "2g",
max_cpu_quota: int = 200000, max_cpu_quota: int = 100000,
idle_timeout: int = 3600,
max_sessions: int = 10, max_sessions: int = 10,
resource_limits: Optional[ResourceLimits] = None, idle_timeout: int = 3600,
audit_logger: Optional[AuditLoggerProtocol] = None audit_logger: Optional[AuditLoggerProtocol] = None
): ):
""" """
Initialize Jupyter backend. Initialize Jupyter backend.
Args: Args:
config: Forge configuration
container_manager: Container lifecycle manager container_manager: Container lifecycle manager
image: Docker/Podman image with ipykernel installed audit_logger: Audit logging instance
default_timeout: Default execution timeout in seconds
default_memory: Default memory limit (e.g., "512m", "1g")
default_cpu_quota: Default CPU quota (100000 = 1 CPU)
max_timeout: Maximum allowed timeout
max_memory: Maximum allowed memory
max_cpu_quota: Maximum allowed CPU quota
idle_timeout: Session idle timeout in seconds
max_sessions: Maximum concurrent sessions
resource_limits: Default resource limits for kernels
audit_logger: Optional audit logger (uses NullAuditLogger if None)
""" """
self.container_manager = container_manager
self.image = image self.image = image
self.default_timeout = default_timeout self.default_timeout = default_timeout
self.default_memory = default_memory self.default_memory = default_memory
@ -54,23 +44,23 @@ class JupyterBackend:
self.max_timeout = max_timeout self.max_timeout = max_timeout
self.max_memory = max_memory self.max_memory = max_memory
self.max_cpu_quota = max_cpu_quota self.max_cpu_quota = max_cpu_quota
self.idle_timeout = idle_timeout
self.max_sessions = max_sessions self.max_sessions = max_sessions
self.audit_logger = audit_logger or NullAuditLogger() self.idle_timeout = idle_timeout
self.container_manager = container_manager
self.audit_logger = audit_logger
# Initialize kernel manager # Initialize kernel manager
kernel_manager = JupyterKernelManager( kernel_manager = JupyterKernelManager(
container_manager=container_manager, container_manager=container_manager,
image=image, image=config.images.jupyter,
resource_limits=resource_limits resource_limits=self._default_resource_limits()
) )
# Initialize session manager # Initialize session manager
self.session_manager = SessionManager( self.session_manager = SessionManager(
config=config.sessions,
kernel_manager=kernel_manager, kernel_manager=kernel_manager,
idle_timeout=idle_timeout, audit_logger=audit_logger
max_sessions=max_sessions,
audit_logger=self.audit_logger
) )
def execute( def execute(
@ -80,6 +70,7 @@ class JupyterBackend:
timeout: Optional[int] = None, timeout: Optional[int] = None,
memory: Optional[str] = None, memory: Optional[str] = None,
cpu_quota: Optional[int] = None, cpu_quota: Optional[int] = None,
custom_image: Optional[str] = None,
volumes: Optional[Dict[str, dict]] = None, volumes: Optional[Dict[str, dict]] = None,
injection_code: Optional[str] = None, injection_code: Optional[str] = None,
bridge_socket_path: Optional[str] = None bridge_socket_path: Optional[str] = None
@ -93,12 +84,13 @@ class JupyterBackend:
Args: Args:
code: Python code to execute code: Python code to execute
session_id: Unique session identifier session_id: Unique session identifier
timeout: Max execution time in seconds (uses default if None) timeout: Max execution time in seconds (uses config default if None)
memory: Memory limit string (uses default if None) memory: Memory limit string (uses config default if None)
cpu_quota: CPU quota (uses 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 volumes: Volume mounts dict
injection_code: Optional code executed once at session start injection_code: Optional MCP tool injection code (executed once at session start)
bridge_socket_path: Optional path to socket for mounting bridge_socket_path: Optional path to MCP bridge socket for mounting
Returns: Returns:
ExecutionResult with execution output and metadata ExecutionResult with execution output and metadata
@ -107,7 +99,7 @@ class JupyterBackend:
ValueError: If limits exceed configured maximums ValueError: If limits exceed configured maximums
SessionError: If session operation fails SessionError: If session operation fails
""" """
# Use defaults if not specified # Use defaults from config if not specified
timeout = timeout if timeout is not None else self.default_timeout timeout = timeout if timeout is not None else self.default_timeout
memory = memory if memory is not None else self.default_memory memory = memory if memory is not None else self.default_memory
cpu_quota = cpu_quota if cpu_quota is not None else self.default_cpu_quota cpu_quota = cpu_quota if cpu_quota is not None else self.default_cpu_quota
@ -134,11 +126,11 @@ class JupyterBackend:
try: try:
self.session_manager.get_session(session_id) self.session_manager.get_session(session_id)
except SessionError: except SessionError:
# Session doesn't exist, create it # Session doesn't exist, create it with MCP injection
resource_limits = ResourceLimits( resource_limits = ResourceLimits(
memory=memory, memory=memory,
cpu_quota=cpu_quota, cpu_quota=cpu_quota,
storage="1g", storage="1g", # Default storage quota
timeout=timeout timeout=timeout
) )
@ -177,6 +169,9 @@ class JupyterBackend:
Returns: Returns:
Dictionary with updated state info Dictionary with updated state info
Raises:
SessionError: If session doesn't exist
""" """
self.session_manager.document_state( self.session_manager.document_state(
session_id=session_id, session_id=session_id,
@ -185,37 +180,91 @@ class JupyterBackend:
clear=clear clear=clear
) )
# Return updated state
state = self.session_manager.get_session_state(session_id) state = self.session_manager.get_session_state(session_id)
return state.to_dict() return state.to_dict()
def get_session_state(self, session_id: str) -> SessionState: def get_session_state(self, session_id: str) -> SessionState:
"""Get documented state for session.""" """
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) return self.session_manager.get_session_state(session_id)
def destroy_session(self, session_id: str) -> None: def destroy_session(self, session_id: str) -> None:
"""Destroy session and cleanup kernel.""" """
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) self.session_manager.destroy_session(session_id)
def list_sessions(self) -> List[dict]: def list_sessions(self) -> List[dict]:
"""List all active sessions with metadata.""" """
List all active sessions with metadata.
Returns:
List of session dictionaries
"""
return self.session_manager.list_sessions() return self.session_manager.list_sessions()
def cleanup_idle_sessions(self) -> int: def cleanup_idle_sessions(self) -> int:
"""Cleanup sessions idle beyond configured timeout.""" """
Cleanup sessions idle beyond configured timeout.
Returns:
Number of sessions cleaned up
"""
return self.session_manager.cleanup_idle_sessions() 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.default_memory,
cpu_quota=self.default_cpu_quota,
storage="1g",
timeout=self.default_timeout
)
def _validate_limits(self, timeout: int, memory: str, cpu_quota: int) -> None: def _validate_limits(self, timeout: int, memory: str, cpu_quota: int) -> None:
""" """
Validate resource limits against configured maximums. Validate resource limits against configured maximums.
Args:
timeout: Timeout in seconds
memory: Memory limit string
cpu_quota: CPU quota value
Raises: Raises:
ValueError: If any limit exceeds maximum ValueError: If any limit exceeds maximum
""" """
# Validate timeout
if timeout > self.max_timeout: if timeout > self.max_timeout:
raise ValueError( raise ValueError(
f"Timeout {timeout} exceeds maximum {self.max_timeout}" f"Timeout {timeout} exceeds maximum {self.max_timeout}"
) )
# Validate memory
memory_bytes = parse_memory_string(memory) memory_bytes = parse_memory_string(memory)
max_memory_bytes = parse_memory_string(self.max_memory) max_memory_bytes = parse_memory_string(self.max_memory)
if memory_bytes > max_memory_bytes: if memory_bytes > max_memory_bytes:
@ -223,6 +272,7 @@ class JupyterBackend:
f"Memory {memory} exceeds maximum {self.max_memory}" f"Memory {memory} exceeds maximum {self.max_memory}"
) )
# Validate CPU quota
if cpu_quota > self.max_cpu_quota: if cpu_quota > self.max_cpu_quota:
raise ValueError( raise ValueError(
f"CPU quota {cpu_quota} exceeds maximum {self.max_cpu_quota}" f"CPU quota {cpu_quota} exceeds maximum {self.max_cpu_quota}"

22
src/pod_executor/jupyter/sessions.py
View file

@ -108,15 +108,15 @@ class SessionManager:
Initialize session manager. Initialize session manager.
Args: Args:
kernel_manager: Kernel lifecycle manager kernel_manager: Jupyter kernel manager instance
idle_timeout: Session idle timeout in seconds (default: 3600) idle_timeout: Seconds before idle session cleanup (default: 3600)
max_sessions: Maximum concurrent sessions (default: 10) max_sessions: Maximum concurrent sessions (default: 10)
audit_logger: Optional audit logger (uses NullAuditLogger if None) audit_logger: Optional audit logger (defaults to NullAuditLogger)
""" """
self.kernel_manager = kernel_manager
self.idle_timeout = idle_timeout self.idle_timeout = idle_timeout
self.max_sessions = max_sessions self.max_sessions = max_sessions
self.audit_logger = audit_logger or NullAuditLogger() self.kernel_manager = kernel_manager
self.audit_logger = audit_logger if audit_logger is not None else NullAuditLogger()
self.sessions: Dict[str, Session] = {} self.sessions: Dict[str, Session] = {}
def create_session( def create_session(
@ -171,8 +171,8 @@ class SessionManager:
# Log session creation # Log session creation
self.audit_logger.log( self.audit_logger.log(
event_type="SESSION_CREATE, event_type="session.create",
severity="INFO, severity="info",
message=f"Session created: {session_id}", message=f"Session created: {session_id}",
session_id=session_id, session_id=session_id,
details={ details={
@ -373,8 +373,8 @@ class SessionManager:
except Exception as e: except Exception as e:
# Log but continue with cleanup # Log but continue with cleanup
self.audit_logger.log( self.audit_logger.log(
event_type="SESSION_DESTROY, event_type="session.destroy",
severity="WARNING, severity="warning",
message=f"Error shutting down kernel for session {session_id}", message=f"Error shutting down kernel for session {session_id}",
session_id=session_id, session_id=session_id,
error=str(e) error=str(e)
@ -385,8 +385,8 @@ class SessionManager:
# Log destruction # Log destruction
self.audit_logger.log( self.audit_logger.log(
event_type="SESSION_DESTROY, event_type="session.destroy",
severity="INFO, severity="info",
message=f"Session destroyed: {session_id}", message=f"Session destroyed: {session_id}",
session_id=session_id session_id=session_id
) )

278
test_cli.py Executable file
View file

@ -0,0 +1,278 @@
#!/usr/bin/env python3
"""
Simple test CLI for MCP-Forge execution backends.
This is a minimal testing tool for development.
Usage:
# Simple (stateless) backend:
./test_cli.py
# Jupyter (stateful) backend:
./test_cli.py --jupyter
# One-shot execution:
./test_cli.py --execute "print(2 + 2)"
"""
import sys
import os
import tempfile
import argparse
import time
from pathlib import Path
# Add src to path
sys.path.insert(0, str(Path(__file__).parent / "src"))
from mcp_forge.execution.simple.executor import CodeExecutor
from mcp_forge.execution.jupyter.kernel import JupyterKernelManager
from mcp_forge.podman.client import PodmanClient
from mcp_forge.podman.containers import SecureContainerManager
from mcp_forge.security.audit import AuditLogger
from mcp_forge.security.allowlist import OperationValidator, AllowlistManager
class PassthroughValidator(OperationValidator):
"""Dummy validator that allows all operations (for testing only)."""
def __init__(self):
"""Initialize without allowlist manager."""
pass # Skip parent __init__
def validate_operation(self, operation: str, target: str) -> tuple:
"""Allow all operations."""
return (True, None)
# ANSI colors
GREEN = "\033[92m"
RED = "\033[91m"
GRAY = "\033[90m"
RESET = "\033[0m"
class SimpleCLI:
"""Minimal CLI for testing executors."""
def __init__(self, use_jupyter: bool = False):
"""Initialize CLI with minimal setup."""
self.use_jupyter = use_jupyter
# Create temp directory for logs and connection files
self.temp_dir = Path(tempfile.mkdtemp(prefix="mcp_forge_cli_"))
print(f"{GRAY}Using temp dir: {self.temp_dir}{RESET}")
# Setup audit logger
log_file = self.temp_dir / "audit.log"
self.audit_logger = AuditLogger(log_file)
# Setup Podman client (without validator for simplicity)
self.podman_client = PodmanClient(
socket_path="/run/podman/podman.sock",
validator=None, # Skip validation for testing
audit_logger=self.audit_logger
)
# Setup container manager
self.container_manager = SecureContainerManager(
podman_client=self.podman_client,
validator=None, # Skip validation for testing
audit_logger=self.audit_logger
)
# Setup executor or kernel manager
if use_jupyter:
print(f"{GRAY}Using Jupyter (stateful) backend{RESET}")
self.kernel_manager = JupyterKernelManager(
container_manager=self.container_manager,
audit_logger=self.audit_logger,
connection_dir=self.temp_dir
)
self.kernel_id = None
else:
print(f"{GRAY}Using Simple (stateless) backend{RESET}")
self.executor = CodeExecutor(
container_manager=self.container_manager,
audit_logger=self.audit_logger
)
def execute_code(self, code: str) -> dict:
"""Execute code using the configured backend."""
if self.use_jupyter:
# Start kernel if not already started
if not self.kernel_id:
print(f"{GRAY}Starting Jupyter kernel...{RESET}")
self.kernel_id = self.kernel_manager.start_kernel(
image="mcp-forge/jupyter:latest"
)
print(f"{GRAY}Kernel started: {self.kernel_id}{RESET}")
# Execute code
result = self.kernel_manager.execute_code(self.kernel_id, code)
return {
"output": result.output,
"success": result.status == "ok",
"execution_time": result.execution_time,
"error": result.error if hasattr(result, "error") else None
}
else:
# Simple executor
result = self.executor.execute(
code=code,
image="mcp-forge/python:3.12",
timeout=30
)
return {
"output": result.output,
"success": result.exit_code == 0,
"execution_time": result.execution_time,
"error": result.error if result.exit_code != 0 else None
}
def repl(self):
"""Run interactive REPL."""
print(f"\n{GREEN}MCP-Forge Interactive Shell{RESET}")
print(f"{GRAY}Type '.exit' or '.quit' to exit, '.help' for help{RESET}\n")
while True:
try:
# Get code input (support multi-line with empty prompt on continuation)
code_lines = []
while True:
if not code_lines:
line = input(">>> ")
else:
line = input("... ")
code_lines.append(line)
# Check if more lines are needed
code = "\n".join(code_lines)
if not line.strip() or not self._needs_more_lines(code):
break
code = code.strip()
if not code:
continue
# Check for special commands
if code in [".exit", ".quit"]:
print("Goodbye!")
break
elif code == ".help":
print(f"""
{GREEN}Special commands:{RESET}
.exit, .quit - Exit the shell
.help - Show this help
.restart - Restart Jupyter kernel (Jupyter mode only)
.vars - Show variables (Jupyter mode only)
.clear - Clear screen
""".strip())
continue
elif code == ".restart":
if self.use_jupyter and self.kernel_id:
print(f"{GRAY}Restarting kernel...{RESET}")
self.kernel_manager.shutdown_kernel(self.kernel_id)
self.kernel_id = None
print(f"{GREEN}Kernel will restart on next execution{RESET}")
else:
print(f"{RED}Restart only available in Jupyter mode{RESET}")
continue
elif code == ".vars":
if self.use_jupyter:
result = self.execute_code("dir()")
print(f"{GREEN}{result['output']}{RESET}")
else:
print(f"{RED}Variables only available in Jupyter mode (stateless backend){RESET}")
continue
elif code == ".clear":
os.system('clear' if os.name == 'posix' else 'cls')
continue
# Execute code
start_time = time.time()
result = self.execute_code(code)
elapsed = time.time() - start_time
# Display result
if result["success"]:
if result["output"].strip():
print(f"{GREEN}{result['output']}{RESET}")
else:
error_msg = result.get("error") or result["output"]
print(f"{RED}{error_msg}{RESET}")
print(f"{GRAY}({elapsed:.2f}s){RESET}")
except KeyboardInterrupt:
print("\n(Use .exit or .quit to exit)")
except EOFError:
print("\nGoodbye!")
break
except Exception as e:
print(f"{RED}CLI Error: {e}{RESET}")
def cleanup(self):
"""Clean up resources."""
if self.use_jupyter and self.kernel_id:
try:
self.kernel_manager.shutdown_kernel(self.kernel_id)
except Exception as e:
print(f"{GRAY}Warning: Error shutting down kernel: {e}{RESET}")
def _needs_more_lines(self, code: str) -> bool:
"""Check if code needs more lines (simple heuristic)."""
# Check for unclosed brackets/parens
opens = code.count('(') + code.count('[') + code.count('{')
closes = code.count(')') + code.count(']') + code.count('}')
if opens > closes:
return True
# Check for continuation indicators
if code.rstrip().endswith((':', '\\')):
return True
return False
def main():
"""Main entry point."""
parser = argparse.ArgumentParser(
description="Simple CLI for testing MCP-Forge backends"
)
parser.add_argument(
"--jupyter",
action="store_true",
help="Use Jupyter (stateful) backend instead of simple executor"
)
parser.add_argument(
"--execute", "-e",
type=str,
help="Execute code and exit (non-interactive)"
)
args = parser.parse_args()
# Create CLI
cli = SimpleCLI(use_jupyter=args.jupyter)
try:
if args.execute:
# One-shot execution
result = cli.execute_code(args.execute)
if result["success"]:
print(result["output"])
sys.exit(0)
else:
print(result.get("error") or result["output"], file=sys.stderr)
sys.exit(1)
else:
# Interactive REPL
cli.repl()
finally:
cli.cleanup()
if __name__ == "__main__":
main()

1
tests/pod_executor/__init__.py Normal file
View file

@ -0,0 +1 @@
"""Tests for pod_executor package."""

1
tests/pod_executor/security/__init__.py Normal file
View file

@ -0,0 +1 @@
"""Tests for pod_executor.security module."""

275
tests/pod_executor/security/test_resource_limits.py Normal file
View file

@ -0,0 +1,275 @@
"""
Tests for pod_executor resource limits module.
Tests cover all parsing and validation requirements.
"""
import pytest
def test_parse_memory_string_megabytes():
"""Test parsing memory string with megabytes suffix."""
from pod_executor.security.resource_limits import parse_memory_string
result = parse_memory_string("512m")
assert result == 536870912 # 512 * 1024 * 1024
def test_parse_memory_string_gigabytes():
"""Test parsing memory string with gigabytes suffix."""
from pod_executor.security.resource_limits import parse_memory_string
result = parse_memory_string("2g")
assert result == 2147483648 # 2 * 1024 * 1024 * 1024
def test_parse_memory_string_kilobytes():
"""Test parsing memory string with kilobytes suffix."""
from pod_executor.security.resource_limits import parse_memory_string
result = parse_memory_string("1024k")
assert result == 1048576 # 1024 * 1024
def test_parse_memory_string_case_insensitive():
"""Test that memory string parsing is case-insensitive."""
from pod_executor.security.resource_limits import parse_memory_string
assert parse_memory_string("512M") == 536870912
assert parse_memory_string("2G") == 2147483648
assert parse_memory_string("1024K") == 1048576
def test_parse_memory_string_invalid_format_raises_value_error():
"""Test that invalid format raises ValueError."""
from pod_executor.security.resource_limits import parse_memory_string
with pytest.raises(ValueError) as exc_info:
parse_memory_string("invalid")
assert "invalid" in str(exc_info.value).lower() or "format" in str(exc_info.value).lower()
with pytest.raises(ValueError):
parse_memory_string("512x") # Invalid suffix
with pytest.raises(ValueError):
parse_memory_string("abc") # Not a number
def test_parse_memory_string_negative_value_raises_value_error():
"""Test that negative values raise ValueError."""
from pod_executor.security.resource_limits import parse_memory_string
with pytest.raises(ValueError) as exc_info:
parse_memory_string("-512m")
assert "positive" in str(exc_info.value).lower() or "negative" in str(exc_info.value).lower()
def test_parse_memory_string_zero_value_raises_value_error():
"""Test that zero value raises ValueError."""
from pod_executor.security.resource_limits import parse_memory_string
with pytest.raises(ValueError) as exc_info:
parse_memory_string("0m")
assert "positive" in str(exc_info.value).lower() or "zero" in str(exc_info.value).lower()
def test_parse_cpu_quota_valid_value():
"""Test that valid CPU quota values are accepted."""
from pod_executor.security.resource_limits import parse_cpu_quota
result = parse_cpu_quota(50000)
assert result == 50000
result = parse_cpu_quota(100000) # 100% of one core
assert result == 100000
def test_parse_cpu_quota_max_limit():
"""Test that CPU quota has a reasonable maximum (10 cores)."""
from pod_executor.security.resource_limits import parse_cpu_quota
# Should accept up to 1000000 (10 cores)
result = parse_cpu_quota(1000000)
assert result == 1000000
# Should reject more than 10 cores
with pytest.raises(ValueError) as exc_info:
parse_cpu_quota(1000001)
assert "1000000" in str(exc_info.value) or "maximum" in str(exc_info.value).lower()
def test_parse_cpu_quota_negative_raises_value_error():
"""Test that negative CPU quota raises ValueError."""
from pod_executor.security.resource_limits import parse_cpu_quota
with pytest.raises(ValueError) as exc_info:
parse_cpu_quota(-1)
assert "positive" in str(exc_info.value).lower() or "negative" in str(exc_info.value).lower()
def test_parse_cpu_quota_zero_raises_value_error():
"""Test that zero CPU quota raises ValueError."""
from pod_executor.security.resource_limits import parse_cpu_quota
with pytest.raises(ValueError) as exc_info:
parse_cpu_quota(0)
assert "positive" in str(exc_info.value).lower() or "zero" in str(exc_info.value).lower()
def test_parse_storage_string_same_as_memory():
"""Test that storage parsing works the same as memory parsing."""
from pod_executor.security.resource_limits import parse_storage_string
assert parse_storage_string("1g") == 1073741824
assert parse_storage_string("512m") == 536870912
assert parse_storage_string("2048k") == 2097152
def test_resource_limits_class_initialization():
"""Test ResourceLimits class initializes correctly."""
from pod_executor.security.resource_limits import ResourceLimits
limits = ResourceLimits(
memory="512m",
storage="1g",
cpu_quota=50000,
timeout=300
)
assert limits.memory_bytes == 536870912
assert limits.storage_bytes == 1073741824
assert limits.cpu_quota == 50000
assert limits.timeout == 300
def test_resource_limits_validates_memory():
"""Test that ResourceLimits validates memory string."""
from pod_executor.security.resource_limits import ResourceLimits
with pytest.raises(ValueError):
ResourceLimits(
memory="invalid",
storage="1g",
cpu_quota=50000,
timeout=300
)
def test_resource_limits_validates_storage():
"""Test that ResourceLimits validates storage string."""
from pod_executor.security.resource_limits import ResourceLimits
with pytest.raises(ValueError):
ResourceLimits(
memory="512m",
storage="invalid",
cpu_quota=50000,
timeout=300
)
def test_resource_limits_validates_cpu_quota():
"""Test that ResourceLimits validates CPU quota."""
from pod_executor.security.resource_limits import ResourceLimits
with pytest.raises(ValueError):
ResourceLimits(
memory="512m",
storage="1g",
cpu_quota=-1,
timeout=300
)
def test_resource_limits_to_podman_params():
"""Test conversion to Podman container parameters."""
from pod_executor.security.resource_limits import ResourceLimits
limits = ResourceLimits(
memory="512m",
storage="1g",
cpu_quota=50000,
timeout=300
)
params = limits.to_podman_params()
assert isinstance(params, dict)
assert "mem_limit" in params
assert params["mem_limit"] == "536870912" # Should be string for Podman
# CPU quota is set via cpu_quota parameter
assert "cpu_quota" in params
assert params["cpu_quota"] == 50000
def test_resource_limits_default_timeout():
"""Test that ResourceLimits has a default timeout."""
from pod_executor.security.resource_limits import ResourceLimits
limits = ResourceLimits(
memory="512m",
storage="1g",
cpu_quota=50000
)
assert limits.timeout == 300 # Default from signature
def test_parse_memory_string_with_spaces():
"""Test parsing memory strings that have spaces."""
from pod_executor.security.resource_limits import parse_memory_string
# Should handle spaces gracefully (strip them)
result = parse_memory_string(" 512m ")
assert result == 536870912
def test_parse_memory_string_bytes_suffix():
"""Test parsing memory string with bytes suffix (no multiplier)."""
from pod_executor.security.resource_limits import parse_memory_string
# Just a number (bytes) - should this be supported?
# Based on architecture, we support k, m, g suffixes
# Plain numbers should probably raise an error for safety
with pytest.raises(ValueError):
parse_memory_string("1024")
def test_resource_limits_storage_quota_in_podman_params():
"""Test that storage limits are included in Podman params."""
from pod_executor.security.resource_limits import ResourceLimits
limits = ResourceLimits(
memory="512m",
storage="1g",
cpu_quota=50000,
timeout=300
)
params = limits.to_podman_params()
# Storage limit might be set via storage_opt or similar
# The exact parameter depends on Podman API
assert "storage_bytes" in params or "storage_opt" in params
def test_cpu_quota_explanation():
"""Test that CPU quota values have clear meaning."""
from pod_executor.security.resource_limits import parse_cpu_quota
# 100000 = 100% of one CPU core
# 50000 = 50% of one CPU core
# 200000 = 200% = 2 CPU cores
assert parse_cpu_quota(50000) == 50000 # 0.5 cores
assert parse_cpu_quota(100000) == 100000 # 1 core
assert parse_cpu_quota(200000) == 200000 # 2 cores
def test_parse_memory_with_decimal():
"""Test parsing memory strings with decimal values."""
from pod_executor.security.resource_limits import parse_memory_string
# Should handle decimals
result = parse_memory_string("1.5g")
assert result == 1610612736 # 1.5 * 1024 * 1024 * 1024

1
tests/pod_executor/simple/__init__.py Normal file
View file

@ -0,0 +1 @@
"""Tests for pod_executor.simple module."""

299
tests/pod_executor/simple/test_executor.py Normal file
View file

@ -0,0 +1,299 @@
"""Tests for the pod_executor Code Executor module."""
import pytest
from unittest.mock import Mock
import json
from pod_executor.simple.executor import CodeExecutor, ExecutionResult
from pod_executor.containers.manager import SecureContainerManager
from pod_executor.security.resource_limits import ResourceLimits
@pytest.fixture
def resource_limits():
"""Standard resource limits for testing."""
return ResourceLimits(
memory="512m",
cpu_quota=50000,
storage="1g",
timeout=30
)
@pytest.fixture
def mock_container_manager():
"""Mock SecureContainerManager."""
manager = Mock(spec=SecureContainerManager)
# Mock container lifecycle
manager.create_container.return_value = "test-container-123"
manager.start_container.return_value = None
manager.stop_container.return_value = None
manager.remove_container.return_value = None
manager.wait_for_container.return_value = 0 # exit code
manager.get_container_logs.return_value = ("", "") # (stdout, stderr)
return manager
@pytest.fixture
def executor(mock_container_manager, resource_limits):
"""CodeExecutor instance with mocked dependencies."""
return CodeExecutor(
container_manager=mock_container_manager,
image="mcp-forge/python:3.11",
resource_limits=resource_limits
)
def test_execute_simple_python_code_returns_result(executor, mock_container_manager):
"""Test executing simple Python code returns the result."""
# Mock successful execution with result
mock_container_manager.get_container_logs.return_value = (
json.dumps({"result": 42, "error": None}),
""
)
result = executor.execute("2 + 2")
assert result.success is True
assert result.result == 42
assert result.exit_code == 0
assert result.error is None
# Verify container lifecycle
mock_container_manager.create_container.assert_called_once()
mock_container_manager.start_container.assert_called_once_with("test-container-123")
mock_container_manager.wait_for_container.assert_called_once_with("test-container-123", timeout=30)
mock_container_manager.remove_container.assert_called_once_with("test-container-123")
def test_execute_code_with_stdout_capture(executor, mock_container_manager):
"""Test code execution captures stdout."""
mock_container_manager.get_container_logs.return_value = (
json.dumps({"result": None, "error": None}) + "\n" + "Hello, World!",
""
)
result = executor.execute('print("Hello, World!")')
assert result.success is True
assert "Hello, World!" in result.stdout
assert result.stderr == ""
def test_execute_code_with_stderr_capture(executor, mock_container_manager):
"""Test code execution captures stderr."""
mock_container_manager.get_container_logs.return_value = (
json.dumps({"result": None, "error": None}),
"Warning: something happened"
)
result = executor.execute('import sys; print("warning", file=sys.stderr)')
assert result.success is True
assert result.stderr == "Warning: something happened"
def test_execute_code_timeout_enforcement(executor, mock_container_manager):
"""Test code execution enforces timeout."""
# Simulate timeout by having wait_for_container take too long
mock_container_manager.wait_for_container.side_effect = TimeoutError("Container exceeded timeout")
result = executor.execute("import time; time.sleep(60)", timeout=1)
assert result.success is False
assert result.error is not None
assert "timeout" in result.error.lower()
# Verify cleanup still happens
mock_container_manager.remove_container.assert_called_once_with("test-container-123")
def test_execute_code_with_exception_handling(executor, mock_container_manager):
"""Test code execution handles exceptions gracefully."""
error_msg = "ZeroDivisionError: division by zero"
mock_container_manager.get_container_logs.return_value = (
json.dumps({"result": None, "error": error_msg}),
""
)
mock_container_manager.wait_for_container.return_value = 1 # non-zero exit
result = executor.execute("1 / 0")
assert result.success is False
assert result.error == error_msg
assert result.exit_code == 1
def test_execute_code_with_syntax_error_returns_clear_error(executor, mock_container_manager):
"""Test code with syntax error returns clear error message."""
error_msg = "SyntaxError: invalid syntax"
mock_container_manager.get_container_logs.return_value = (
json.dumps({"result": None, "error": error_msg}),
""
)
mock_container_manager.wait_for_container.return_value = 1
result = executor.execute("def foo( :")
assert result.success is False
assert "SyntaxError" in result.error
def test_execute_code_with_runtime_error_returns_clear_error(executor, mock_container_manager):
"""Test code with runtime error returns clear error with traceback."""
error_msg = "NameError: name 'undefined_var' is not defined"
mock_container_manager.get_container_logs.return_value = (
json.dumps({"result": None, "error": error_msg}),
""
)
mock_container_manager.wait_for_container.return_value = 1
result = executor.execute("print(undefined_var)")
assert result.success is False
assert "NameError" in result.error
def test_result_serialization_json_compatible_types(executor, mock_container_manager):
"""Test execution result contains only JSON-serializable data."""
mock_container_manager.get_container_logs.return_value = (
json.dumps({"result": [1, 2, {"key": "value"}], "error": None}),
""
)
result = executor.execute('[1, 2, {"key": "value"}]')
# Verify result can be serialized to JSON
result_dict = result.to_dict()
json_str = json.dumps(result_dict)
assert json_str is not None
# Verify result data
assert result.result == [1, 2, {"key": "value"}]
def test_large_output_handling(executor, mock_container_manager):
"""Test execution handles large output without issues."""
large_output = "x" * 10000 # 10KB of output
mock_container_manager.get_container_logs.return_value = (
json.dumps({"result": None, "error": None}) + "\n" + large_output,
""
)
result = executor.execute('print("x" * 10000)')
assert result.success is True
assert len(result.stdout) >= 10000
def test_execution_result_to_dict(resource_limits):
"""Test ExecutionResult.to_dict() returns proper dictionary."""
result = ExecutionResult(
success=True,
stdout="output",
stderr="",
result=42,
execution_time=0.5,
exit_code=0,
error=None
)
result_dict = result.to_dict()
assert isinstance(result_dict, dict)
assert result_dict["success"] is True
assert result_dict["stdout"] == "output"
assert result_dict["stderr"] == ""
assert result_dict["result"] == 42
assert result_dict["execution_time"] == 0.5
assert result_dict["exit_code"] == 0
assert result_dict["error"] is None
def test_prepare_code_wraps_code_properly(executor):
"""Test _prepare_code wraps code to capture result."""
code = "x = 2 + 2\nx"
wrapped = executor._prepare_code(code)
# Wrapped code should be executable Python
assert "import" in wrapped
assert "json" in wrapped
assert code in wrapped or "2 + 2" in wrapped
def test_parse_output_extracts_result_and_error(executor):
"""Test _parse_output correctly extracts result and error from JSON."""
# Test successful result
stdout = json.dumps({"result": 42, "error": None})
result, error = executor._parse_output(stdout)
assert result == 42
assert error is None
# Test error
stdout = json.dumps({"result": None, "error": "ValueError: invalid"})
result, error = executor._parse_output(stdout)
assert result is None
assert error == "ValueError: invalid"
def test_cleanup_happens_even_on_create_failure(executor, mock_container_manager):
"""Test container cleanup happens even if create fails."""
mock_container_manager.create_container.side_effect = Exception("Create failed")
with pytest.raises(Exception, match="Create failed"):
executor.execute("print('test')")
# No container to remove since create failed
mock_container_manager.remove_container.assert_not_called()
def test_cleanup_happens_even_on_start_failure(executor, mock_container_manager):
"""Test container cleanup happens even if start fails."""
mock_container_manager.start_container.side_effect = Exception("Start failed")
with pytest.raises(Exception, match="Start failed"):
executor.execute("print('test')")
# Container should still be removed
mock_container_manager.remove_container.assert_called_once_with("test-container-123")
def test_execution_time_tracking(executor, mock_container_manager):
"""Test execution time is tracked accurately."""
mock_container_manager.get_container_logs.return_value = (
json.dumps({"result": None, "error": None}),
""
)
result = executor.execute("pass")
assert result.execution_time >= 0
assert isinstance(result.execution_time, float)
def test_execute_with_custom_timeout(executor, mock_container_manager):
"""Test execute respects custom timeout parameter."""
mock_container_manager.get_container_logs.return_value = (
json.dumps({"result": None, "error": None}),
""
)
executor.execute("pass", timeout=60)
# Verify wait was called with custom timeout
mock_container_manager.wait_for_container.assert_called_with("test-container-123", timeout=60)
def test_execute_uses_default_timeout_from_resource_limits(executor, mock_container_manager):
"""Test execute uses default timeout from resource limits when not specified."""
mock_container_manager.get_container_logs.return_value = (
json.dumps({"result": None, "error": None}),
""
)
executor.execute("pass") # No timeout specified
# Should use resource_limits.timeout (30)
mock_container_manager.wait_for_container.assert_called_with("test-container-123", timeout=30)