From 3a6bd0127240ae232146c65c57fb7608a997ece5 Mon Sep 17 00:00:00 2001 From: Hans Aschauer Date: Sat, 7 Feb 2026 11:05:32 +0100 Subject: [PATCH] Adapt tests for pod_executor package and clean up mcp_forge - 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 --- mcp_forge_cli.py | 321 +++++++++ src/mcp_forge/execution/__init__.py | 1 - src/mcp_forge/execution/jupyter/__init__.py | 16 - src/mcp_forge/execution/jupyter/backend.py | 263 -------- src/mcp_forge/execution/jupyter/kernel.py | 631 ------------------ src/mcp_forge/execution/jupyter/kernel_old.py | 418 ------------ src/mcp_forge/execution/jupyter/sessions.py | 435 ------------ src/mcp_forge/execution/simple/__init__.py | 6 - src/mcp_forge/execution/simple/backend.py | 184 ----- src/mcp_forge/execution/simple/executor.py | 222 ------ src/mcp_forge/podman/__init__.py | 11 - src/mcp_forge/podman/client.py | 157 ----- src/mcp_forge/podman/containers.py | 508 -------------- src/mcp_forge/security/__init__.py | 13 +- src/mcp_forge/security/resource_limits.py | 150 ----- src/pod_executor/containers/client.py | 10 +- src/pod_executor/containers/manager.py | 207 +++--- src/pod_executor/jupyter/backend.py | 126 ++-- src/pod_executor/jupyter/sessions.py | 22 +- test_cli.py | 278 ++++++++ tests/pod_executor/__init__.py | 1 + tests/pod_executor/security/__init__.py | 1 + .../security/test_resource_limits.py | 275 ++++++++ tests/pod_executor/simple/__init__.py | 1 + tests/pod_executor/simple/test_executor.py | 299 +++++++++ 25 files changed, 1390 insertions(+), 3166 deletions(-) create mode 100755 mcp_forge_cli.py delete mode 100644 src/mcp_forge/execution/__init__.py delete mode 100644 src/mcp_forge/execution/jupyter/__init__.py delete mode 100644 src/mcp_forge/execution/jupyter/backend.py delete mode 100644 src/mcp_forge/execution/jupyter/kernel.py delete mode 100644 src/mcp_forge/execution/jupyter/kernel_old.py delete mode 100644 src/mcp_forge/execution/jupyter/sessions.py delete mode 100644 src/mcp_forge/execution/simple/__init__.py delete mode 100644 src/mcp_forge/execution/simple/backend.py delete mode 100644 src/mcp_forge/execution/simple/executor.py delete mode 100644 src/mcp_forge/podman/__init__.py delete mode 100644 src/mcp_forge/podman/client.py delete mode 100644 src/mcp_forge/podman/containers.py delete mode 100644 src/mcp_forge/security/resource_limits.py create mode 100755 test_cli.py create mode 100644 tests/pod_executor/__init__.py create mode 100644 tests/pod_executor/security/__init__.py create mode 100644 tests/pod_executor/security/test_resource_limits.py create mode 100644 tests/pod_executor/simple/__init__.py create mode 100644 tests/pod_executor/simple/test_executor.py diff --git a/mcp_forge_cli.py b/mcp_forge_cli.py new file mode 100755 index 0000000..43a11ee --- /dev/null +++ b/mcp_forge_cli.py @@ -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, '', '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() diff --git a/src/mcp_forge/execution/__init__.py b/src/mcp_forge/execution/__init__.py deleted file mode 100644 index c25381d..0000000 --- a/src/mcp_forge/execution/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Execution backends for running code in containers.""" diff --git a/src/mcp_forge/execution/jupyter/__init__.py b/src/mcp_forge/execution/jupyter/__init__.py deleted file mode 100644 index b327bdd..0000000 --- a/src/mcp_forge/execution/jupyter/__init__.py +++ /dev/null @@ -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" -] diff --git a/src/mcp_forge/execution/jupyter/backend.py b/src/mcp_forge/execution/jupyter/backend.py deleted file mode 100644 index d51b3ec..0000000 --- a/src/mcp_forge/execution/jupyter/backend.py +++ /dev/null @@ -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}" - ) diff --git a/src/mcp_forge/execution/jupyter/kernel.py b/src/mcp_forge/execution/jupyter/kernel.py deleted file mode 100644 index 668f1bd..0000000 --- a/src/mcp_forge/execution/jupyter/kernel.py +++ /dev/null @@ -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 diff --git a/src/mcp_forge/execution/jupyter/kernel_old.py b/src/mcp_forge/execution/jupyter/kernel_old.py deleted file mode 100644 index 3540426..0000000 --- a/src/mcp_forge/execution/jupyter/kernel_old.py +++ /dev/null @@ -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 diff --git a/src/mcp_forge/execution/jupyter/sessions.py b/src/mcp_forge/execution/jupyter/sessions.py deleted file mode 100644 index 5b0a494..0000000 --- a/src/mcp_forge/execution/jupyter/sessions.py +++ /dev/null @@ -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" - ) diff --git a/src/mcp_forge/execution/simple/__init__.py b/src/mcp_forge/execution/simple/__init__.py deleted file mode 100644 index 874a3ba..0000000 --- a/src/mcp_forge/execution/simple/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -"""Simple (stateless) execution backend.""" - -from .executor import CodeExecutor, ExecutionResult -from .backend import SimpleBackend - -__all__ = ["CodeExecutor", "ExecutionResult", "SimpleBackend"] diff --git a/src/mcp_forge/execution/simple/backend.py b/src/mcp_forge/execution/simple/backend.py deleted file mode 100644 index 7b58959..0000000 --- a/src/mcp_forge/execution/simple/backend.py +++ /dev/null @@ -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 diff --git a/src/mcp_forge/execution/simple/executor.py b/src/mcp_forge/execution/simple/executor.py deleted file mode 100644 index f7f4c9a..0000000 --- a/src/mcp_forge/execution/simple/executor.py +++ /dev/null @@ -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 diff --git a/src/mcp_forge/podman/__init__.py b/src/mcp_forge/podman/__init__.py deleted file mode 100644 index 1427354..0000000 --- a/src/mcp_forge/podman/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -"""Podman integration module.""" - -from .client import PodmanClient, PodmanConnectionError -from .containers import ContainerConfig, SecureContainerManager - -__all__ = [ - "PodmanClient", - "PodmanConnectionError", - "ContainerConfig", - "SecureContainerManager", -] diff --git a/src/mcp_forge/podman/client.py b/src/mcp_forge/podman/client.py deleted file mode 100644 index e46d34d..0000000 --- a/src/mcp_forge/podman/client.py +++ /dev/null @@ -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 diff --git a/src/mcp_forge/podman/containers.py b/src/mcp_forge/podman/containers.py deleted file mode 100644 index e84f3f7..0000000 --- a/src/mcp_forge/podman/containers.py +++ /dev/null @@ -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 diff --git a/src/mcp_forge/security/__init__.py b/src/mcp_forge/security/__init__.py index 5123bd3..701bacd 100644 --- a/src/mcp_forge/security/__init__.py +++ b/src/mcp_forge/security/__init__.py @@ -1,11 +1,8 @@ """Security module for resource validation and enforcement.""" -from .resource_limits import ( - parse_memory_string, - parse_cpu_quota, - parse_storage_string, - ResourceLimits, -) +# Resource limits moved to pod_executor package +# Import from pod_executor.security.resource_limits instead + from .allowlist import ( SecurityError, OperationValidator, @@ -17,10 +14,6 @@ from .audit import ( ) __all__ = [ - "parse_memory_string", - "parse_cpu_quota", - "parse_storage_string", - "ResourceLimits", "SecurityError", "OperationValidator", "AuditEventType", diff --git a/src/mcp_forge/security/resource_limits.py b/src/mcp_forge/security/resource_limits.py deleted file mode 100644 index 7540dfe..0000000 --- a/src/mcp_forge/security/resource_limits.py +++ /dev/null @@ -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: (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) - } diff --git a/src/pod_executor/containers/client.py b/src/pod_executor/containers/client.py index 2ad3e43..4418912 100644 --- a/src/pod_executor/containers/client.py +++ b/src/pod_executor/containers/client.py @@ -7,7 +7,7 @@ All container operations are validated against security policy. import os from pathlib import Path -from typing import Optional, Union +from typing import Optional from podman import PodmanClient as BasePodmanClient from pod_executor.security.validation import OperationValidatorProtocol @@ -30,9 +30,9 @@ class PodmanClient: def __init__( self, - socket_path: Union[str, Path], + socket_path: Path, validator: OperationValidatorProtocol, - audit_logger: Optional[AuditLoggerProtocol] = None + audit_logger: AuditLoggerProtocol ): """ Initialize Podman client wrapper. @@ -40,11 +40,11 @@ class PodmanClient: Args: socket_path: Path to Podman socket 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.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 def connect(self) -> None: diff --git a/src/pod_executor/containers/manager.py b/src/pod_executor/containers/manager.py index 78872a2..8e57ccc 100644 --- a/src/pod_executor/containers/manager.py +++ b/src/pod_executor/containers/manager.py @@ -1,12 +1,12 @@ -"" +""" 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, Union +from typing import Optional, Dict, List from datetime import datetime, timedelta from pathlib import Path @@ -17,21 +17,21 @@ from pod_executor.security.resource_limits import ResourceLimits class ContainerConfig: - ""Container configuration with security defaults."" + """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", + 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: @@ -44,7 +44,7 @@ class ContainerConfig: 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 {} @@ -56,7 +56,7 @@ class ContainerConfig: self.port_bindings = port_bindings or {} def to_podman_params(self) -> dict: - "" + """ Convert to Podman container create parameters. Ensures all security requirements are included: @@ -68,16 +68,16 @@ class ContainerConfig: 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", + "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"]", + "network_mode": self.network_mode, + "read_only": True, + "security_opt": ["no-new-privileges"], } # Add working_dir only if explicitly set @@ -108,34 +108,34 @@ class ContainerConfig: class SecureContainerManager: - ""Manages container lifecycle with security enforcement."" + """Manages container lifecycle with security enforcement.""" def __init__( - self", - podman_client: PodmanClient", - validator: OperationValidatorProtocol", - audit_logger: Optional[AuditLoggerProtocol] = None + self, + podman_client: PodmanClient, + validator: OperationValidatorProtocol, + audit_logger: AuditLoggerProtocol = None ): - "" + """ Initialize secure container manager. Args: podman_client: Podman client wrapper 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.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( - self", - config: ContainerConfig", - session_id: Optional[str] = None", - name: Optional[str] = None", + self, + config: ContainerConfig, + session_id: Optional[str] = None, + name: Optional[str] = None, **extra_params ) -> str: - "" + """ Create a container with security validation. Args: @@ -149,7 +149,7 @@ class SecureContainerManager: Raises: SecurityError: If configuration violates security policy - "" + """ # Convert config to Podman parameters params = config.to_podman_params() @@ -170,13 +170,16 @@ class SecureContainerManager: try: # Extract image from params for validation self.validator.validate_container_create( - image=config.image", - params=params", + image=config.image, + params=params, session_id=session_id ) except SecurityError as e: # 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 @@ -190,12 +193,13 @@ class SecureContainerManager: self.validator.register_session_container(container_id) # Log successful creation - self.audit_logger.log(event_type="container.operation", severity="info", operation="create", - container_id=container_id", - image=config.image", - session_id=session_id", + self.audit_logger.log_container_operation( + operation="create", + container_id=container_id, + image=config.image, + session_id=session_id, details={ - "name": name", + "name": name, "command": config.command } ) @@ -204,19 +208,19 @@ class SecureContainerManager: except Exception as e: self.audit_logger.log( - event_type="CONTAINER_CREATE", - severity="ERROR", + event_type="container.create", + severity="error", message=f"Container creation failed: {e}", details={ - "image": config.image", - "session_id": session_id", + "image": config.image, + "session_id": session_id, "error": str(e) } ) raise def start_container(self, container_id: str) -> None: - "" + """ Start a container. Args: @@ -224,10 +228,10 @@ class SecureContainerManager: 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( + self.audit_logger.log(event_type="security.violation", severity="critical", operation="container_start", reason=f"Attempted to start unregistered container: {container_id}" ) @@ -239,35 +243,36 @@ class SecureContainerManager: container = self.podman.client.containers.get(container_id) container.start() - self.audit_logger.log(event_type="container.operation", severity="info", operation="start", - container_id=container_id", - image=" # Not available without extra lookup + 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="CONTAINER_START", - severity="ERROR", + event_type="container.start", + severity="error", message=f"Container start failed: {e}", details={"container_id": container_id, "error": str(e)} ) raise def stop_container( - self", - container_id: str", + 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( + self.audit_logger.log(event_type="security.violation", severity="critical", operation="container_stop", reason=f"Attempted to stop unregistered container: {container_id}" ) @@ -279,36 +284,37 @@ class SecureContainerManager: container = self.podman.client.containers.get(container_id) container.stop(timeout=timeout) - self.audit_logger.log(event_type="container.operation", severity="info", operation="stop", - container_id=container_id", + 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="CONTAINER_STOP", - severity="ERROR", + event_type="container.stop", + severity="error", message=f"Container stop failed: {e}", details={"container_id": container_id, "error": str(e)} ) raise def remove_container( - self", - container_id: str", + 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( + self.audit_logger.log(event_type="security.violation", severity="critical", operation="container_remove", reason=f"Attempted to remove unregistered container: {container_id}" ) @@ -323,27 +329,28 @@ class SecureContainerManager: # Unregister from validator self.validator.unregister_session_container(container_id) - self.audit_logger.log(event_type="container.operation", severity="info", operation="remove", - container_id=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="CONTAINER_REMOVE", - severity="ERROR", + event_type="container.remove", + severity="error", message=f"Container removal failed: {e}", details={"container_id": container_id, "error": str(e)} ) raise def get_container_logs( - self", - container_id: str", + self, + container_id: str, tail: int = 100 ) -> tuple[str, str]: - "" + """ Get container stdout and stderr logs. Args: @@ -352,7 +359,7 @@ class SecureContainerManager: 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" @@ -373,23 +380,23 @@ class SecureContainerManager: logs_str = str(logs) # For simplicity, return all logs in stdout (Podman combines them) - return logs_str, " + return logs_str, "" except Exception as e: self.audit_logger.log( - event_type="EXECUTION_REQUEST", - severity="ERROR", + event_type="execution.request", + severity="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", + self, + container_id: str, timeout: int = 300 ) -> int: - "" + """ Wait for container to exit. Args: @@ -401,7 +408,7 @@ class SecureContainerManager: 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" @@ -421,18 +428,18 @@ class SecureContainerManager: except Exception as e: self.audit_logger.log( - event_type="EXECUTION_REQUEST", - severity="ERROR", + event_type="execution.request", + severity="error", message=f"Failed to wait for container: {e}", details={"container_id": container_id, "error": str(e)} ) raise def cleanup_old_containers( - self", + self, max_age: timedelta = timedelta(hours=24) ) -> int: - "" + """ Cleanup containers older than max_age. Args: @@ -440,11 +447,11 @@ class SecureContainerManager: Returns: Number of containers removed - "" + """ try: # Get all containers with mcp-forge.session label containers = self.podman.client.containers.list( - all=True", + all=True, filters={"label": ["mcp-forge.session"]} ) @@ -453,7 +460,7 @@ class SecureContainerManager: for container in containers: # Get creation time - created_str = container.attrs.get("Created", ") + created_str = container.attrs.get("Created", "") if not created_str: continue @@ -473,18 +480,18 @@ class SecureContainerManager: removed_count += 1 self.audit_logger.log( - event_type="CONTAINER_REMOVE", - severity="INFO", + event_type="container.remove", + severity="info", message=f"Cleaned up old container: {container.id}", details={ - "container_id": container.id", + "container_id": container.id, "age_hours": age.total_seconds() / 3600 } ) except Exception as e: self.audit_logger.log( - event_type="CONTAINER_REMOVE", - severity="WARNING", + event_type="container.remove", + severity="warning", message=f"Failed to remove old container: {e}", details={"container_id": container.id, "error": str(e)} ) @@ -493,8 +500,8 @@ class SecureContainerManager: except Exception as e: self.audit_logger.log( - event_type="CONTAINER_REMOVE", - severity="ERROR", + event_type="container.remove", + severity="error", message=f"Cleanup failed: {e}", details={"error": str(e)} ) diff --git a/src/pod_executor/jupyter/backend.py b/src/pod_executor/jupyter/backend.py index dd27681..71b06e7 100644 --- a/src/pod_executor/jupyter/backend.py +++ b/src/pod_executor/jupyter/backend.py @@ -3,9 +3,10 @@ from typing import Optional, Dict, List import hashlib + 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.resource_limits import ResourceLimits, parse_memory_string from pod_executor.simple.executor import ExecutionResult from pod_executor.jupyter.kernel import JupyterKernelManager from pod_executor.jupyter.sessions import SessionManager, SessionState, SessionError @@ -17,36 +18,25 @@ class JupyterBackend: def __init__( self, container_manager: SecureContainerManager, - image: str, + image: str = "mcp-forge/jupyter:latest", default_timeout: int = 300, default_memory: str = "512m", - default_cpu_quota: int = 100000, - max_timeout: int = 3600, + default_cpu_quota: int = 50000, + max_timeout: int = 1800, max_memory: str = "2g", - max_cpu_quota: int = 200000, - idle_timeout: int = 3600, + max_cpu_quota: int = 100000, max_sessions: int = 10, - resource_limits: Optional[ResourceLimits] = None, + idle_timeout: int = 3600, audit_logger: Optional[AuditLoggerProtocol] = None ): """ Initialize Jupyter backend. Args: + config: Forge configuration container_manager: Container lifecycle manager - image: Docker/Podman image with ipykernel installed - 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) + audit_logger: Audit logging instance """ - self.container_manager = container_manager self.image = image self.default_timeout = default_timeout self.default_memory = default_memory @@ -54,23 +44,23 @@ class JupyterBackend: self.max_timeout = max_timeout self.max_memory = max_memory self.max_cpu_quota = max_cpu_quota - self.idle_timeout = idle_timeout 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 kernel_manager = JupyterKernelManager( container_manager=container_manager, - image=image, - resource_limits=resource_limits + image=config.images.jupyter, + resource_limits=self._default_resource_limits() ) # Initialize session manager self.session_manager = SessionManager( + config=config.sessions, kernel_manager=kernel_manager, - idle_timeout=idle_timeout, - max_sessions=max_sessions, - audit_logger=self.audit_logger + audit_logger=audit_logger ) def execute( @@ -80,6 +70,7 @@ class JupyterBackend: 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 @@ -93,12 +84,13 @@ class JupyterBackend: Args: code: Python code to execute session_id: Unique session identifier - timeout: Max execution time in seconds (uses default if None) - memory: Memory limit string (uses default if None) - cpu_quota: CPU quota (uses default if None) + 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 code executed once at session start - bridge_socket_path: Optional path to socket for mounting + 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 @@ -107,7 +99,7 @@ class JupyterBackend: ValueError: If limits exceed configured maximums 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 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 @@ -134,11 +126,11 @@ class JupyterBackend: try: self.session_manager.get_session(session_id) except SessionError: - # Session doesn't exist, create it + # Session doesn't exist, create it with MCP injection resource_limits = ResourceLimits( memory=memory, cpu_quota=cpu_quota, - storage="1g", + storage="1g", # Default storage quota timeout=timeout ) @@ -177,6 +169,9 @@ class JupyterBackend: Returns: Dictionary with updated state info + + Raises: + SessionError: If session doesn't exist """ self.session_manager.document_state( session_id=session_id, @@ -185,37 +180,91 @@ class JupyterBackend: 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.""" + """ + 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.""" + """ + 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.""" + """ + 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.""" + """ + 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.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: """ 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.max_timeout: raise ValueError( f"Timeout {timeout} exceeds maximum {self.max_timeout}" ) + # Validate memory memory_bytes = parse_memory_string(memory) max_memory_bytes = parse_memory_string(self.max_memory) if memory_bytes > max_memory_bytes: @@ -223,6 +272,7 @@ class JupyterBackend: f"Memory {memory} exceeds maximum {self.max_memory}" ) + # Validate CPU quota if cpu_quota > self.max_cpu_quota: raise ValueError( f"CPU quota {cpu_quota} exceeds maximum {self.max_cpu_quota}" diff --git a/src/pod_executor/jupyter/sessions.py b/src/pod_executor/jupyter/sessions.py index 34ccf8c..7342d81 100644 --- a/src/pod_executor/jupyter/sessions.py +++ b/src/pod_executor/jupyter/sessions.py @@ -108,15 +108,15 @@ class SessionManager: Initialize session manager. Args: - kernel_manager: Kernel lifecycle manager - idle_timeout: Session idle timeout in seconds (default: 3600) + kernel_manager: Jupyter kernel manager instance + idle_timeout: Seconds before idle session cleanup (default: 3600) 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.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] = {} def create_session( @@ -171,8 +171,8 @@ class SessionManager: # Log session creation self.audit_logger.log( - event_type="SESSION_CREATE, - severity="INFO, + event_type="session.create", + severity="info", message=f"Session created: {session_id}", session_id=session_id, details={ @@ -373,8 +373,8 @@ class SessionManager: except Exception as e: # Log but continue with cleanup self.audit_logger.log( - event_type="SESSION_DESTROY, - severity="WARNING, + event_type="session.destroy", + severity="warning", message=f"Error shutting down kernel for session {session_id}", session_id=session_id, error=str(e) @@ -385,8 +385,8 @@ class SessionManager: # Log destruction self.audit_logger.log( - event_type="SESSION_DESTROY, - severity="INFO, + event_type="session.destroy", + severity="info", message=f"Session destroyed: {session_id}", session_id=session_id ) diff --git a/test_cli.py b/test_cli.py new file mode 100755 index 0000000..36efb86 --- /dev/null +++ b/test_cli.py @@ -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() diff --git a/tests/pod_executor/__init__.py b/tests/pod_executor/__init__.py new file mode 100644 index 0000000..20c1144 --- /dev/null +++ b/tests/pod_executor/__init__.py @@ -0,0 +1 @@ +"""Tests for pod_executor package.""" diff --git a/tests/pod_executor/security/__init__.py b/tests/pod_executor/security/__init__.py new file mode 100644 index 0000000..4993e70 --- /dev/null +++ b/tests/pod_executor/security/__init__.py @@ -0,0 +1 @@ +"""Tests for pod_executor.security module.""" diff --git a/tests/pod_executor/security/test_resource_limits.py b/tests/pod_executor/security/test_resource_limits.py new file mode 100644 index 0000000..e5251b8 --- /dev/null +++ b/tests/pod_executor/security/test_resource_limits.py @@ -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 diff --git a/tests/pod_executor/simple/__init__.py b/tests/pod_executor/simple/__init__.py new file mode 100644 index 0000000..9f7d943 --- /dev/null +++ b/tests/pod_executor/simple/__init__.py @@ -0,0 +1 @@ +"""Tests for pod_executor.simple module.""" diff --git a/tests/pod_executor/simple/test_executor.py b/tests/pod_executor/simple/test_executor.py new file mode 100644 index 0000000..710f553 --- /dev/null +++ b/tests/pod_executor/simple/test_executor.py @@ -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)