Compare commits
10 commits
db677ae537
...
403b48bc8e
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
403b48bc8e | ||
|
|
114381b71c | ||
|
|
17577d3fba | ||
|
|
ac91362bd2 | ||
|
|
7d9efc5a38 | ||
|
|
9ceeaa1eda | ||
|
|
3a6bd01272 | ||
|
|
63d9b55a00 | ||
|
|
db75b822f4 | ||
|
|
8b6b237be9 |
59 changed files with 6375 additions and 776 deletions
1475
docs/architecture2.md
Normal file
1475
docs/architecture2.md
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -68,6 +68,13 @@
|
||||||
- Phase 5.3: Integration & End-to-End Testing
|
- Phase 5.3: Integration & End-to-End Testing
|
||||||
- Phase 6: Documentation & Deployment
|
- Phase 6: Documentation & Deployment
|
||||||
|
|
||||||
|
**New Workstream (2026-03-04): R execution environments**
|
||||||
|
- Add language selection to execution tooling (`python` + `r`)
|
||||||
|
- Add simple backend R shell execution path
|
||||||
|
- Add stateful Jupyter R kernel support (`IRkernel`)
|
||||||
|
- Add MCP tool injection generator for R wrappers
|
||||||
|
- Add tests first for R simple/stateful execution and server dispatch
|
||||||
|
|
||||||
**Test Count:** 387 tests passing
|
**Test Count:** 387 tests passing
|
||||||
|
|
||||||
**Last Updated:** 2026-02-06
|
**Last Updated:** 2026-02-06
|
||||||
|
|
|
||||||
321
mcp_forge_cli.py
Executable file
321
mcp_forge_cli.py
Executable file
|
|
@ -0,0 +1,321 @@
|
||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
MCP-Forge CLI - Interactive Python REPL with MCP Tools
|
||||||
|
|
||||||
|
A standalone command-line interface for testing and using MCP-Forge backends:
|
||||||
|
- Simple executor (stateless)
|
||||||
|
- Jupyter kernel (stateful)
|
||||||
|
|
||||||
|
With MCP tools injected for enhanced capabilities.
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
python mcp_forge_cli.py # Interactive mode
|
||||||
|
python mcp_forge_cli.py --jupyter # Use Jupyter backend
|
||||||
|
python mcp_forge_cli.py --execute "code" # Execute and exit
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import sys
|
||||||
|
import os
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
|
# Add src to path
|
||||||
|
sys.path.insert(0, str(Path(__file__).parent / "src"))
|
||||||
|
|
||||||
|
from mcp_forge.config.loader import load_config
|
||||||
|
from pod_executor.containers.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 pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
from mcp_forge.adapters import SimpleBackend
|
||||||
|
from mcp_forge.execution.jupyter.backend import JupyterBackend
|
||||||
|
|
||||||
|
|
||||||
|
class MCPForgeCLI:
|
||||||
|
"""Interactive CLI for MCP-Forge."""
|
||||||
|
|
||||||
|
def __init__(self, use_jupyter: bool = False, config_path: Optional[str] = None):
|
||||||
|
"""
|
||||||
|
Initialize CLI.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
use_jupyter: Use Jupyter backend instead of simple executor
|
||||||
|
config_path: Path to config file (default: config.yaml)
|
||||||
|
"""
|
||||||
|
self.use_jupyter = use_jupyter
|
||||||
|
self.session_id = "cli-session"
|
||||||
|
|
||||||
|
# Load configuration
|
||||||
|
config_file = config_path or "config.yaml"
|
||||||
|
if not Path(config_file).exists():
|
||||||
|
config_file = "config.example.yaml"
|
||||||
|
|
||||||
|
print(f"Loading config from: {config_file}")
|
||||||
|
self.config = load_config(config_file)
|
||||||
|
|
||||||
|
# Initialize components
|
||||||
|
self.audit_store = MemoryAuditStore()
|
||||||
|
self.audit_logger = AuditLogger(self.audit_store)
|
||||||
|
|
||||||
|
self.podman_client = PodmanClient()
|
||||||
|
|
||||||
|
allowlist_manager = AllowlistManager(self.config.allowlist_file)
|
||||||
|
self.validator = OperationValidator(allowlist_manager)
|
||||||
|
|
||||||
|
self.container_manager = SecureContainerManager(
|
||||||
|
podman_client=self.podman_client,
|
||||||
|
validator=self.validator,
|
||||||
|
audit_logger=self.audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
# Initialize backend
|
||||||
|
if use_jupyter:
|
||||||
|
print("🚀 Initializing Jupyter backend (stateful)...")
|
||||||
|
self.backend = JupyterBackend(
|
||||||
|
config=self.config,
|
||||||
|
container_manager=self.container_manager,
|
||||||
|
audit_logger=self.audit_logger
|
||||||
|
)
|
||||||
|
self.backend_name = "Jupyter"
|
||||||
|
else:
|
||||||
|
print("🚀 Initializing Simple backend (stateless)...")
|
||||||
|
self.backend = SimpleBackend(
|
||||||
|
config=self.config,
|
||||||
|
container_manager=self.container_manager,
|
||||||
|
audit_logger=self.audit_logger
|
||||||
|
)
|
||||||
|
self.backend_name = "Simple"
|
||||||
|
|
||||||
|
print(f"✓ {self.backend_name} backend ready")
|
||||||
|
print()
|
||||||
|
|
||||||
|
def execute_code(self, code: str, show_result: bool = True) -> dict:
|
||||||
|
"""
|
||||||
|
Execute Python code.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
code: Python code to execute
|
||||||
|
show_result: Whether to print the result
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Execution result dictionary
|
||||||
|
"""
|
||||||
|
if self.use_jupyter:
|
||||||
|
# Stateful execution with session
|
||||||
|
result = self.backend.execute(
|
||||||
|
code=code,
|
||||||
|
session_id=self.session_id
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
# Stateless execution
|
||||||
|
result = self.backend.execute(code=code)
|
||||||
|
|
||||||
|
if show_result:
|
||||||
|
self._display_result(result)
|
||||||
|
|
||||||
|
return result
|
||||||
|
|
||||||
|
def _display_result(self, result):
|
||||||
|
"""Display execution result."""
|
||||||
|
if result.stdout:
|
||||||
|
print(result.stdout, end='')
|
||||||
|
|
||||||
|
if result.stderr:
|
||||||
|
print(f"\033[31m{result.stderr}\033[0m", end='')
|
||||||
|
|
||||||
|
if result.result and result.result != 'None':
|
||||||
|
print(f"\033[32m{result.result}\033[0m")
|
||||||
|
|
||||||
|
if not result.success and result.error:
|
||||||
|
print(f"\033[31mError: {result.error}\033[0m")
|
||||||
|
|
||||||
|
if result.execution_time > 0.1:
|
||||||
|
print(f"\033[90m({result.execution_time:.3f}s)\033[0m")
|
||||||
|
|
||||||
|
def repl(self):
|
||||||
|
"""Run interactive REPL."""
|
||||||
|
print("=" * 60)
|
||||||
|
print(f" MCP-Forge CLI - {self.backend_name} Backend")
|
||||||
|
print("=" * 60)
|
||||||
|
print()
|
||||||
|
print("Type Python code and press Enter to execute.")
|
||||||
|
print("Special commands:")
|
||||||
|
print(" .exit, .quit - Exit")
|
||||||
|
print(" .clear - Clear screen")
|
||||||
|
print(" .help - Show this help")
|
||||||
|
if self.use_jupyter:
|
||||||
|
print(" .restart - Restart kernel")
|
||||||
|
print(" .vars - Show variables")
|
||||||
|
print()
|
||||||
|
print("Tip: Code execution is " +
|
||||||
|
("stateful (variables persist)" if self.use_jupyter else "stateless"))
|
||||||
|
print()
|
||||||
|
|
||||||
|
buffer = []
|
||||||
|
|
||||||
|
try:
|
||||||
|
while True:
|
||||||
|
try:
|
||||||
|
if buffer:
|
||||||
|
prompt = "... "
|
||||||
|
else:
|
||||||
|
prompt = ">>> "
|
||||||
|
|
||||||
|
line = input(prompt)
|
||||||
|
|
||||||
|
# Handle special commands
|
||||||
|
if not buffer and line.startswith('.'):
|
||||||
|
if line in ['.exit', '.quit']:
|
||||||
|
print("\nGoodbye!")
|
||||||
|
break
|
||||||
|
elif line == '.clear':
|
||||||
|
os.system('clear' if os.name != 'nt' else 'cls')
|
||||||
|
continue
|
||||||
|
elif line == '.help':
|
||||||
|
self.repl()
|
||||||
|
return
|
||||||
|
elif line == '.restart' and self.use_jupyter:
|
||||||
|
self._restart_session()
|
||||||
|
continue
|
||||||
|
elif line == '.vars' and self.use_jupyter:
|
||||||
|
self._show_variables()
|
||||||
|
continue
|
||||||
|
else:
|
||||||
|
print(f"Unknown command: {line}")
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Build multi-line code
|
||||||
|
buffer.append(line)
|
||||||
|
|
||||||
|
# Check if we should execute
|
||||||
|
code = '\n'.join(buffer)
|
||||||
|
|
||||||
|
# Simple heuristic: execute if line is complete
|
||||||
|
# (doesn't end with :, \, or is blank after content)
|
||||||
|
if line and not line.rstrip().endswith((':', '\\')):
|
||||||
|
try:
|
||||||
|
compile(code, '<stdin>', 'exec')
|
||||||
|
# Code is complete, execute it
|
||||||
|
self.execute_code(code)
|
||||||
|
buffer = []
|
||||||
|
except SyntaxError as e:
|
||||||
|
if 'unexpected EOF' in str(e):
|
||||||
|
# Need more input
|
||||||
|
continue
|
||||||
|
else:
|
||||||
|
# Real syntax error
|
||||||
|
self.execute_code(code)
|
||||||
|
buffer = []
|
||||||
|
elif not line and buffer:
|
||||||
|
# Empty line after content, execute
|
||||||
|
self.execute_code(code)
|
||||||
|
buffer = []
|
||||||
|
|
||||||
|
except KeyboardInterrupt:
|
||||||
|
print("\nKeyboardInterrupt")
|
||||||
|
buffer = []
|
||||||
|
continue
|
||||||
|
except EOFError:
|
||||||
|
print("\nGoodbye!")
|
||||||
|
break
|
||||||
|
|
||||||
|
finally:
|
||||||
|
self._cleanup()
|
||||||
|
|
||||||
|
def _restart_session(self):
|
||||||
|
"""Restart Jupyter session."""
|
||||||
|
if self.use_jupyter:
|
||||||
|
try:
|
||||||
|
self.backend.session_manager.close_session(self.session_id)
|
||||||
|
print("✓ Session restarted")
|
||||||
|
except Exception as e:
|
||||||
|
print(f"Error restarting session: {e}")
|
||||||
|
|
||||||
|
def _show_variables(self):
|
||||||
|
"""Show variables in Jupyter session."""
|
||||||
|
if self.use_jupyter:
|
||||||
|
try:
|
||||||
|
session = self.backend.session_manager.get_session(self.session_id)
|
||||||
|
kernel_id = session.kernel_id
|
||||||
|
variables = self.backend.session_manager.kernel_manager.inspect_namespace(kernel_id)
|
||||||
|
|
||||||
|
if variables:
|
||||||
|
print("Variables:")
|
||||||
|
for var in variables:
|
||||||
|
print(f" - {var}")
|
||||||
|
else:
|
||||||
|
print("No variables defined")
|
||||||
|
except Exception as e:
|
||||||
|
print(f"Error showing variables: {e}")
|
||||||
|
|
||||||
|
def _cleanup(self):
|
||||||
|
"""Cleanup resources."""
|
||||||
|
print("\nCleaning up...")
|
||||||
|
if self.use_jupyter:
|
||||||
|
try:
|
||||||
|
self.backend.session_manager.close_session(self.session_id)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
"""Main entry point."""
|
||||||
|
parser = argparse.ArgumentParser(
|
||||||
|
description="MCP-Forge CLI - Interactive Python with MCP tools",
|
||||||
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||||
|
epilog="""
|
||||||
|
Examples:
|
||||||
|
%(prog)s # Simple backend, interactive
|
||||||
|
%(prog)s --jupyter # Jupyter backend, interactive
|
||||||
|
%(prog)s --execute "print(2+2)" # Execute code and exit
|
||||||
|
%(prog)s -j -e "x=5; print(x)" # Jupyter backend, execute code
|
||||||
|
"""
|
||||||
|
)
|
||||||
|
|
||||||
|
parser.add_argument(
|
||||||
|
'-j', '--jupyter',
|
||||||
|
action='store_true',
|
||||||
|
help='Use Jupyter backend (stateful) instead of simple backend'
|
||||||
|
)
|
||||||
|
|
||||||
|
parser.add_argument(
|
||||||
|
'-e', '--execute',
|
||||||
|
metavar='CODE',
|
||||||
|
help='Execute code and exit (non-interactive)'
|
||||||
|
)
|
||||||
|
|
||||||
|
parser.add_argument(
|
||||||
|
'-c', '--config',
|
||||||
|
metavar='FILE',
|
||||||
|
help='Path to config file (default: config.yaml)'
|
||||||
|
)
|
||||||
|
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
try:
|
||||||
|
cli = MCPForgeCLI(
|
||||||
|
use_jupyter=args.jupyter,
|
||||||
|
config_path=args.config
|
||||||
|
)
|
||||||
|
|
||||||
|
if args.execute:
|
||||||
|
# Execute code and exit
|
||||||
|
cli.execute_code(args.execute)
|
||||||
|
else:
|
||||||
|
# Interactive mode
|
||||||
|
cli.repl()
|
||||||
|
|
||||||
|
except KeyboardInterrupt:
|
||||||
|
print("\nInterrupted")
|
||||||
|
sys.exit(1)
|
||||||
|
except Exception as e:
|
||||||
|
print(f"\nError: {e}", file=sys.stderr)
|
||||||
|
import traceback
|
||||||
|
traceback.print_exc()
|
||||||
|
sys.exit(1)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
main()
|
||||||
432
simple_test_cli.py
Executable file
432
simple_test_cli.py
Executable file
|
|
@ -0,0 +1,432 @@
|
||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
Simplified test CLI for MCP-Forge execution backends.
|
||||||
|
|
||||||
|
This is a minimal testing tool that bypasses full security stack.
|
||||||
|
MCP injection is not supported - this is purely for testing execution.
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
# Simple (stateless) backend:
|
||||||
|
./simple_test_cli.py
|
||||||
|
|
||||||
|
# Jupyter (stateful) backend:
|
||||||
|
./simple_test_cli.py --jupyter
|
||||||
|
|
||||||
|
# One-shot execution:
|
||||||
|
./simple_test_cli.py --execute "print(2 + 2)"
|
||||||
|
"""
|
||||||
|
|
||||||
|
import sys
|
||||||
|
import os
|
||||||
|
import tempfile
|
||||||
|
import argparse
|
||||||
|
import time
|
||||||
|
import uuid
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
# Add src to path
|
||||||
|
sys.path.insert(0, str(Path(__file__).parent / "src"))
|
||||||
|
|
||||||
|
from pod_executor.simple.executor import ExecutionResult
|
||||||
|
from mcp_forge.execution.jupyter.kernel import JupyterKernelManager
|
||||||
|
from pod_executor.containers.client import PodmanClient
|
||||||
|
from mcp_forge.podman.containers import SecureContainerManager, ContainerConfig
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
import time
|
||||||
|
from typing import Optional, Dict
|
||||||
|
|
||||||
|
|
||||||
|
# ANSI colors
|
||||||
|
GREEN = "\033[92m"
|
||||||
|
RED = "\033[91m"
|
||||||
|
GRAY = "\033[90m"
|
||||||
|
YELLOW = "\033[93m"
|
||||||
|
RESET = "\033[0m"
|
||||||
|
|
||||||
|
|
||||||
|
class PassthroughValidator(OperationValidator):
|
||||||
|
"""Dummy validator that allows all operations (for testing only)."""
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
"""Initialize without allowlist manager."""
|
||||||
|
# Don't call parent __init__ to avoid needing AllowlistManager
|
||||||
|
# But we need to initialize the session_containers set
|
||||||
|
self.session_containers = set()
|
||||||
|
|
||||||
|
def validate_operation(self, operation: str, target: str) -> tuple:
|
||||||
|
"""Allow all operations without validation."""
|
||||||
|
return (True, None)
|
||||||
|
|
||||||
|
def validate_container_create(self, image: str, params: dict, session_id: str = None) -> None:
|
||||||
|
"""Allow all container creations."""
|
||||||
|
pass
|
||||||
|
|
||||||
|
def register_session_container(self, container_id: str) -> None:
|
||||||
|
"""Track session containers."""
|
||||||
|
self.session_containers.add(container_id)
|
||||||
|
|
||||||
|
def unregister_session_container(self, container_id: str) -> None:
|
||||||
|
"""Untrack session containers."""
|
||||||
|
self.session_containers.discard(container_id)
|
||||||
|
|
||||||
|
|
||||||
|
class SimpleExecutorWrapper:
|
||||||
|
"""Wrapper for CodeExecutor that bypasses resource limits."""
|
||||||
|
|
||||||
|
def __init__(self, container_manager: SecureContainerManager, image: str):
|
||||||
|
self.container_manager = container_manager
|
||||||
|
self.image = image
|
||||||
|
|
||||||
|
def execute(self, code: str, timeout: int = 30) -> ExecutionResult:
|
||||||
|
"""Execute code without resource limits."""
|
||||||
|
start_time = time.time()
|
||||||
|
container_id = None
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Prepare code with result extraction
|
||||||
|
wrapped_code = f"""
|
||||||
|
import sys
|
||||||
|
import json
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Execute code
|
||||||
|
exec_globals = {{'__builtins__': __builtins__}}
|
||||||
|
exec({repr(code)}, exec_globals)
|
||||||
|
|
||||||
|
# Get result (last expression value if any)
|
||||||
|
result = exec_globals.get('_', None)
|
||||||
|
except Exception as e:
|
||||||
|
import traceback
|
||||||
|
print(traceback.format_exc(), file=sys.stderr)
|
||||||
|
sys.exit(1)
|
||||||
|
"""
|
||||||
|
|
||||||
|
# Create container configuration WITHOUT resource limits
|
||||||
|
config = ContainerConfig(
|
||||||
|
image=self.image,
|
||||||
|
command=["python", "-c", wrapped_code],
|
||||||
|
resource_limits=None # Skip resource limits for rootless testing
|
||||||
|
)
|
||||||
|
|
||||||
|
# Create and start container
|
||||||
|
session_id = str(uuid.uuid4())
|
||||||
|
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=timeout)
|
||||||
|
|
||||||
|
# Get logs (returns tuple of (stdout, stderr))
|
||||||
|
stdout, stderr = self.container_manager.get_container_logs(container_id)
|
||||||
|
|
||||||
|
execution_time = time.time() - start_time
|
||||||
|
|
||||||
|
return ExecutionResult(
|
||||||
|
success=(exit_code == 0),
|
||||||
|
stdout=stdout,
|
||||||
|
stderr=stderr,
|
||||||
|
exit_code=exit_code,
|
||||||
|
execution_time=execution_time,
|
||||||
|
result=None, # Don't try to extract result for simplicity
|
||||||
|
error=stderr if exit_code != 0 else None
|
||||||
|
)
|
||||||
|
|
||||||
|
finally:
|
||||||
|
# Cleanup
|
||||||
|
if container_id:
|
||||||
|
try:
|
||||||
|
self.container_manager.remove_container(container_id)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class JupyterKernelWrapper:
|
||||||
|
"""Wrapper for JupyterKernelManager that bypasses resource limits."""
|
||||||
|
|
||||||
|
def __init__(self, container_manager: SecureContainerManager, image: str, connection_dir: Path):
|
||||||
|
# Create minimal resource limits
|
||||||
|
minimal_limits = ResourceLimits(
|
||||||
|
memory="512m",
|
||||||
|
cpu_quota=100000,
|
||||||
|
storage="1g",
|
||||||
|
timeout=30
|
||||||
|
)
|
||||||
|
self.manager = JupyterKernelManager(
|
||||||
|
container_manager=container_manager,
|
||||||
|
image=image,
|
||||||
|
resource_limits=minimal_limits
|
||||||
|
)
|
||||||
|
# Override resource_limits to None after construction to avoid cgroup issues
|
||||||
|
self.manager.resource_limits = None
|
||||||
|
self.connection_dir = connection_dir
|
||||||
|
|
||||||
|
def start_kernel(self, session_id: str) -> str:
|
||||||
|
"""Start kernel."""
|
||||||
|
return self.manager.start_kernel(session_id)
|
||||||
|
|
||||||
|
def execute_code(self, kernel_id: str, code: str) -> ExecutionResult:
|
||||||
|
"""Execute code in kernel."""
|
||||||
|
return self.manager.execute_code(kernel_id, code)
|
||||||
|
|
||||||
|
def shutdown_kernel(self, kernel_id: str) -> None:
|
||||||
|
"""Shutdown kernel."""
|
||||||
|
self.manager.shutdown_kernel(kernel_id)
|
||||||
|
|
||||||
|
|
||||||
|
class SimpleCLI:
|
||||||
|
"""Minimal CLI for testing executors without full security stack."""
|
||||||
|
|
||||||
|
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 with passthrough validator
|
||||||
|
# Use user socket (rootless Podman)
|
||||||
|
runtime_dir = os.environ.get("XDG_RUNTIME_DIR", f"/run/user/{os.getuid()}")
|
||||||
|
socket_path = Path(runtime_dir) / "podman" / "podman.sock"
|
||||||
|
|
||||||
|
self.validator = PassthroughValidator()
|
||||||
|
self.podman_client = PodmanClient(
|
||||||
|
socket_path=socket_path,
|
||||||
|
validator=self.validator,
|
||||||
|
audit_logger=self.audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
# Setup container manager
|
||||||
|
self.container_manager = SecureContainerManager(
|
||||||
|
podman_client=self.podman_client,
|
||||||
|
validator=self.validator,
|
||||||
|
audit_logger=self.audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
# Setup executor or kernel manager
|
||||||
|
if use_jupyter:
|
||||||
|
print(f"{GRAY}Using Jupyter (stateful) backend{RESET}")
|
||||||
|
self.kernel_manager = JupyterKernelWrapper(
|
||||||
|
container_manager=self.container_manager,
|
||||||
|
image="mcp-forge/jupyter:latest",
|
||||||
|
connection_dir=self.temp_dir
|
||||||
|
)
|
||||||
|
self.session_id = str(uuid.uuid4())
|
||||||
|
self.kernel_id = None
|
||||||
|
else:
|
||||||
|
print(f"{GRAY}Using Simple (stateless) backend{RESET}")
|
||||||
|
self.executor = SimpleExecutorWrapper(
|
||||||
|
container_manager=self.container_manager,
|
||||||
|
image="mcp-forge/python:3.12"
|
||||||
|
)
|
||||||
|
|
||||||
|
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(
|
||||||
|
session_id=self.session_id
|
||||||
|
)
|
||||||
|
print(f"{GRAY}Kernel started: {self.kernel_id}{RESET}")
|
||||||
|
|
||||||
|
# Execute code
|
||||||
|
result = self.kernel_manager.execute_code(self.kernel_id, code)
|
||||||
|
return {
|
||||||
|
"output": result.stdout + result.stderr,
|
||||||
|
"result": result.result,
|
||||||
|
"success": result.exit_code == 0,
|
||||||
|
"execution_time": result.execution_time,
|
||||||
|
"exit_code": result.exit_code
|
||||||
|
}
|
||||||
|
else:
|
||||||
|
# Simple executor
|
||||||
|
result = self.executor.execute(code=code)
|
||||||
|
return {
|
||||||
|
"output": result.stdout + result.stderr,
|
||||||
|
"result": result.result,
|
||||||
|
"success": result.exit_code == 0,
|
||||||
|
"execution_time": result.execution_time,
|
||||||
|
"exit_code": result.exit_code
|
||||||
|
}
|
||||||
|
|
||||||
|
def repl(self):
|
||||||
|
"""Run interactive REPL."""
|
||||||
|
print(f"\n{GREEN}MCP-Forge Interactive Shell{RESET}")
|
||||||
|
if self.use_jupyter:
|
||||||
|
print(f"{YELLOW}Note: Variables persist between executions in Jupyter mode{RESET}")
|
||||||
|
else:
|
||||||
|
print(f"{YELLOW}Note: Each execution is isolated in Simple mode{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":
|
||||||
|
self._show_help()
|
||||||
|
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:
|
||||||
|
# Show variables by running dir()
|
||||||
|
result = self.execute_code("print([v for v in dir() if not v.startswith('_')])")
|
||||||
|
if result["success"]:
|
||||||
|
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"]:
|
||||||
|
# Show output (stdout/stderr)
|
||||||
|
if result["output"].strip():
|
||||||
|
print(result["output"])
|
||||||
|
# Show result value if different from output
|
||||||
|
if result["result"] and result["result"] not in [None, "None"]:
|
||||||
|
print(f"{GREEN}→ {result['result']}{RESET}")
|
||||||
|
else:
|
||||||
|
# Show error
|
||||||
|
print(f"{RED}{result['output']}{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}")
|
||||||
|
import traceback
|
||||||
|
traceback.print_exc()
|
||||||
|
|
||||||
|
def cleanup(self):
|
||||||
|
"""Clean up resources."""
|
||||||
|
if self.use_jupyter and self.kernel_id:
|
||||||
|
try:
|
||||||
|
print(f"\n{GRAY}Shutting down kernel...{RESET}")
|
||||||
|
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 _show_help(self):
|
||||||
|
"""Show help message."""
|
||||||
|
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
|
||||||
|
|
||||||
|
{GREEN}Tips:{RESET}
|
||||||
|
- Multi-line input: Press Enter on a blank line to execute
|
||||||
|
- Jupyter mode: Variables persist between executions
|
||||||
|
- Simple mode: Each execution is in a fresh container
|
||||||
|
""".strip())
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
"""Main entry point."""
|
||||||
|
parser = argparse.ArgumentParser(
|
||||||
|
description="Simple CLI for testing MCP-Forge backends"
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--jupyter", "-j",
|
||||||
|
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 output
|
||||||
|
if result["output"].strip():
|
||||||
|
print(result["output"])
|
||||||
|
# Print result if available
|
||||||
|
if result["result"] and result["result"] not in [None, "None"]:
|
||||||
|
print(result["result"])
|
||||||
|
sys.exit(0)
|
||||||
|
else:
|
||||||
|
print(result["output"], file=sys.stderr)
|
||||||
|
sys.exit(1)
|
||||||
|
else:
|
||||||
|
# Interactive REPL
|
||||||
|
cli.repl()
|
||||||
|
finally:
|
||||||
|
cli.cleanup()
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
84
simple_test_cli_README.md
Normal file
84
simple_test_cli_README.md
Normal file
|
|
@ -0,0 +1,84 @@
|
||||||
|
# Simple Test CLI for MCP-Forge
|
||||||
|
|
||||||
|
A minimal testing tool for MCP-Forge execution backends.
|
||||||
|
|
||||||
|
## Features
|
||||||
|
|
||||||
|
✅ **Working:**
|
||||||
|
- Simple (stateless) backend with one-shot and interactive execution
|
||||||
|
- No security restrictions (for testing only)
|
||||||
|
- No resource limits (avoids cgroupv2 issues in rootless Podman)
|
||||||
|
|
||||||
|
❌ **Not Working:**
|
||||||
|
- Jupyter (stateful) backend - has connection file timing issues
|
||||||
|
- MCP tool injection - not supported in this simplified version
|
||||||
|
|
||||||
|
## Usage
|
||||||
|
|
||||||
|
### One-Shot Execution (Simple Backend)
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Execute Python code and exit
|
||||||
|
./simple_test_cli.py --execute 'print(2 + 2)'
|
||||||
|
|
||||||
|
# More complex code
|
||||||
|
./simple_test_cli.py --execute 'x = [1, 2, 3]; print(sum(x))'
|
||||||
|
```
|
||||||
|
|
||||||
|
### Interactive REPL (Simple Backend)
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Start interactive shell
|
||||||
|
./simple_test_cli.py
|
||||||
|
|
||||||
|
# In the shell:
|
||||||
|
>>> x = 42
|
||||||
|
>>> print(x * 2)
|
||||||
|
84
|
||||||
|
>>> .exit
|
||||||
|
```
|
||||||
|
|
||||||
|
### Special Commands
|
||||||
|
|
||||||
|
- `.exit`, `.quit` - Exit the shell
|
||||||
|
- `.help` - Show help
|
||||||
|
- `.clear` - Clear screen
|
||||||
|
|
||||||
|
## Requirements
|
||||||
|
|
||||||
|
- Rootless Podman with user socket at `$XDG_RUNTIME_DIR/podman/podman.sock`
|
||||||
|
- Container images built:
|
||||||
|
- `mcp-forge/python:3.12` - For simple backend
|
||||||
|
- `mcp-forge/jupyter:latest` - For Jupyter backend (not working yet)
|
||||||
|
|
||||||
|
## Limitations
|
||||||
|
|
||||||
|
This is a **testing tool only**:
|
||||||
|
- No security validation
|
||||||
|
- No resource limits
|
||||||
|
- No MCP tool injection
|
||||||
|
- No proper error handling for production use
|
||||||
|
- Jupyter backend has connection file issues
|
||||||
|
|
||||||
|
For production use, use the full MCP-Forge server with proper configuration.
|
||||||
|
|
||||||
|
## Testing Containers
|
||||||
|
|
||||||
|
To test if containers work:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Test simple execution
|
||||||
|
./test_containers.py
|
||||||
|
|
||||||
|
# Should output:
|
||||||
|
# ✓ Simple Execution: PASS
|
||||||
|
# ✓ Jupyter Kernel: PASS
|
||||||
|
```
|
||||||
|
|
||||||
|
## Architecture
|
||||||
|
|
||||||
|
- `PassthroughValidator` - Dummy validator that allows all operations
|
||||||
|
- `SimpleExecutorWrapper` - Wraps CodeExecutor, skips resource limits
|
||||||
|
- `JupyterKernelWrapper` - Wraps JupyterKernelManager (has issues)
|
||||||
|
|
||||||
|
The CLI creates temporary directories for logs and connection files, cleans up on exit.
|
||||||
13
src/mcp_forge/adapters/__init__.py
Normal file
13
src/mcp_forge/adapters/__init__.py
Normal file
|
|
@ -0,0 +1,13 @@
|
||||||
|
"""Adapter layer between MCP-Forge and pod_executor.
|
||||||
|
|
||||||
|
This module provides wrappers that adapt pod_executor components
|
||||||
|
to work with MCP-Forge's configuration and infrastructure.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from .executor_adapter import SimpleBackend
|
||||||
|
from .jupyter_adapter import JupyterBackend
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"SimpleBackend",
|
||||||
|
"JupyterBackend",
|
||||||
|
]
|
||||||
121
src/mcp_forge/adapters/executor_adapter.py
Normal file
121
src/mcp_forge/adapters/executor_adapter.py
Normal file
|
|
@ -0,0 +1,121 @@
|
||||||
|
"""Simple executor backend adapter for MCP-Forge.
|
||||||
|
|
||||||
|
Wraps pod_executor.CodeExecutor with MCP-Forge configuration.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import Optional
|
||||||
|
from pathlib import Path
|
||||||
|
import logging
|
||||||
|
|
||||||
|
from pod_executor import CodeExecutor, ExecutionResult, ResourceLimits
|
||||||
|
from pod_executor.containers.manager import SecureContainerManager
|
||||||
|
|
||||||
|
from ..config.schema import ForgeConfig
|
||||||
|
from ..security.audit import AuditLogger
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class SimpleBackend:
|
||||||
|
"""Adapter for stateless Python code execution using pod_executor.
|
||||||
|
|
||||||
|
This wraps pod_executor.CodeExecutor and adapts it to MCP-Forge's
|
||||||
|
configuration system.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
container_manager: SecureContainerManager,
|
||||||
|
audit_logger: AuditLogger,
|
||||||
|
config: ForgeConfig
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize simple backend adapter.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
container_manager: Container lifecycle manager
|
||||||
|
audit_logger: Audit logging instance
|
||||||
|
config: MCP-Forge configuration
|
||||||
|
"""
|
||||||
|
self.container_manager = container_manager
|
||||||
|
self.audit_logger = audit_logger
|
||||||
|
self.config = config
|
||||||
|
|
||||||
|
# Create default resource limits from config
|
||||||
|
# When not enforcing limits, use very high values to effectively disable
|
||||||
|
if config.security.enforce_resource_limits:
|
||||||
|
self.default_limits = ResourceLimits(
|
||||||
|
memory=config.execution.default_memory,
|
||||||
|
storage="10g",
|
||||||
|
cpu_quota=config.execution.default_cpu_quota,
|
||||||
|
timeout=config.execution.default_timeout
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
# No resource enforcement - use very high limits (effectively unlimited)
|
||||||
|
self.default_limits = ResourceLimits(
|
||||||
|
memory="16g", # Very high memory limit
|
||||||
|
storage="100g", # Very high storage limit
|
||||||
|
cpu_quota=1000000, # Effectively unlimited CPU
|
||||||
|
timeout=config.execution.default_timeout
|
||||||
|
)
|
||||||
|
|
||||||
|
# Create executor with default image (prefer python_3_12)
|
||||||
|
self.executor = CodeExecutor(
|
||||||
|
container_manager=container_manager,
|
||||||
|
image=config.images.python_3_12,
|
||||||
|
resource_limits=self.default_limits
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.debug(f"SimpleBackend initialized with image={config.images.python_3_12}")
|
||||||
|
|
||||||
|
def execute(
|
||||||
|
self,
|
||||||
|
code: str,
|
||||||
|
timeout: Optional[int] = None,
|
||||||
|
memory: Optional[str] = None,
|
||||||
|
cpu_quota: Optional[int] = None,
|
||||||
|
injection_code: Optional[str] = None,
|
||||||
|
bridge_socket_path: Optional[str] = None
|
||||||
|
) -> ExecutionResult:
|
||||||
|
"""
|
||||||
|
Execute Python code in isolated container.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
code: Python code to execute
|
||||||
|
timeout: Optional timeout override (seconds)
|
||||||
|
memory: Optional memory limit override (e.g., "512m")
|
||||||
|
cpu_quota: Optional CPU quota override
|
||||||
|
injection_code: Optional code to inject before user code (for MCP tools)
|
||||||
|
bridge_socket_path: Optional path to MCP bridge socket
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
ExecutionResult with stdout, stderr, result, etc.
|
||||||
|
"""
|
||||||
|
# Create custom resource limits if any overrides provided
|
||||||
|
if memory or cpu_quota:
|
||||||
|
limits = ResourceLimits(
|
||||||
|
memory=memory or self.config.execution.default_memory,
|
||||||
|
storage="10g",
|
||||||
|
cpu_quota=cpu_quota or self.config.execution.default_cpu_quota,
|
||||||
|
timeout=timeout or self.config.execution.default_timeout
|
||||||
|
)
|
||||||
|
# Create temporary executor with custom limits
|
||||||
|
executor = CodeExecutor(
|
||||||
|
container_manager=self.container_manager,
|
||||||
|
image=self.config.images.python_3_12,
|
||||||
|
resource_limits=limits
|
||||||
|
)
|
||||||
|
return executor.execute(
|
||||||
|
code=code,
|
||||||
|
timeout=timeout,
|
||||||
|
injection_code=injection_code,
|
||||||
|
bridge_socket_path=bridge_socket_path
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
# Use default executor
|
||||||
|
return self.executor.execute(
|
||||||
|
code=code,
|
||||||
|
timeout=timeout,
|
||||||
|
injection_code=injection_code,
|
||||||
|
bridge_socket_path=bridge_socket_path
|
||||||
|
)
|
||||||
108
src/mcp_forge/adapters/jupyter_adapter.py
Normal file
108
src/mcp_forge/adapters/jupyter_adapter.py
Normal file
|
|
@ -0,0 +1,108 @@
|
||||||
|
"""Jupyter backend adapter for MCP-Forge.
|
||||||
|
|
||||||
|
Wraps pod_executor.JupyterBackend with MCP-Forge configuration.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import Optional, List, Dict, Any
|
||||||
|
import logging
|
||||||
|
|
||||||
|
from pod_executor import JupyterBackend as PodJupyterBackend, ExecutionResult
|
||||||
|
from pod_executor.containers.manager import SecureContainerManager
|
||||||
|
from pod_executor.jupyter.sessions import SessionState, SessionError
|
||||||
|
|
||||||
|
from ..config.schema import ForgeConfig
|
||||||
|
from ..security.audit import AuditLogger
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
# Re-export for compatibility
|
||||||
|
__all__ = ["JupyterBackend", "SessionState", "SessionError"]
|
||||||
|
|
||||||
|
|
||||||
|
class JupyterBackend:
|
||||||
|
"""Adapter for stateful Python code execution using Jupyter kernels.
|
||||||
|
|
||||||
|
This wraps pod_executor.JupyterBackend and adapts it to MCP-Forge's
|
||||||
|
configuration system.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
container_manager: SecureContainerManager,
|
||||||
|
audit_logger: AuditLogger,
|
||||||
|
config: ForgeConfig
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize Jupyter backend adapter.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
container_manager: Container lifecycle manager
|
||||||
|
audit_logger: Audit logging instance
|
||||||
|
config: MCP-Forge configuration
|
||||||
|
"""
|
||||||
|
self.container_manager = container_manager
|
||||||
|
self.audit_logger = audit_logger
|
||||||
|
self.config = config
|
||||||
|
|
||||||
|
# Create pod_executor backend with config parameters
|
||||||
|
self.backend = PodJupyterBackend(
|
||||||
|
container_manager=container_manager,
|
||||||
|
image=config.images.jupyter,
|
||||||
|
default_timeout=config.execution.default_timeout,
|
||||||
|
default_memory=config.execution.default_memory,
|
||||||
|
default_cpu_quota=config.execution.default_cpu_quota,
|
||||||
|
max_timeout=config.execution.max_timeout,
|
||||||
|
max_memory=config.execution.max_memory,
|
||||||
|
max_cpu_quota=config.execution.max_cpu_quota,
|
||||||
|
max_sessions=config.sessions.max_concurrent,
|
||||||
|
idle_timeout=config.sessions.idle_timeout,
|
||||||
|
audit_logger=audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.debug(f"JupyterBackend initialized with image={config.images.jupyter}")
|
||||||
|
|
||||||
|
def execute(
|
||||||
|
self,
|
||||||
|
code: str,
|
||||||
|
session_id: str,
|
||||||
|
timeout: Optional[int] = None,
|
||||||
|
memory: Optional[str] = None,
|
||||||
|
cpu_quota: Optional[int] = None
|
||||||
|
) -> ExecutionResult:
|
||||||
|
"""
|
||||||
|
Execute code in a stateful Jupyter session.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
code: Python code to execute
|
||||||
|
session_id: Session identifier
|
||||||
|
timeout: Optional timeout override (seconds)
|
||||||
|
memory: Optional memory limit override
|
||||||
|
cpu_quota: Optional CPU quota override
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
ExecutionResult with stdout, stderr, result, etc.
|
||||||
|
"""
|
||||||
|
return self.backend.execute(
|
||||||
|
code=code,
|
||||||
|
session_id=session_id,
|
||||||
|
timeout=timeout,
|
||||||
|
memory=memory,
|
||||||
|
cpu_quota=cpu_quota
|
||||||
|
)
|
||||||
|
|
||||||
|
def list_sessions(self) -> List[Dict[str, Any]]:
|
||||||
|
"""List all active sessions."""
|
||||||
|
return self.backend.list_sessions()
|
||||||
|
|
||||||
|
def get_session(self, session_id: str) -> Optional[Dict[str, Any]]:
|
||||||
|
"""Get session information."""
|
||||||
|
return self.backend.get_session(session_id)
|
||||||
|
|
||||||
|
def destroy_session(self, session_id: str) -> bool:
|
||||||
|
"""Destroy a session."""
|
||||||
|
return self.backend.destroy_session(session_id)
|
||||||
|
|
||||||
|
def cleanup_idle_sessions(self) -> int:
|
||||||
|
"""Clean up idle sessions."""
|
||||||
|
return self.backend.cleanup_idle_sessions()
|
||||||
|
|
@ -7,8 +7,9 @@ from typing import List, Optional, Dict, Set
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from datetime import datetime, timedelta
|
from datetime import datetime, timedelta
|
||||||
|
|
||||||
|
from pod_executor.containers.client import PodmanClient
|
||||||
|
|
||||||
from mcp_forge.config.schema import EnvironmentBuilderConfig
|
from mcp_forge.config.schema import EnvironmentBuilderConfig
|
||||||
from mcp_forge.podman.client import PodmanClient
|
|
||||||
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||||
from mcp_forge.builder.package_validator import PackageValidator
|
from mcp_forge.builder.package_validator import PackageValidator
|
||||||
from mcp_forge.builder.uv_installer import UVInstaller
|
from mcp_forge.builder.uv_installer import UVInstaller
|
||||||
|
|
|
||||||
|
|
@ -8,10 +8,11 @@ from typing import List, Optional
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
|
|
||||||
from mcp_forge.podman.client import PodmanClient
|
from pod_executor.containers.client import PodmanClient
|
||||||
|
from pod_executor.security.resource_limits import parse_memory_string
|
||||||
|
|
||||||
from mcp_forge.config.schema import EnvironmentBuilderConfig
|
from mcp_forge.config.schema import EnvironmentBuilderConfig
|
||||||
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||||
from mcp_forge.security.resource_limits import parse_memory_string
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
|
|
|
||||||
|
|
@ -1 +0,0 @@
|
||||||
"""Execution backends for running code in containers."""
|
|
||||||
|
|
@ -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"
|
|
||||||
]
|
|
||||||
|
|
@ -6,8 +6,8 @@ import hashlib
|
||||||
from mcp_forge.config.schema import ForgeConfig
|
from mcp_forge.config.schema import ForgeConfig
|
||||||
from mcp_forge.podman.containers import SecureContainerManager
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||||
from mcp_forge.security.resource_limits import ResourceLimits, parse_memory_string
|
from pod_executor.security.resource_limits import ResourceLimits, parse_memory_string
|
||||||
from mcp_forge.execution.simple.executor import ExecutionResult
|
from pod_executor.simple.executor import ExecutionResult
|
||||||
from mcp_forge.execution.jupyter.kernel import JupyterKernelManager
|
from mcp_forge.execution.jupyter.kernel import JupyterKernelManager
|
||||||
from mcp_forge.execution.jupyter.sessions import SessionManager, SessionState, SessionError
|
from mcp_forge.execution.jupyter.sessions import SessionManager, SessionState, SessionError
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -22,8 +22,8 @@ from jupyter_client.blocking.client import BlockingKernelClient
|
||||||
import zmq
|
import zmq
|
||||||
|
|
||||||
from mcp_forge.podman.containers import SecureContainerManager, ContainerConfig
|
from mcp_forge.podman.containers import SecureContainerManager, ContainerConfig
|
||||||
from mcp_forge.security.resource_limits import ResourceLimits
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
from mcp_forge.execution.simple.executor import ExecutionResult
|
from pod_executor.simple.executor import ExecutionResult
|
||||||
|
|
||||||
|
|
||||||
class KernelError(Exception):
|
class KernelError(Exception):
|
||||||
|
|
|
||||||
|
|
@ -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
|
|
||||||
|
|
@ -7,8 +7,8 @@ from datetime import datetime, timedelta
|
||||||
from mcp_forge.execution.jupyter.kernel import JupyterKernelManager
|
from mcp_forge.execution.jupyter.kernel import JupyterKernelManager
|
||||||
from mcp_forge.config.schema import SessionConfig
|
from mcp_forge.config.schema import SessionConfig
|
||||||
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||||
from mcp_forge.security.resource_limits import ResourceLimits
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
from mcp_forge.execution.simple.executor import ExecutionResult
|
from pod_executor.simple.executor import ExecutionResult
|
||||||
|
|
||||||
|
|
||||||
class SessionError(Exception):
|
class SessionError(Exception):
|
||||||
|
|
|
||||||
|
|
@ -1,6 +0,0 @@
|
||||||
"""Simple (stateless) execution backend."""
|
|
||||||
|
|
||||||
from .executor import CodeExecutor, ExecutionResult
|
|
||||||
from .backend import SimpleBackend
|
|
||||||
|
|
||||||
__all__ = ["CodeExecutor", "ExecutionResult", "SimpleBackend"]
|
|
||||||
|
|
@ -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
|
|
||||||
|
|
@ -1,11 +0,0 @@
|
||||||
"""Podman integration module."""
|
|
||||||
|
|
||||||
from .client import PodmanClient, PodmanConnectionError
|
|
||||||
from .containers import ContainerConfig, SecureContainerManager
|
|
||||||
|
|
||||||
__all__ = [
|
|
||||||
"PodmanClient",
|
|
||||||
"PodmanConnectionError",
|
|
||||||
"ContainerConfig",
|
|
||||||
"SecureContainerManager",
|
|
||||||
]
|
|
||||||
|
|
@ -10,10 +10,10 @@ from typing import Optional, Dict, List
|
||||||
from datetime import datetime, timedelta
|
from datetime import datetime, timedelta
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
from mcp_forge.podman.client import PodmanClient
|
from pod_executor.containers.client import PodmanClient
|
||||||
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||||
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||||
from mcp_forge.security.resource_limits import ResourceLimits
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
|
|
||||||
class ContainerConfig:
|
class ContainerConfig:
|
||||||
|
|
|
||||||
|
|
@ -1,11 +1,8 @@
|
||||||
"""Security module for resource validation and enforcement."""
|
"""Security module for resource validation and enforcement."""
|
||||||
|
|
||||||
from .resource_limits import (
|
# Resource limits moved to pod_executor package
|
||||||
parse_memory_string,
|
# Import from pod_executor.security.resource_limits instead
|
||||||
parse_cpu_quota,
|
|
||||||
parse_storage_string,
|
|
||||||
ResourceLimits,
|
|
||||||
)
|
|
||||||
from .allowlist import (
|
from .allowlist import (
|
||||||
SecurityError,
|
SecurityError,
|
||||||
OperationValidator,
|
OperationValidator,
|
||||||
|
|
@ -17,10 +14,6 @@ from .audit import (
|
||||||
)
|
)
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"parse_memory_string",
|
|
||||||
"parse_cpu_quota",
|
|
||||||
"parse_storage_string",
|
|
||||||
"ResourceLimits",
|
|
||||||
"SecurityError",
|
"SecurityError",
|
||||||
"OperationValidator",
|
"OperationValidator",
|
||||||
"AuditEventType",
|
"AuditEventType",
|
||||||
|
|
|
||||||
|
|
@ -25,7 +25,7 @@ class ResourceHandler:
|
||||||
config: ForgeConfig instance for configuration info
|
config: ForgeConfig instance for configuration info
|
||||||
"""
|
"""
|
||||||
self.client_manager = client_manager
|
self.client_manager = client_manager
|
||||||
self.session_manager = session_manager
|
self.jupyter_backend = session_manager
|
||||||
self.environment_builder = environment_builder
|
self.environment_builder = environment_builder
|
||||||
self.config = config
|
self.config = config
|
||||||
|
|
||||||
|
|
@ -88,7 +88,7 @@ class ResourceHandler:
|
||||||
Raises:
|
Raises:
|
||||||
KeyError: If session doesn't exist
|
KeyError: If session doesn't exist
|
||||||
"""
|
"""
|
||||||
state = self.session_manager.get_session_state(session_id)
|
state = self.jupyter_backend.get_session_state(session_id)
|
||||||
content = json.dumps(state.to_dict())
|
content = json.dumps(state.to_dict())
|
||||||
|
|
||||||
return TextResourceContents(
|
return TextResourceContents(
|
||||||
|
|
@ -109,7 +109,7 @@ class ResourceHandler:
|
||||||
Raises:
|
Raises:
|
||||||
KeyError: If session doesn't exist
|
KeyError: If session doesn't exist
|
||||||
"""
|
"""
|
||||||
state = self.session_manager.get_session_state(session_id)
|
state = self.jupyter_backend.get_session_state(session_id)
|
||||||
content = json.dumps({"variables": state.all_variables})
|
content = json.dumps({"variables": state.all_variables})
|
||||||
|
|
||||||
return TextResourceContents(
|
return TextResourceContents(
|
||||||
|
|
@ -120,7 +120,7 @@ class ResourceHandler:
|
||||||
|
|
||||||
async def _handle_sessions_list(self) -> TextResourceContents:
|
async def _handle_sessions_list(self) -> TextResourceContents:
|
||||||
"""Return list of active sessions."""
|
"""Return list of active sessions."""
|
||||||
sessions = self.session_manager.list_sessions()
|
sessions = self.jupyter_backend.list_sessions()
|
||||||
content = json.dumps({"sessions": sessions})
|
content = json.dumps({"sessions": sessions})
|
||||||
|
|
||||||
return TextResourceContents(
|
return TextResourceContents(
|
||||||
|
|
|
||||||
|
|
@ -7,18 +7,18 @@ import logging
|
||||||
|
|
||||||
from fastmcp import FastMCP
|
from fastmcp import FastMCP
|
||||||
|
|
||||||
|
from pod_executor.containers.client import PodmanClient
|
||||||
|
from pod_executor.containers.manager import SecureContainerManager
|
||||||
|
from pod_executor.security.validation import BasicValidator
|
||||||
|
from pod_executor.security.audit import SimpleFileAuditLogger
|
||||||
|
|
||||||
from ..config.schema import ForgeConfig
|
from ..config.schema import ForgeConfig
|
||||||
from ..security.audit import AuditLogger
|
from ..security.audit import AuditLogger
|
||||||
from ..security.allowlist import OperationValidator
|
from ..security.allowlist import OperationValidator
|
||||||
from ..podman.client import PodmanClient
|
|
||||||
from ..podman.containers import SecureContainerManager
|
|
||||||
from ..mcp.manager import MCPClientManager
|
from ..mcp.manager import MCPClientManager
|
||||||
from ..mcp.bridge import ToolBridgeServer
|
from ..mcp.bridge import ToolBridgeServer
|
||||||
from ..mcp.injection import ToolInjectionGenerator
|
from ..mcp.injection import ToolInjectionGenerator
|
||||||
from ..execution.simple.backend import SimpleBackend
|
from ..adapters import SimpleBackend, JupyterBackend
|
||||||
from ..execution.jupyter.backend import JupyterBackend
|
|
||||||
from ..execution.jupyter.kernel import JupyterKernelManager
|
|
||||||
from ..execution.jupyter.sessions import SessionManager
|
|
||||||
from ..builder.environment_builder import EnvironmentBuilder
|
from ..builder.environment_builder import EnvironmentBuilder
|
||||||
from .resources import ResourceHandler
|
from .resources import ResourceHandler
|
||||||
from .tools.execute_python import ExecutePythonTool
|
from .tools.execute_python import ExecutePythonTool
|
||||||
|
|
@ -143,15 +143,30 @@ for record in data:
|
||||||
|
|
||||||
def _init_podman(self) -> None:
|
def _init_podman(self) -> None:
|
||||||
"""Initialize Podman client and container manager."""
|
"""Initialize Podman client and container manager."""
|
||||||
|
# Create pod_executor compatible audit logger from MCP-Forge logger
|
||||||
|
pod_audit_logger = SimpleFileAuditLogger(
|
||||||
|
log_path=self.config.security.audit_log
|
||||||
|
)
|
||||||
|
|
||||||
|
# Create pod_executor validator from MCP-Forge validator
|
||||||
|
# Using BasicValidator with same allowed images
|
||||||
|
pod_validator = BasicValidator(
|
||||||
|
allowed_images=[
|
||||||
|
f"{self.config.images.python_3_12}*",
|
||||||
|
f"{self.config.images.python_3_11}*",
|
||||||
|
f"{self.config.images.jupyter}*"
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
self.podman_client = PodmanClient(
|
self.podman_client = PodmanClient(
|
||||||
socket_path=self.config.server.podman_socket,
|
socket_path=self.config.server.podman_socket,
|
||||||
validator=self.operation_validator,
|
validator=pod_validator,
|
||||||
audit_logger=self.audit_logger
|
audit_logger=pod_audit_logger
|
||||||
)
|
)
|
||||||
self.container_manager = SecureContainerManager(
|
self.container_manager = SecureContainerManager(
|
||||||
podman_client=self.podman_client,
|
podman_client=self.podman_client,
|
||||||
validator=self.operation_validator,
|
validator=pod_validator,
|
||||||
audit_logger=self.audit_logger
|
audit_logger=pod_audit_logger
|
||||||
)
|
)
|
||||||
logger.debug("Podman components initialized")
|
logger.debug("Podman components initialized")
|
||||||
|
|
||||||
|
|
@ -185,38 +200,18 @@ for record in data:
|
||||||
|
|
||||||
def _init_backends(self) -> None:
|
def _init_backends(self) -> None:
|
||||||
"""Initialize execution backends."""
|
"""Initialize execution backends."""
|
||||||
from ..security.resource_limits import ResourceLimits
|
# Simple backend for stateless execution (uses adapter)
|
||||||
|
|
||||||
# Simple backend for stateless execution
|
|
||||||
self.simple_backend = SimpleBackend(
|
self.simple_backend = SimpleBackend(
|
||||||
container_manager=self.container_manager,
|
container_manager=self.container_manager,
|
||||||
audit_logger=self.audit_logger,
|
audit_logger=self.audit_logger,
|
||||||
config=self.config
|
config=self.config
|
||||||
)
|
)
|
||||||
|
|
||||||
# Kernel manager with proper resource limits
|
# Jupyter backend for stateful execution (uses adapter)
|
||||||
try:
|
self.jupyter_backend = JupyterBackend(
|
||||||
resource_limits = ResourceLimits(
|
container_manager=self.container_manager,
|
||||||
memory=self.config.execution.max_memory,
|
|
||||||
timeout=self.config.execution.max_timeout,
|
|
||||||
storage="10g",
|
|
||||||
cpu_quota=100000 # 1 CPU
|
|
||||||
)
|
|
||||||
self.kernel_manager = JupyterKernelManager(
|
|
||||||
container_manager=self.container_manager,
|
|
||||||
image="python:3.11",
|
|
||||||
resource_limits=resource_limits
|
|
||||||
)
|
|
||||||
except Exception:
|
|
||||||
# Use a mock if initialization fails (e.g., in tests)
|
|
||||||
from unittest.mock import Mock
|
|
||||||
self.kernel_manager = Mock()
|
|
||||||
|
|
||||||
# Session manager for stateful execution
|
|
||||||
self.session_manager = SessionManager(
|
|
||||||
kernel_manager=self.kernel_manager,
|
|
||||||
audit_logger=self.audit_logger,
|
audit_logger=self.audit_logger,
|
||||||
config=self.config.sessions
|
config=self.config
|
||||||
)
|
)
|
||||||
|
|
||||||
# Jupyter backend for stateful execution
|
# Jupyter backend for stateful execution
|
||||||
|
|
@ -258,7 +253,7 @@ for record in data:
|
||||||
)
|
)
|
||||||
|
|
||||||
self.document_state_tool = DocumentStateTool(
|
self.document_state_tool = DocumentStateTool(
|
||||||
session_manager=self.session_manager
|
jupyter_backend=self.jupyter_backend
|
||||||
)
|
)
|
||||||
|
|
||||||
self.build_environment_tool = BuildEnvironmentTool(
|
self.build_environment_tool = BuildEnvironmentTool(
|
||||||
|
|
@ -411,7 +406,7 @@ print(f"Found {len(records)} records")
|
||||||
# Create resource handler
|
# Create resource handler
|
||||||
self.resource_handler = ResourceHandler(
|
self.resource_handler = ResourceHandler(
|
||||||
client_manager=self.client_manager,
|
client_manager=self.client_manager,
|
||||||
session_manager=self.session_manager,
|
session_manager=self.jupyter_backend,
|
||||||
environment_builder=self.environment_builder,
|
environment_builder=self.environment_builder,
|
||||||
config=self.config
|
config=self.config
|
||||||
)
|
)
|
||||||
|
|
@ -439,7 +434,7 @@ print(f"Found {len(records)} records")
|
||||||
across execute_python calls when session_id parameter is provided.
|
across execute_python calls when session_id parameter is provided.
|
||||||
"""
|
"""
|
||||||
import json
|
import json
|
||||||
sessions = self.session_manager.list_sessions()
|
sessions = self.jupyter_backend.list_sessions()
|
||||||
return json.dumps(sessions, indent=2)
|
return json.dumps(sessions, indent=2)
|
||||||
|
|
||||||
@self.mcp_server.resource("mcp://forge/environments/list")
|
@self.mcp_server.resource("mcp://forge/environments/list")
|
||||||
|
|
|
||||||
|
|
@ -8,14 +8,14 @@ import json
|
||||||
class DocumentStateTool:
|
class DocumentStateTool:
|
||||||
"""MCP tool for documenting important variables in stateful sessions."""
|
"""MCP tool for documenting important variables in stateful sessions."""
|
||||||
|
|
||||||
def __init__(self, session_manager):
|
def __init__(self, jupyter_backend):
|
||||||
"""
|
"""
|
||||||
Initialize Document State Tool.
|
Initialize Document State Tool.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
session_manager: Session manager for accessing session state
|
jupyter_backend: Jupyter backend for accessing session state
|
||||||
"""
|
"""
|
||||||
self.session_manager = session_manager
|
self.jupyter_backend = jupyter_backend
|
||||||
|
|
||||||
def get_tool_definition(self) -> Tool:
|
def get_tool_definition(self) -> Tool:
|
||||||
"""
|
"""
|
||||||
|
|
@ -78,11 +78,11 @@ class DocumentStateTool:
|
||||||
clear = arguments.get("clear", False)
|
clear = arguments.get("clear", False)
|
||||||
|
|
||||||
# Verify session exists
|
# Verify session exists
|
||||||
if not self.session_manager.session_exists(session_id):
|
if not self.jupyter_backend.session_exists(session_id):
|
||||||
raise ValueError(f"Session '{session_id}' not found")
|
raise ValueError(f"Session '{session_id}' not found")
|
||||||
|
|
||||||
# Document variables
|
# Document variables
|
||||||
result = await self.session_manager.document_variables(
|
result = await self.jupyter_backend.document_variables(
|
||||||
session_id=session_id,
|
session_id=session_id,
|
||||||
variables=variables,
|
variables=variables,
|
||||||
note=note,
|
note=note,
|
||||||
|
|
|
||||||
|
|
@ -6,8 +6,7 @@ import json
|
||||||
import logging
|
import logging
|
||||||
import asyncio
|
import asyncio
|
||||||
|
|
||||||
from ...execution.simple.backend import SimpleBackend
|
from ...adapters import SimpleBackend, JupyterBackend
|
||||||
from ...execution.jupyter.backend import JupyterBackend
|
|
||||||
from ...mcp.manager import MCPClientManager
|
from ...mcp.manager import MCPClientManager
|
||||||
from ...mcp.bridge import ToolBridgeServer
|
from ...mcp.bridge import ToolBridgeServer
|
||||||
from ...mcp.injection import ToolInjectionGenerator
|
from ...mcp.injection import ToolInjectionGenerator
|
||||||
|
|
|
||||||
217
src/pod_executor/README.md
Normal file
217
src/pod_executor/README.md
Normal file
|
|
@ -0,0 +1,217 @@
|
||||||
|
# pod_executor
|
||||||
|
|
||||||
|
Standalone Python code execution in Podman containers with security isolation.
|
||||||
|
|
||||||
|
## Overview
|
||||||
|
|
||||||
|
`pod_executor` provides both stateless and stateful (Jupyter) code execution backends that run Python code in isolated Podman containers. It's designed to be usable standalone or as part of larger systems like MCP-Forge.
|
||||||
|
|
||||||
|
## Features
|
||||||
|
|
||||||
|
- ✅ **Stateless execution**: Each code snippet runs in a fresh container
|
||||||
|
- ✅ **Stateful execution**: Jupyter kernels maintain namespace across executions
|
||||||
|
- ✅ **Security isolation**: Containers with configurable resource limits
|
||||||
|
- ✅ **Protocol-based design**: Pluggable audit logging and validation
|
||||||
|
- ✅ **No external configuration**: All parameters explicit
|
||||||
|
- ✅ **Rootless Podman support**: Works with user-level Podman
|
||||||
|
|
||||||
|
## Installation
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Requires Python 3.11+ and Podman
|
||||||
|
pip install podman jupyter-client pyzmq
|
||||||
|
```
|
||||||
|
|
||||||
|
## Quick Start
|
||||||
|
|
||||||
|
### Simple Stateless Execution
|
||||||
|
|
||||||
|
```python
|
||||||
|
from pathlib import Path
|
||||||
|
from pod_executor import (
|
||||||
|
CodeExecutor,
|
||||||
|
ResourceLimits,
|
||||||
|
SecureContainerManager,
|
||||||
|
PodmanClient,
|
||||||
|
NoOpValidator,
|
||||||
|
NullAuditLogger
|
||||||
|
)
|
||||||
|
|
||||||
|
# Setup Podman client (user socket)
|
||||||
|
client = PodmanClient(
|
||||||
|
socket_path=Path("/run/user/1000/podman/podman.sock"),
|
||||||
|
validator=NoOpValidator(),
|
||||||
|
audit_logger=NullAuditLogger()
|
||||||
|
)
|
||||||
|
|
||||||
|
# Setup container manager
|
||||||
|
container_manager = SecureContainerManager(
|
||||||
|
podman_client=client,
|
||||||
|
validator=NoOpValidator(),
|
||||||
|
audit_logger=NullAuditLogger()
|
||||||
|
)
|
||||||
|
|
||||||
|
# Setup executor with resource limits
|
||||||
|
limits = ResourceLimits(
|
||||||
|
memory="512m",
|
||||||
|
cpu_quota=100000, # 100% of 1 CPU
|
||||||
|
storage="1g",
|
||||||
|
timeout=30
|
||||||
|
)
|
||||||
|
|
||||||
|
executor = CodeExecutor(
|
||||||
|
container_manager=container_manager,
|
||||||
|
image="python:3.12",
|
||||||
|
resource_limits=limits
|
||||||
|
)
|
||||||
|
|
||||||
|
# Execute code
|
||||||
|
result = executor.execute("print('Hello from container!')")
|
||||||
|
print(result.stdout) # "Hello from container!\n"
|
||||||
|
print(result.exit_code) # 0
|
||||||
|
print(result.execution_time) # e.g., 0.523
|
||||||
|
```
|
||||||
|
|
||||||
|
### Stateful Jupyter Execution
|
||||||
|
|
||||||
|
```python
|
||||||
|
from pod_executor import JupyterBackend
|
||||||
|
|
||||||
|
# Setup backend
|
||||||
|
backend = JupyterBackend(
|
||||||
|
container_manager=container_manager,
|
||||||
|
image="mcp-forge/jupyter:latest",
|
||||||
|
default_timeout=300,
|
||||||
|
default_memory="512m",
|
||||||
|
max_sessions=10,
|
||||||
|
idle_timeout=3600
|
||||||
|
)
|
||||||
|
|
||||||
|
# Execute in session - variables persist
|
||||||
|
result1 = backend.execute("x = 42", session_id="my-session")
|
||||||
|
result2 = backend.execute("print(x * 2)", session_id="my-session")
|
||||||
|
print(result2.stdout) # "84\n"
|
||||||
|
|
||||||
|
# List active sessions
|
||||||
|
sessions = backend.list_sessions()
|
||||||
|
|
||||||
|
# Cleanup
|
||||||
|
backend.destroy_session("my-session")
|
||||||
|
```
|
||||||
|
|
||||||
|
## Architecture
|
||||||
|
|
||||||
|
```
|
||||||
|
pod_executor/
|
||||||
|
├── security/ # Security components
|
||||||
|
│ ├── resource_limits.py # Memory, CPU, storage limits
|
||||||
|
│ ├── audit.py # Audit logging protocols
|
||||||
|
│ └── validation.py # Security validation protocols
|
||||||
|
├── containers/ # Container management
|
||||||
|
│ ├── client.py # Podman client wrapper
|
||||||
|
│ └── manager.py # Container lifecycle
|
||||||
|
├── simple/ # Stateless execution
|
||||||
|
│ └── executor.py # CodeExecutor
|
||||||
|
└── jupyter/ # Stateful execution
|
||||||
|
├── backend.py # JupyterBackend
|
||||||
|
├── kernel.py # Kernel management
|
||||||
|
└── sessions.py # Session management
|
||||||
|
```
|
||||||
|
|
||||||
|
## Security
|
||||||
|
|
||||||
|
### Validators
|
||||||
|
|
||||||
|
Three validator implementations:
|
||||||
|
|
||||||
|
1. **NoOpValidator**: No validation (testing only!)
|
||||||
|
2. **BasicValidator**: Minimal checks (image allowlist, forbidden params)
|
||||||
|
3. **Custom**: Implement `OperationValidatorProtocol`
|
||||||
|
|
||||||
|
```python
|
||||||
|
from pod_executor.security import BasicValidator
|
||||||
|
|
||||||
|
validator = BasicValidator(allowed_images=["python:3.12*", "jupyter/*"])
|
||||||
|
```
|
||||||
|
|
||||||
|
### Audit Logging
|
||||||
|
|
||||||
|
Three audit logger implementations:
|
||||||
|
|
||||||
|
1. **NullAuditLogger**: No logging
|
||||||
|
2. **SimpleFileAuditLogger**: JSON Lines file logging
|
||||||
|
3. **Custom**: Implement `AuditLoggerProtocol`
|
||||||
|
|
||||||
|
```python
|
||||||
|
from pod_executor.security import SimpleFileAuditLogger
|
||||||
|
|
||||||
|
logger = SimpleFileAuditLogger(Path("/var/log/executor/audit.log"))
|
||||||
|
```
|
||||||
|
|
||||||
|
## Resource Limits
|
||||||
|
|
||||||
|
Control container resources:
|
||||||
|
|
||||||
|
```python
|
||||||
|
from pod_executor import ResourceLimits
|
||||||
|
|
||||||
|
limits = ResourceLimits(
|
||||||
|
memory="2g", # Memory limit
|
||||||
|
cpu_quota=200000, # CPU quota (200% = 2 CPUs)
|
||||||
|
storage="5g", # Storage limit (tracked, not enforced)
|
||||||
|
timeout=600 # Max execution time in seconds
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
|
## Container Images
|
||||||
|
|
||||||
|
Requires Python-capable container images:
|
||||||
|
|
||||||
|
- **Simple executor**: Any Python image (`python:3.12`, `python:3.11-slim`, etc.)
|
||||||
|
- **Jupyter backend**: Image with `ipykernel` installed
|
||||||
|
|
||||||
|
Build Jupyter image:
|
||||||
|
```dockerfile
|
||||||
|
FROM python:3.12-slim
|
||||||
|
RUN pip install ipykernel==6.29.0
|
||||||
|
```
|
||||||
|
|
||||||
|
## Error Handling
|
||||||
|
|
||||||
|
```python
|
||||||
|
from pod_executor import SecurityError, KernelError, SessionError
|
||||||
|
|
||||||
|
try:
|
||||||
|
result = executor.execute("import os; os.system('bad')")
|
||||||
|
except SecurityError as e:
|
||||||
|
print(f"Security violation: {e}")
|
||||||
|
except KernelError as e:
|
||||||
|
print(f"Kernel error: {e}")
|
||||||
|
```
|
||||||
|
|
||||||
|
## Dependencies
|
||||||
|
|
||||||
|
- **podman** (Python library) - Podman API client
|
||||||
|
- **jupyter-client** - Jupyter kernel protocol (for stateful execution)
|
||||||
|
- **pyzmq** - ZMQ messaging (for Jupyter communication)
|
||||||
|
|
||||||
|
## Limitations
|
||||||
|
|
||||||
|
- Requires Podman (not Docker)
|
||||||
|
- Resource limits may not work in all rootless configurations
|
||||||
|
- Jupyter backend needs host networking for ZMQ communication
|
||||||
|
- No automatic image pulling (images must exist)
|
||||||
|
|
||||||
|
## Development
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Run syntax checks
|
||||||
|
python3 -m py_compile src/pod_executor/**/*.py
|
||||||
|
|
||||||
|
# Test simple execution
|
||||||
|
python3 -c "from pod_executor import CodeExecutor; print('Import OK')"
|
||||||
|
```
|
||||||
|
|
||||||
|
## License
|
||||||
|
|
||||||
|
Part of MCP-Forge project.
|
||||||
82
src/pod_executor/__init__.py
Normal file
82
src/pod_executor/__init__.py
Normal file
|
|
@ -0,0 +1,82 @@
|
||||||
|
"""
|
||||||
|
Pod Executor - Standalone Python code execution in Podman containers.
|
||||||
|
|
||||||
|
Provides stateless and stateful (Jupyter) code execution backends with
|
||||||
|
security isolation via Podman containers.
|
||||||
|
|
||||||
|
Examples:
|
||||||
|
# Simple stateless execution
|
||||||
|
from pod_executor import CodeExecutor, ResourceLimits
|
||||||
|
from pod_executor.containers import SecureContainerManager, PodmanClient
|
||||||
|
from pod_executor.security import NoOpValidator, NullAuditLogger
|
||||||
|
|
||||||
|
client = PodmanClient(socket_path="/run/podman/podman.sock",
|
||||||
|
validator=NoOpValidator(),
|
||||||
|
audit_logger=NullAuditLogger())
|
||||||
|
container_manager = SecureContainerManager(client, NoOpValidator(), NullAuditLogger())
|
||||||
|
limits = ResourceLimits(memory="512m", cpu_quota=100000, storage="1g", timeout=30)
|
||||||
|
executor = CodeExecutor(container_manager, "python:3.12", limits)
|
||||||
|
|
||||||
|
result = executor.execute("print('Hello World')")
|
||||||
|
print(result.stdout)
|
||||||
|
|
||||||
|
# Stateful Jupyter execution
|
||||||
|
from pod_executor import JupyterBackend
|
||||||
|
|
||||||
|
backend = JupyterBackend(
|
||||||
|
container_manager=container_manager,
|
||||||
|
image="jupyter/base-notebook",
|
||||||
|
default_timeout=300
|
||||||
|
)
|
||||||
|
|
||||||
|
result = backend.execute("x = 42", session_id="my-session")
|
||||||
|
result = backend.execute("print(x * 2)", session_id="my-session")
|
||||||
|
"""
|
||||||
|
|
||||||
|
from pod_executor.simple.executor import CodeExecutor, ExecutionResult
|
||||||
|
from pod_executor.jupyter.backend import JupyterBackend
|
||||||
|
from pod_executor.jupyter.kernel import JupyterKernelManager, KernelError
|
||||||
|
from pod_executor.jupyter.sessions import SessionManager, SessionState, SessionError
|
||||||
|
from pod_executor.containers.manager import SecureContainerManager, ContainerConfig
|
||||||
|
from pod_executor.containers.client import PodmanClient, PodmanConnectionError
|
||||||
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
from pod_executor.security.audit import (
|
||||||
|
AuditLoggerProtocol,
|
||||||
|
NullAuditLogger,
|
||||||
|
SimpleFileAuditLogger,
|
||||||
|
)
|
||||||
|
from pod_executor.security.validation import (
|
||||||
|
SecurityError,
|
||||||
|
OperationValidatorProtocol,
|
||||||
|
NoOpValidator,
|
||||||
|
BasicValidator,
|
||||||
|
)
|
||||||
|
|
||||||
|
__version__ = "0.1.0"
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
# Simple execution
|
||||||
|
"CodeExecutor",
|
||||||
|
"ExecutionResult",
|
||||||
|
# Jupyter execution
|
||||||
|
"JupyterBackend",
|
||||||
|
"JupyterKernelManager",
|
||||||
|
"KernelError",
|
||||||
|
"SessionManager",
|
||||||
|
"SessionState",
|
||||||
|
"SessionError",
|
||||||
|
# Container management
|
||||||
|
"SecureContainerManager",
|
||||||
|
"ContainerConfig",
|
||||||
|
"PodmanClient",
|
||||||
|
"PodmanConnectionError",
|
||||||
|
# Security
|
||||||
|
"ResourceLimits",
|
||||||
|
"AuditLoggerProtocol",
|
||||||
|
"NullAuditLogger",
|
||||||
|
"SimpleFileAuditLogger",
|
||||||
|
"SecurityError",
|
||||||
|
"OperationValidatorProtocol",
|
||||||
|
"NoOpValidator",
|
||||||
|
"BasicValidator",
|
||||||
|
]
|
||||||
14
src/pod_executor/containers/__init__.py
Normal file
14
src/pod_executor/containers/__init__.py
Normal file
|
|
@ -0,0 +1,14 @@
|
||||||
|
"""Container management for pod_executor."""
|
||||||
|
|
||||||
|
from pod_executor.containers.client import PodmanClient, PodmanConnectionError
|
||||||
|
from pod_executor.containers.manager import (
|
||||||
|
ContainerConfig,
|
||||||
|
SecureContainerManager,
|
||||||
|
)
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"PodmanClient",
|
||||||
|
"PodmanConnectionError",
|
||||||
|
"ContainerConfig",
|
||||||
|
"SecureContainerManager",
|
||||||
|
]
|
||||||
|
|
@ -10,8 +10,8 @@ from pathlib import Path
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
from podman import PodmanClient as BasePodmanClient
|
from podman import PodmanClient as BasePodmanClient
|
||||||
|
|
||||||
from mcp_forge.security.allowlist import OperationValidator
|
from pod_executor.security.validation import OperationValidatorProtocol
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from pod_executor.security.audit import AuditLoggerProtocol, NullAuditLogger
|
||||||
|
|
||||||
|
|
||||||
class PodmanConnectionError(Exception):
|
class PodmanConnectionError(Exception):
|
||||||
|
|
@ -31,8 +31,8 @@ class PodmanClient:
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
socket_path: Path,
|
socket_path: Path,
|
||||||
validator: OperationValidator,
|
validator: OperationValidatorProtocol,
|
||||||
audit_logger: AuditLogger
|
audit_logger: AuditLoggerProtocol
|
||||||
):
|
):
|
||||||
"""
|
"""
|
||||||
Initialize Podman client wrapper.
|
Initialize Podman client wrapper.
|
||||||
|
|
@ -44,7 +44,7 @@ class PodmanClient:
|
||||||
"""
|
"""
|
||||||
self.socket_path = Path(socket_path)
|
self.socket_path = Path(socket_path)
|
||||||
self.validator = validator
|
self.validator = validator
|
||||||
self.audit_logger = audit_logger
|
self.audit_logger = audit_logger if audit_logger is not None else NullAuditLogger()
|
||||||
self._client: Optional[BasePodmanClient] = None
|
self._client: Optional[BasePodmanClient] = None
|
||||||
|
|
||||||
def connect(self) -> None:
|
def connect(self) -> None:
|
||||||
514
src/pod_executor/containers/manager.py
Normal file
514
src/pod_executor/containers/manager.py
Normal file
|
|
@ -0,0 +1,514 @@
|
||||||
|
"""
|
||||||
|
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 pod_executor.containers.client import PodmanClient
|
||||||
|
from pod_executor.security.validation import OperationValidatorProtocol, SecurityError
|
||||||
|
from pod_executor.security.audit import AuditLoggerProtocol, NullAuditLogger
|
||||||
|
from pod_executor.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
|
||||||
|
# Skip resource limits if using very high values (indicates no enforcement)
|
||||||
|
if self.resource_limits:
|
||||||
|
limit_params = self.resource_limits.to_podman_params()
|
||||||
|
# Only apply limits if they're reasonable (not "no enforcement" markers)
|
||||||
|
# Check if mem_limit looks like an enforcement bypass (>= 16GB)
|
||||||
|
mem_limit = limit_params.get("mem_limit", "0")
|
||||||
|
mem_bytes = int(mem_limit) if mem_limit != "0" else 0
|
||||||
|
if mem_bytes < 16 * 1024 * 1024 * 1024: # Less than 16GB = real limit
|
||||||
|
params.update(limit_params)
|
||||||
|
# Disable swap to avoid cgroup swap.max issues on some systems
|
||||||
|
if "mem_limit" in params:
|
||||||
|
params["memswap_limit"] = -1 # Disable swap
|
||||||
|
|
||||||
|
return params
|
||||||
|
|
||||||
|
|
||||||
|
class SecureContainerManager:
|
||||||
|
"""Manages container lifecycle with security enforcement."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
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: Audit logger for operation logging
|
||||||
|
"""
|
||||||
|
self.podman = podman_client
|
||||||
|
self.validator = validator
|
||||||
|
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,
|
||||||
|
**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(event_type="security.violation", severity="critical",
|
||||||
|
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="container.create",
|
||||||
|
severity="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(event_type="security.violation", severity="critical",
|
||||||
|
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="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,
|
||||||
|
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(event_type="security.violation", severity="critical",
|
||||||
|
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="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,
|
||||||
|
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(event_type="security.violation", severity="critical",
|
||||||
|
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="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,
|
||||||
|
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="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,
|
||||||
|
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="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,
|
||||||
|
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="container.remove",
|
||||||
|
severity="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="container.remove",
|
||||||
|
severity="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="container.remove",
|
||||||
|
severity="error",
|
||||||
|
message=f"Cleanup failed: {e}",
|
||||||
|
details={"error": str(e)}
|
||||||
|
)
|
||||||
|
raise
|
||||||
14
src/pod_executor/jupyter/__init__.py
Normal file
14
src/pod_executor/jupyter/__init__.py
Normal file
|
|
@ -0,0 +1,14 @@
|
||||||
|
"""Jupyter stateful code executor."""
|
||||||
|
|
||||||
|
from pod_executor.jupyter.backend import JupyterBackend
|
||||||
|
from pod_executor.jupyter.kernel import JupyterKernelManager, KernelError
|
||||||
|
from pod_executor.jupyter.sessions import SessionManager, SessionState, SessionError
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"JupyterBackend",
|
||||||
|
"JupyterKernelManager",
|
||||||
|
"KernelError",
|
||||||
|
"SessionManager",
|
||||||
|
"SessionState",
|
||||||
|
"SessionError",
|
||||||
|
]
|
||||||
279
src/pod_executor/jupyter/backend.py
Normal file
279
src/pod_executor/jupyter/backend.py
Normal file
|
|
@ -0,0 +1,279 @@
|
||||||
|
"""Jupyter backend for stateful code execution."""
|
||||||
|
|
||||||
|
from typing import Optional, Dict, List
|
||||||
|
import hashlib
|
||||||
|
|
||||||
|
|
||||||
|
from pod_executor.containers.manager import SecureContainerManager
|
||||||
|
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
|
||||||
|
|
||||||
|
|
||||||
|
class JupyterBackend:
|
||||||
|
"""Stateful code execution backend using Jupyter kernels."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
container_manager: SecureContainerManager,
|
||||||
|
image: str = "mcp-forge/jupyter:latest",
|
||||||
|
default_timeout: int = 300,
|
||||||
|
default_memory: str = "512m",
|
||||||
|
default_cpu_quota: int = 50000,
|
||||||
|
max_timeout: int = 1800,
|
||||||
|
max_memory: str = "2g",
|
||||||
|
max_cpu_quota: int = 100000,
|
||||||
|
max_sessions: int = 10,
|
||||||
|
idle_timeout: int = 3600,
|
||||||
|
audit_logger: Optional[AuditLoggerProtocol] = None
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize Jupyter backend.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
config: Forge configuration
|
||||||
|
container_manager: Container lifecycle manager
|
||||||
|
audit_logger: Audit logging instance
|
||||||
|
"""
|
||||||
|
self.image = image
|
||||||
|
self.default_timeout = default_timeout
|
||||||
|
self.default_memory = default_memory
|
||||||
|
self.default_cpu_quota = default_cpu_quota
|
||||||
|
self.max_timeout = max_timeout
|
||||||
|
self.max_memory = max_memory
|
||||||
|
self.max_cpu_quota = max_cpu_quota
|
||||||
|
self.max_sessions = max_sessions
|
||||||
|
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=self._default_resource_limits()
|
||||||
|
)
|
||||||
|
|
||||||
|
# Initialize session manager
|
||||||
|
self.session_manager = SessionManager(
|
||||||
|
kernel_manager=kernel_manager,
|
||||||
|
idle_timeout=idle_timeout,
|
||||||
|
max_sessions=max_sessions,
|
||||||
|
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.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
|
||||||
|
|
||||||
|
# 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()
|
||||||
|
if self.audit_logger:
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type="execution.request",
|
||||||
|
severity="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
|
||||||
|
"""
|
||||||
|
# Always return resource limits in pod_executor
|
||||||
|
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:
|
||||||
|
raise ValueError(
|
||||||
|
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}"
|
||||||
|
)
|
||||||
631
src/pod_executor/jupyter/kernel.py
Normal file
631
src/pod_executor/jupyter/kernel.py
Normal file
|
|
@ -0,0 +1,631 @@
|
||||||
|
"""
|
||||||
|
Real Jupyter kernel management for stateful execution.
|
||||||
|
|
||||||
|
This module implements proper Jupyter kernel management:
|
||||||
|
- jupyter-client runs on host (MCP-Forge server)
|
||||||
|
- ipykernel runs inside Podman containers
|
||||||
|
- Communication via ZMQ protocol
|
||||||
|
- 1:1 mapping: one container per session, one kernel per container
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import Dict, Optional, List, Any
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from datetime import datetime, timedelta
|
||||||
|
import uuid
|
||||||
|
import json
|
||||||
|
import tempfile
|
||||||
|
import time
|
||||||
|
import socket
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from jupyter_client.blocking.client import BlockingKernelClient
|
||||||
|
import zmq
|
||||||
|
|
||||||
|
from pod_executor.containers.manager import SecureContainerManager, ContainerConfig
|
||||||
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
from pod_executor.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
|
||||||
437
src/pod_executor/jupyter/sessions.py
Normal file
437
src/pod_executor/jupyter/sessions.py
Normal file
|
|
@ -0,0 +1,437 @@
|
||||||
|
"""Session management for stateful execution."""
|
||||||
|
|
||||||
|
from typing import Dict, Optional, List, Any
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
from datetime import datetime, timedelta
|
||||||
|
|
||||||
|
from pod_executor.jupyter.kernel import JupyterKernelManager
|
||||||
|
from pod_executor.security.audit import AuditLoggerProtocol, NullAuditLogger
|
||||||
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
from pod_executor.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,
|
||||||
|
kernel_manager: JupyterKernelManager,
|
||||||
|
idle_timeout: int = 3600,
|
||||||
|
max_sessions: int = 10,
|
||||||
|
audit_logger: Optional[AuditLoggerProtocol] = None
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize session manager.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
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 (defaults to NullAuditLogger)
|
||||||
|
"""
|
||||||
|
self.idle_timeout = idle_timeout
|
||||||
|
self.max_sessions = max_sessions
|
||||||
|
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(
|
||||||
|
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="session.create",
|
||||||
|
severity="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="session.destroy",
|
||||||
|
severity="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="session.destroy",
|
||||||
|
severity="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.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.max_sessions:
|
||||||
|
raise SessionError(
|
||||||
|
f"Maximum concurrent sessions ({self.max_sessions}) reached"
|
||||||
|
)
|
||||||
33
src/pod_executor/security/__init__.py
Normal file
33
src/pod_executor/security/__init__.py
Normal file
|
|
@ -0,0 +1,33 @@
|
||||||
|
"""Security components for pod_executor."""
|
||||||
|
|
||||||
|
from pod_executor.security.resource_limits import (
|
||||||
|
ResourceLimits,
|
||||||
|
parse_memory_string,
|
||||||
|
parse_cpu_quota,
|
||||||
|
parse_storage_string,
|
||||||
|
)
|
||||||
|
from pod_executor.security.audit import (
|
||||||
|
AuditLoggerProtocol,
|
||||||
|
NullAuditLogger,
|
||||||
|
SimpleFileAuditLogger,
|
||||||
|
)
|
||||||
|
from pod_executor.security.validation import (
|
||||||
|
SecurityError,
|
||||||
|
OperationValidatorProtocol,
|
||||||
|
NoOpValidator,
|
||||||
|
BasicValidator,
|
||||||
|
)
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"ResourceLimits",
|
||||||
|
"parse_memory_string",
|
||||||
|
"parse_cpu_quota",
|
||||||
|
"parse_storage_string",
|
||||||
|
"AuditLoggerProtocol",
|
||||||
|
"NullAuditLogger",
|
||||||
|
"SimpleFileAuditLogger",
|
||||||
|
"SecurityError",
|
||||||
|
"OperationValidatorProtocol",
|
||||||
|
"NoOpValidator",
|
||||||
|
"BasicValidator",
|
||||||
|
]
|
||||||
175
src/pod_executor/security/audit.py
Normal file
175
src/pod_executor/security/audit.py
Normal file
|
|
@ -0,0 +1,175 @@
|
||||||
|
"""
|
||||||
|
Audit logger protocol for pod_executor.
|
||||||
|
|
||||||
|
Provides a protocol (interface) for audit logging that can be implemented
|
||||||
|
by consuming applications. A default no-op implementation is provided.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import Protocol, Any, Optional, Dict
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
class AuditLoggerProtocol(Protocol):
|
||||||
|
"""Protocol for audit logging (optional dependency)."""
|
||||||
|
|
||||||
|
def log(
|
||||||
|
self,
|
||||||
|
event_type: str,
|
||||||
|
severity: str,
|
||||||
|
message: str,
|
||||||
|
session_id: Optional[str] = None,
|
||||||
|
user_id: Optional[str] = None,
|
||||||
|
details: Optional[Dict[str, Any]] = None,
|
||||||
|
error: Optional[str] = None,
|
||||||
|
**kwargs: Any
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Log an audit event.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
event_type: Type of event (e.g., "container.create", "execution.request")
|
||||||
|
severity: Severity level ("info", "warning", "error", "critical")
|
||||||
|
message: Human-readable message describing the event
|
||||||
|
session_id: Optional session ID associated with event
|
||||||
|
user_id: Optional user ID associated with event
|
||||||
|
details: Optional dictionary of additional details
|
||||||
|
error: Optional error message if event represents an error
|
||||||
|
**kwargs: Additional keyword arguments for extensibility
|
||||||
|
"""
|
||||||
|
...
|
||||||
|
|
||||||
|
|
||||||
|
class NullAuditLogger:
|
||||||
|
"""No-op audit logger for standalone usage without audit requirements."""
|
||||||
|
|
||||||
|
def log(
|
||||||
|
self,
|
||||||
|
event_type: str = "",
|
||||||
|
severity: str = "info",
|
||||||
|
message: str = "",
|
||||||
|
session_id: Optional[str] = None,
|
||||||
|
user_id: Optional[str] = None,
|
||||||
|
details: Optional[Dict[str, Any]] = None,
|
||||||
|
error: Optional[str] = None,
|
||||||
|
**kwargs: Any
|
||||||
|
) -> None:
|
||||||
|
"""Do nothing - audit logging disabled."""
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class SimpleFileAuditLogger:
|
||||||
|
"""
|
||||||
|
Simple file-based audit logger for basic use cases.
|
||||||
|
|
||||||
|
Logs events to a JSON Lines file (one JSON object per line).
|
||||||
|
Thread-safe via file locking.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, log_path: Path):
|
||||||
|
"""
|
||||||
|
Initialize file audit logger.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
log_path: Path to audit log file
|
||||||
|
"""
|
||||||
|
self.log_path = Path(log_path)
|
||||||
|
self._ensure_log_file()
|
||||||
|
|
||||||
|
def _ensure_log_file(self) -> None:
|
||||||
|
"""Ensure log file and directory exist."""
|
||||||
|
self.log_path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
if not self.log_path.exists():
|
||||||
|
self.log_path.touch()
|
||||||
|
|
||||||
|
def log(
|
||||||
|
self,
|
||||||
|
event_type: str = "",
|
||||||
|
severity: str = "info",
|
||||||
|
message: str = "",
|
||||||
|
session_id: Optional[str] = None,
|
||||||
|
user_id: Optional[str] = None,
|
||||||
|
details: Optional[Dict[str, Any]] = None,
|
||||||
|
error: Optional[str] = None,
|
||||||
|
**kwargs: Any
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Log an audit event to JSON Lines file.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
event_type: Type of event
|
||||||
|
severity: Severity level
|
||||||
|
message: Human-readable message
|
||||||
|
session_id: Optional session ID
|
||||||
|
user_id: Optional user ID
|
||||||
|
details: Optional details dictionary
|
||||||
|
error: Optional error message
|
||||||
|
**kwargs: Additional fields
|
||||||
|
"""
|
||||||
|
import json
|
||||||
|
from datetime import datetime, timezone
|
||||||
|
|
||||||
|
entry = {
|
||||||
|
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||||
|
"event_type": event_type,
|
||||||
|
"severity": severity,
|
||||||
|
"message": message,
|
||||||
|
}
|
||||||
|
|
||||||
|
if session_id is not None:
|
||||||
|
entry["session_id"] = session_id
|
||||||
|
if user_id is not None:
|
||||||
|
entry["user_id"] = user_id
|
||||||
|
if details is not None:
|
||||||
|
entry["details"] = details
|
||||||
|
if error is not None:
|
||||||
|
entry["error"] = error
|
||||||
|
|
||||||
|
# Add any additional kwargs
|
||||||
|
entry.update(kwargs)
|
||||||
|
|
||||||
|
# Write to file (append mode, file locking via 'a' mode)
|
||||||
|
with open(self.log_path, 'a') as f:
|
||||||
|
f.write(json.dumps(entry) + '\n')
|
||||||
|
|
||||||
|
def log_container_operation(
|
||||||
|
self,
|
||||||
|
operation: str,
|
||||||
|
container_id: str,
|
||||||
|
image: str,
|
||||||
|
session_id: Optional[str] = None,
|
||||||
|
user_id: Optional[str] = None,
|
||||||
|
details: Optional[Dict[str, Any]] = None,
|
||||||
|
error: Optional[str] = None
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Convenience method to log container operations.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
operation: Operation type (create, start, stop, remove)
|
||||||
|
container_id: Container ID
|
||||||
|
image: Container image name
|
||||||
|
session_id: Optional session ID
|
||||||
|
user_id: Optional user ID
|
||||||
|
details: Optional additional details
|
||||||
|
error: Optional error message
|
||||||
|
"""
|
||||||
|
severity = "error" if error else "info"
|
||||||
|
message = f"Container {operation}: {container_id[:12]} (image: {image})"
|
||||||
|
|
||||||
|
op_details = {
|
||||||
|
"operation": operation,
|
||||||
|
"container_id": container_id,
|
||||||
|
"image": image
|
||||||
|
}
|
||||||
|
if details:
|
||||||
|
op_details.update(details)
|
||||||
|
|
||||||
|
self.log(
|
||||||
|
event_type=f"container.{operation}",
|
||||||
|
severity=severity,
|
||||||
|
message=message,
|
||||||
|
session_id=session_id,
|
||||||
|
user_id=user_id,
|
||||||
|
details=op_details,
|
||||||
|
error=error
|
||||||
|
)
|
||||||
232
src/pod_executor/security/validation.py
Normal file
232
src/pod_executor/security/validation.py
Normal file
|
|
@ -0,0 +1,232 @@
|
||||||
|
"""
|
||||||
|
Security validation protocol for pod_executor.
|
||||||
|
|
||||||
|
Provides protocols (interfaces) for security validation that can be implemented
|
||||||
|
by consuming applications. Default implementations are provided.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import Protocol, Set, Optional, Dict, Any
|
||||||
|
|
||||||
|
|
||||||
|
class SecurityError(Exception):
|
||||||
|
"""Raised when security policy is violated."""
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class OperationValidatorProtocol(Protocol):
|
||||||
|
"""Protocol for validating container operations."""
|
||||||
|
|
||||||
|
session_containers: Set[str]
|
||||||
|
|
||||||
|
def validate_container_create(
|
||||||
|
self,
|
||||||
|
image: str,
|
||||||
|
params: Dict[str, Any],
|
||||||
|
session_id: Optional[str] = None
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Validate container creation parameters against security policy.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
image: Container image name
|
||||||
|
params: Container creation parameters
|
||||||
|
session_id: Optional session ID for volume validation
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SecurityError: If any security policy is violated
|
||||||
|
"""
|
||||||
|
...
|
||||||
|
|
||||||
|
def validate_container_start(self, container_id: str) -> None:
|
||||||
|
"""
|
||||||
|
Validate container start.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
container_id: Container ID to start
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SecurityError: If operation is not allowed
|
||||||
|
"""
|
||||||
|
...
|
||||||
|
|
||||||
|
def validate_container_stop(self, container_id: str) -> None:
|
||||||
|
"""
|
||||||
|
Validate container stop.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
container_id: Container ID to stop
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SecurityError: If operation is not allowed
|
||||||
|
"""
|
||||||
|
...
|
||||||
|
|
||||||
|
def validate_operation(self, operation: str, target: str) -> tuple[bool, Optional[str]]:
|
||||||
|
"""
|
||||||
|
Validate generic operation.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
operation: Operation type
|
||||||
|
target: Operation target
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
(allowed, reason) tuple
|
||||||
|
"""
|
||||||
|
...
|
||||||
|
|
||||||
|
def register_session_container(self, container_id: str) -> None:
|
||||||
|
"""Register a container as belonging to a session."""
|
||||||
|
...
|
||||||
|
|
||||||
|
def unregister_session_container(self, container_id: str) -> None:
|
||||||
|
"""Unregister a session container."""
|
||||||
|
...
|
||||||
|
|
||||||
|
|
||||||
|
class NoOpValidator:
|
||||||
|
"""
|
||||||
|
No-op validator that allows all operations.
|
||||||
|
|
||||||
|
WARNING: This validator provides NO SECURITY. Only use for testing
|
||||||
|
or in fully trusted environments.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
"""Initialize with empty session container set."""
|
||||||
|
self.session_containers: Set[str] = set()
|
||||||
|
|
||||||
|
def validate_container_create(
|
||||||
|
self,
|
||||||
|
image: str,
|
||||||
|
params: Dict[str, Any],
|
||||||
|
session_id: Optional[str] = None
|
||||||
|
) -> None:
|
||||||
|
"""Allow all container creations."""
|
||||||
|
pass
|
||||||
|
|
||||||
|
def validate_container_start(self, container_id: str) -> None:
|
||||||
|
"""Allow all container starts."""
|
||||||
|
pass
|
||||||
|
|
||||||
|
def validate_container_stop(self, container_id: str) -> None:
|
||||||
|
"""Allow all container stops."""
|
||||||
|
pass
|
||||||
|
|
||||||
|
def validate_operation(self, operation: str, target: str) -> tuple[bool, Optional[str]]:
|
||||||
|
"""Allow all operations."""
|
||||||
|
return (True, None)
|
||||||
|
|
||||||
|
def register_session_container(self, container_id: str) -> None:
|
||||||
|
"""Track session container."""
|
||||||
|
self.session_containers.add(container_id)
|
||||||
|
|
||||||
|
def unregister_session_container(self, container_id: str) -> None:
|
||||||
|
"""Untrack session container."""
|
||||||
|
self.session_containers.discard(container_id)
|
||||||
|
|
||||||
|
|
||||||
|
class BasicValidator:
|
||||||
|
"""
|
||||||
|
Basic validator with minimal security checks.
|
||||||
|
|
||||||
|
Enforces:
|
||||||
|
- Allowed image patterns
|
||||||
|
- Required security parameters
|
||||||
|
- Forbidden dangerous parameters
|
||||||
|
- Session container tracking
|
||||||
|
"""
|
||||||
|
|
||||||
|
# Allowed container images with wildcard support
|
||||||
|
DEFAULT_ALLOWED_IMAGES = [
|
||||||
|
"python:3.11*",
|
||||||
|
"python:3.12*",
|
||||||
|
"jupyter/*",
|
||||||
|
"mcp-forge/*",
|
||||||
|
]
|
||||||
|
|
||||||
|
# Parameters that are forbidden
|
||||||
|
FORBIDDEN_PARAMS = [
|
||||||
|
"privileged",
|
||||||
|
"cap_add",
|
||||||
|
"devices",
|
||||||
|
"pid_mode",
|
||||||
|
]
|
||||||
|
|
||||||
|
def __init__(self, allowed_images: Optional[list[str]] = None):
|
||||||
|
"""
|
||||||
|
Initialize basic validator.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
allowed_images: List of allowed image patterns (supports wildcards)
|
||||||
|
"""
|
||||||
|
self.allowed_images = allowed_images or self.DEFAULT_ALLOWED_IMAGES
|
||||||
|
self.session_containers: Set[str] = set()
|
||||||
|
|
||||||
|
def validate_container_create(
|
||||||
|
self,
|
||||||
|
image: str,
|
||||||
|
params: Dict[str, Any],
|
||||||
|
session_id: Optional[str] = None
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Validate container creation.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
image: Container image name
|
||||||
|
params: Container creation parameters
|
||||||
|
session_id: Optional session ID
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SecurityError: If validation fails
|
||||||
|
"""
|
||||||
|
# Validate image is allowed
|
||||||
|
if not self._is_image_allowed(image):
|
||||||
|
raise SecurityError(
|
||||||
|
f"Image '{image}' not in allowlist. "
|
||||||
|
f"Allowed patterns: {self.allowed_images}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Check for forbidden parameters
|
||||||
|
for forbidden in self.FORBIDDEN_PARAMS:
|
||||||
|
if forbidden in params:
|
||||||
|
raise SecurityError(
|
||||||
|
f"Forbidden parameter '{forbidden}' in container creation"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Ensure security_opt includes no-new-privileges
|
||||||
|
security_opts = params.get("security_opt", [])
|
||||||
|
if "no-new-privileges" not in security_opts:
|
||||||
|
raise SecurityError(
|
||||||
|
"Container must include security_opt=['no-new-privileges']"
|
||||||
|
)
|
||||||
|
|
||||||
|
def validate_container_start(self, container_id: str) -> None:
|
||||||
|
"""Validate container start."""
|
||||||
|
if container_id not in self.session_containers:
|
||||||
|
raise SecurityError(
|
||||||
|
f"Container {container_id} is not a registered session container"
|
||||||
|
)
|
||||||
|
|
||||||
|
def validate_container_stop(self, container_id: str) -> None:
|
||||||
|
"""Validate container stop."""
|
||||||
|
if container_id not in self.session_containers:
|
||||||
|
raise SecurityError(
|
||||||
|
f"Container {container_id} is not a registered session container"
|
||||||
|
)
|
||||||
|
|
||||||
|
def validate_operation(self, operation: str, target: str) -> tuple[bool, Optional[str]]:
|
||||||
|
"""Validate generic operation."""
|
||||||
|
return (True, None) # Allow by default
|
||||||
|
|
||||||
|
def register_session_container(self, container_id: str) -> None:
|
||||||
|
"""Register session container."""
|
||||||
|
self.session_containers.add(container_id)
|
||||||
|
|
||||||
|
def unregister_session_container(self, container_id: str) -> None:
|
||||||
|
"""Unregister session container."""
|
||||||
|
self.session_containers.discard(container_id)
|
||||||
|
|
||||||
|
def _is_image_allowed(self, image: str) -> bool:
|
||||||
|
"""Check if image matches any allowed pattern."""
|
||||||
|
import fnmatch
|
||||||
|
return any(fnmatch.fnmatch(image, pattern) for pattern in self.allowed_images)
|
||||||
8
src/pod_executor/simple/__init__.py
Normal file
8
src/pod_executor/simple/__init__.py
Normal file
|
|
@ -0,0 +1,8 @@
|
||||||
|
"""Simple stateless code executor."""
|
||||||
|
|
||||||
|
from pod_executor.simple.executor import CodeExecutor, ExecutionResult
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"CodeExecutor",
|
||||||
|
"ExecutionResult",
|
||||||
|
]
|
||||||
|
|
@ -7,8 +7,8 @@ import time
|
||||||
import textwrap
|
import textwrap
|
||||||
import uuid
|
import uuid
|
||||||
|
|
||||||
from mcp_forge.podman.containers import SecureContainerManager, ContainerConfig
|
from pod_executor.containers.manager import SecureContainerManager, ContainerConfig
|
||||||
from mcp_forge.security.resource_limits import ResourceLimits
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
278
test_cli.py
Executable file
278
test_cli.py
Executable file
|
|
@ -0,0 +1,278 @@
|
||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
Simple test CLI for MCP-Forge execution backends.
|
||||||
|
|
||||||
|
This is a minimal testing tool for development.
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
# Simple (stateless) backend:
|
||||||
|
./test_cli.py
|
||||||
|
|
||||||
|
# Jupyter (stateful) backend:
|
||||||
|
./test_cli.py --jupyter
|
||||||
|
|
||||||
|
# One-shot execution:
|
||||||
|
./test_cli.py --execute "print(2 + 2)"
|
||||||
|
"""
|
||||||
|
|
||||||
|
import sys
|
||||||
|
import os
|
||||||
|
import tempfile
|
||||||
|
import argparse
|
||||||
|
import time
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
# Add src to path
|
||||||
|
sys.path.insert(0, str(Path(__file__).parent / "src"))
|
||||||
|
|
||||||
|
from pod_executor.simple.executor import CodeExecutor
|
||||||
|
from mcp_forge.execution.jupyter.kernel import JupyterKernelManager
|
||||||
|
from pod_executor.containers.client import PodmanClient
|
||||||
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
|
|
||||||
|
|
||||||
|
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()
|
||||||
|
|
@ -12,7 +12,7 @@ from mcp_forge.builder.environment_builder import (
|
||||||
from mcp_forge.builder.package_validator import SecurityError
|
from mcp_forge.builder.package_validator import SecurityError
|
||||||
from mcp_forge.builder.image_builder import BuildResult
|
from mcp_forge.builder.image_builder import BuildResult
|
||||||
from mcp_forge.config.schema import EnvironmentBuilderConfig
|
from mcp_forge.config.schema import EnvironmentBuilderConfig
|
||||||
from mcp_forge.podman.client import PodmanClient
|
from pod_executor.containers.client import PodmanClient
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -5,7 +5,7 @@ from pathlib import Path
|
||||||
from unittest.mock import Mock
|
from unittest.mock import Mock
|
||||||
|
|
||||||
from mcp_forge.builder.image_builder import ImageBuilder, BuildResult
|
from mcp_forge.builder.image_builder import ImageBuilder, BuildResult
|
||||||
from mcp_forge.podman.client import PodmanClient
|
from pod_executor.containers.client import PodmanClient
|
||||||
from mcp_forge.config.schema import EnvironmentBuilderConfig
|
from mcp_forge.config.schema import EnvironmentBuilderConfig
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -6,11 +6,11 @@ from pathlib import Path
|
||||||
|
|
||||||
from mcp_forge.execution.jupyter.backend import JupyterBackend
|
from mcp_forge.execution.jupyter.backend import JupyterBackend
|
||||||
from mcp_forge.execution.jupyter.sessions import SessionManager, Session, SessionState
|
from mcp_forge.execution.jupyter.sessions import SessionManager, Session, SessionState
|
||||||
from mcp_forge.execution.simple.executor import ExecutionResult
|
from pod_executor.simple.executor import ExecutionResult
|
||||||
from mcp_forge.config.schema import ForgeConfig, ExecutionConfig, ImageConfig, SessionConfig
|
from mcp_forge.config.schema import ForgeConfig, ExecutionConfig, ImageConfig, SessionConfig
|
||||||
from mcp_forge.podman.containers import SecureContainerManager
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
from mcp_forge.security.resource_limits import ResourceLimits
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
|
|
|
||||||
|
|
@ -11,8 +11,8 @@ from mcp_forge.execution.jupyter.kernel import (
|
||||||
KernelError
|
KernelError
|
||||||
)
|
)
|
||||||
from mcp_forge.podman.containers import SecureContainerManager
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
from mcp_forge.security.resource_limits import ResourceLimits
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
from mcp_forge.execution.simple.executor import ExecutionResult
|
from pod_executor.simple.executor import ExecutionResult
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
|
|
|
||||||
|
|
@ -14,8 +14,8 @@ from mcp_forge.execution.jupyter.sessions import (
|
||||||
from mcp_forge.execution.jupyter.kernel import JupyterKernelManager
|
from mcp_forge.execution.jupyter.kernel import JupyterKernelManager
|
||||||
from mcp_forge.config.schema import SessionConfig
|
from mcp_forge.config.schema import SessionConfig
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
from mcp_forge.security.resource_limits import ResourceLimits
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
from mcp_forge.execution.simple.executor import ExecutionResult
|
from pod_executor.simple.executor import ExecutionResult
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
|
|
|
||||||
|
|
@ -4,12 +4,12 @@ import pytest
|
||||||
from unittest.mock import Mock, MagicMock, patch
|
from unittest.mock import Mock, MagicMock, patch
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
from mcp_forge.execution.simple.backend import SimpleBackend
|
from mcp_forge.adapters import SimpleBackend
|
||||||
from mcp_forge.execution.simple.executor import ExecutionResult
|
from pod_executor.simple.executor import ExecutionResult
|
||||||
from mcp_forge.config.schema import ForgeConfig, ExecutionConfig, ImageConfig
|
from mcp_forge.config.schema import ForgeConfig, ExecutionConfig, ImageConfig
|
||||||
from mcp_forge.podman.containers import SecureContainerManager
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
from mcp_forge.security.resource_limits import ResourceLimits
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
|
|
|
||||||
|
|
@ -4,9 +4,9 @@ import pytest
|
||||||
from unittest.mock import Mock
|
from unittest.mock import Mock
|
||||||
import json
|
import json
|
||||||
|
|
||||||
from mcp_forge.execution.simple.executor import CodeExecutor, ExecutionResult
|
from pod_executor.simple.executor import CodeExecutor, ExecutionResult
|
||||||
from mcp_forge.podman.containers import SecureContainerManager
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
from mcp_forge.security.resource_limits import ResourceLimits
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
|
|
|
||||||
|
|
@ -1,12 +1,70 @@
|
||||||
"""Integration tests for end-to-end execution flows."""
|
"""Integration tests for end-to-end execution flows."""
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
from unittest.mock import Mock, MagicMock, patch
|
||||||
|
|
||||||
|
from pod_executor import CodeExecutor, ExecutionResult, ResourceLimits
|
||||||
|
from pod_executor.containers.manager import SecureContainerManager
|
||||||
|
from mcp_forge.adapters import SimpleBackend, JupyterBackend
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_simple_backend_execution_flow(real_config):
|
async def test_simple_backend_execution_flow(real_config):
|
||||||
"""Test complete flow: code submission → execution → result return."""
|
"""Test complete flow: code submission → execution → result return."""
|
||||||
pytest.skip("TODO: Phase 5.4 - Requires proper container registration and security setup")
|
# Mock the container manager
|
||||||
|
mock_container_manager = Mock(spec=SecureContainerManager)
|
||||||
|
mock_container_manager.create_container.return_value = "test-container-123"
|
||||||
|
mock_container_manager.wait_for_container.return_value = 0
|
||||||
|
mock_container_manager.get_container_logs.return_value = (
|
||||||
|
'{"result": 42, "error": null}',
|
||||||
|
""
|
||||||
|
)
|
||||||
|
|
||||||
|
# Create audit logger mock
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
mock_audit_logger = Mock(spec=AuditLogger)
|
||||||
|
|
||||||
|
# Create simple backend
|
||||||
|
backend = SimpleBackend(
|
||||||
|
container_manager=mock_container_manager,
|
||||||
|
audit_logger=mock_audit_logger,
|
||||||
|
config=real_config
|
||||||
|
)
|
||||||
|
|
||||||
|
# Execute code
|
||||||
|
result = backend.execute("print(21 * 2)")
|
||||||
|
|
||||||
|
# Verify execution worked
|
||||||
|
assert result.success is True
|
||||||
|
assert result.exit_code == 0
|
||||||
|
|
||||||
|
# Verify container lifecycle was called
|
||||||
|
mock_container_manager.create_container.assert_called_once()
|
||||||
|
mock_container_manager.start_container.assert_called_once()
|
||||||
|
mock_container_manager.remove_container.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_jupyter_backend_initialization(real_config):
|
||||||
|
"""Test Jupyter backend can be initialized with adapters."""
|
||||||
|
# Mock the container manager
|
||||||
|
mock_container_manager = Mock(spec=SecureContainerManager)
|
||||||
|
|
||||||
|
# Create audit logger mock
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
mock_audit_logger = Mock(spec=AuditLogger)
|
||||||
|
|
||||||
|
# Create jupyter backend (should not raise)
|
||||||
|
backend = JupyterBackend(
|
||||||
|
container_manager=mock_container_manager,
|
||||||
|
audit_logger=mock_audit_logger,
|
||||||
|
config=real_config
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify backend was created
|
||||||
|
assert backend is not None
|
||||||
|
assert backend.backend is not None
|
||||||
|
assert backend.config == real_config
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
|
|
|
||||||
1
tests/pod_executor/__init__.py
Normal file
1
tests/pod_executor/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
||||||
|
"""Tests for pod_executor package."""
|
||||||
1
tests/pod_executor/security/__init__.py
Normal file
1
tests/pod_executor/security/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
||||||
|
"""Tests for pod_executor.security module."""
|
||||||
279
tests/pod_executor/security/test_resource_limits.py
Normal file
279
tests/pod_executor/security/test_resource_limits.py
Normal file
|
|
@ -0,0 +1,279 @@
|
||||||
|
"""
|
||||||
|
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 tracked internally but not in Podman params."""
|
||||||
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
|
limits = ResourceLimits(
|
||||||
|
memory="512m",
|
||||||
|
storage="1g",
|
||||||
|
cpu_quota=50000,
|
||||||
|
timeout=300
|
||||||
|
)
|
||||||
|
|
||||||
|
# Storage is tracked internally
|
||||||
|
assert limits.storage_bytes == 1073741824
|
||||||
|
|
||||||
|
params = limits.to_podman_params()
|
||||||
|
|
||||||
|
# Storage is NOT included in Podman params as it's not directly supported
|
||||||
|
# by Podman API for runtime limits. It's tracked for monitoring/validation.
|
||||||
|
assert "mem_limit" in params
|
||||||
|
assert "cpu_quota" in params
|
||||||
|
|
||||||
|
|
||||||
|
def test_cpu_quota_explanation():
|
||||||
|
"""Test that CPU quota values have clear meaning."""
|
||||||
|
from pod_executor.security.resource_limits import parse_cpu_quota
|
||||||
|
|
||||||
|
# 100000 = 100% of one CPU core
|
||||||
|
# 50000 = 50% of one CPU core
|
||||||
|
# 200000 = 200% = 2 CPU cores
|
||||||
|
|
||||||
|
assert parse_cpu_quota(50000) == 50000 # 0.5 cores
|
||||||
|
assert parse_cpu_quota(100000) == 100000 # 1 core
|
||||||
|
assert parse_cpu_quota(200000) == 200000 # 2 cores
|
||||||
|
|
||||||
|
|
||||||
|
def test_parse_memory_with_decimal():
|
||||||
|
"""Test parsing memory strings with decimal values."""
|
||||||
|
from pod_executor.security.resource_limits import parse_memory_string
|
||||||
|
|
||||||
|
# Should handle decimals
|
||||||
|
result = parse_memory_string("1.5g")
|
||||||
|
assert result == 1610612736 # 1.5 * 1024 * 1024 * 1024
|
||||||
1
tests/pod_executor/simple/__init__.py
Normal file
1
tests/pod_executor/simple/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
||||||
|
"""Tests for pod_executor.simple module."""
|
||||||
299
tests/pod_executor/simple/test_executor.py
Normal file
299
tests/pod_executor/simple/test_executor.py
Normal file
|
|
@ -0,0 +1,299 @@
|
||||||
|
"""Tests for the pod_executor Code Executor module."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from unittest.mock import Mock
|
||||||
|
import json
|
||||||
|
|
||||||
|
from pod_executor.simple.executor import CodeExecutor, ExecutionResult
|
||||||
|
from pod_executor.containers.manager import SecureContainerManager
|
||||||
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def resource_limits():
|
||||||
|
"""Standard resource limits for testing."""
|
||||||
|
return ResourceLimits(
|
||||||
|
memory="512m",
|
||||||
|
cpu_quota=50000,
|
||||||
|
storage="1g",
|
||||||
|
timeout=30
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_container_manager():
|
||||||
|
"""Mock SecureContainerManager."""
|
||||||
|
manager = Mock(spec=SecureContainerManager)
|
||||||
|
|
||||||
|
# Mock container lifecycle
|
||||||
|
manager.create_container.return_value = "test-container-123"
|
||||||
|
manager.start_container.return_value = None
|
||||||
|
manager.stop_container.return_value = None
|
||||||
|
manager.remove_container.return_value = None
|
||||||
|
manager.wait_for_container.return_value = 0 # exit code
|
||||||
|
manager.get_container_logs.return_value = ("", "") # (stdout, stderr)
|
||||||
|
|
||||||
|
return manager
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def executor(mock_container_manager, resource_limits):
|
||||||
|
"""CodeExecutor instance with mocked dependencies."""
|
||||||
|
return CodeExecutor(
|
||||||
|
container_manager=mock_container_manager,
|
||||||
|
image="mcp-forge/python:3.11",
|
||||||
|
resource_limits=resource_limits
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_simple_python_code_returns_result(executor, mock_container_manager):
|
||||||
|
"""Test executing simple Python code returns the result."""
|
||||||
|
# Mock successful execution with result
|
||||||
|
mock_container_manager.get_container_logs.return_value = (
|
||||||
|
json.dumps({"result": 42, "error": None}),
|
||||||
|
""
|
||||||
|
)
|
||||||
|
|
||||||
|
result = executor.execute("2 + 2")
|
||||||
|
|
||||||
|
assert result.success is True
|
||||||
|
assert result.result == 42
|
||||||
|
assert result.exit_code == 0
|
||||||
|
assert result.error is None
|
||||||
|
|
||||||
|
# Verify container lifecycle
|
||||||
|
mock_container_manager.create_container.assert_called_once()
|
||||||
|
mock_container_manager.start_container.assert_called_once_with("test-container-123")
|
||||||
|
mock_container_manager.wait_for_container.assert_called_once_with("test-container-123", timeout=30)
|
||||||
|
mock_container_manager.remove_container.assert_called_once_with("test-container-123")
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_code_with_stdout_capture(executor, mock_container_manager):
|
||||||
|
"""Test code execution captures stdout."""
|
||||||
|
mock_container_manager.get_container_logs.return_value = (
|
||||||
|
json.dumps({"result": None, "error": None}) + "\n" + "Hello, World!",
|
||||||
|
""
|
||||||
|
)
|
||||||
|
|
||||||
|
result = executor.execute('print("Hello, World!")')
|
||||||
|
|
||||||
|
assert result.success is True
|
||||||
|
assert "Hello, World!" in result.stdout
|
||||||
|
assert result.stderr == ""
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_code_with_stderr_capture(executor, mock_container_manager):
|
||||||
|
"""Test code execution captures stderr."""
|
||||||
|
mock_container_manager.get_container_logs.return_value = (
|
||||||
|
json.dumps({"result": None, "error": None}),
|
||||||
|
"Warning: something happened"
|
||||||
|
)
|
||||||
|
|
||||||
|
result = executor.execute('import sys; print("warning", file=sys.stderr)')
|
||||||
|
|
||||||
|
assert result.success is True
|
||||||
|
assert result.stderr == "Warning: something happened"
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_code_timeout_enforcement(executor, mock_container_manager):
|
||||||
|
"""Test code execution enforces timeout."""
|
||||||
|
# Simulate timeout by having wait_for_container take too long
|
||||||
|
mock_container_manager.wait_for_container.side_effect = TimeoutError("Container exceeded timeout")
|
||||||
|
|
||||||
|
result = executor.execute("import time; time.sleep(60)", timeout=1)
|
||||||
|
|
||||||
|
assert result.success is False
|
||||||
|
assert result.error is not None
|
||||||
|
assert "timeout" in result.error.lower()
|
||||||
|
|
||||||
|
# Verify cleanup still happens
|
||||||
|
mock_container_manager.remove_container.assert_called_once_with("test-container-123")
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_code_with_exception_handling(executor, mock_container_manager):
|
||||||
|
"""Test code execution handles exceptions gracefully."""
|
||||||
|
error_msg = "ZeroDivisionError: division by zero"
|
||||||
|
mock_container_manager.get_container_logs.return_value = (
|
||||||
|
json.dumps({"result": None, "error": error_msg}),
|
||||||
|
""
|
||||||
|
)
|
||||||
|
mock_container_manager.wait_for_container.return_value = 1 # non-zero exit
|
||||||
|
|
||||||
|
result = executor.execute("1 / 0")
|
||||||
|
|
||||||
|
assert result.success is False
|
||||||
|
assert result.error == error_msg
|
||||||
|
assert result.exit_code == 1
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_code_with_syntax_error_returns_clear_error(executor, mock_container_manager):
|
||||||
|
"""Test code with syntax error returns clear error message."""
|
||||||
|
error_msg = "SyntaxError: invalid syntax"
|
||||||
|
mock_container_manager.get_container_logs.return_value = (
|
||||||
|
json.dumps({"result": None, "error": error_msg}),
|
||||||
|
""
|
||||||
|
)
|
||||||
|
mock_container_manager.wait_for_container.return_value = 1
|
||||||
|
|
||||||
|
result = executor.execute("def foo( :")
|
||||||
|
|
||||||
|
assert result.success is False
|
||||||
|
assert "SyntaxError" in result.error
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_code_with_runtime_error_returns_clear_error(executor, mock_container_manager):
|
||||||
|
"""Test code with runtime error returns clear error with traceback."""
|
||||||
|
error_msg = "NameError: name 'undefined_var' is not defined"
|
||||||
|
mock_container_manager.get_container_logs.return_value = (
|
||||||
|
json.dumps({"result": None, "error": error_msg}),
|
||||||
|
""
|
||||||
|
)
|
||||||
|
mock_container_manager.wait_for_container.return_value = 1
|
||||||
|
|
||||||
|
result = executor.execute("print(undefined_var)")
|
||||||
|
|
||||||
|
assert result.success is False
|
||||||
|
assert "NameError" in result.error
|
||||||
|
|
||||||
|
|
||||||
|
def test_result_serialization_json_compatible_types(executor, mock_container_manager):
|
||||||
|
"""Test execution result contains only JSON-serializable data."""
|
||||||
|
mock_container_manager.get_container_logs.return_value = (
|
||||||
|
json.dumps({"result": [1, 2, {"key": "value"}], "error": None}),
|
||||||
|
""
|
||||||
|
)
|
||||||
|
|
||||||
|
result = executor.execute('[1, 2, {"key": "value"}]')
|
||||||
|
|
||||||
|
# Verify result can be serialized to JSON
|
||||||
|
result_dict = result.to_dict()
|
||||||
|
json_str = json.dumps(result_dict)
|
||||||
|
assert json_str is not None
|
||||||
|
|
||||||
|
# Verify result data
|
||||||
|
assert result.result == [1, 2, {"key": "value"}]
|
||||||
|
|
||||||
|
|
||||||
|
def test_large_output_handling(executor, mock_container_manager):
|
||||||
|
"""Test execution handles large output without issues."""
|
||||||
|
large_output = "x" * 10000 # 10KB of output
|
||||||
|
mock_container_manager.get_container_logs.return_value = (
|
||||||
|
json.dumps({"result": None, "error": None}) + "\n" + large_output,
|
||||||
|
""
|
||||||
|
)
|
||||||
|
|
||||||
|
result = executor.execute('print("x" * 10000)')
|
||||||
|
|
||||||
|
assert result.success is True
|
||||||
|
assert len(result.stdout) >= 10000
|
||||||
|
|
||||||
|
|
||||||
|
def test_execution_result_to_dict(resource_limits):
|
||||||
|
"""Test ExecutionResult.to_dict() returns proper dictionary."""
|
||||||
|
result = ExecutionResult(
|
||||||
|
success=True,
|
||||||
|
stdout="output",
|
||||||
|
stderr="",
|
||||||
|
result=42,
|
||||||
|
execution_time=0.5,
|
||||||
|
exit_code=0,
|
||||||
|
error=None
|
||||||
|
)
|
||||||
|
|
||||||
|
result_dict = result.to_dict()
|
||||||
|
|
||||||
|
assert isinstance(result_dict, dict)
|
||||||
|
assert result_dict["success"] is True
|
||||||
|
assert result_dict["stdout"] == "output"
|
||||||
|
assert result_dict["stderr"] == ""
|
||||||
|
assert result_dict["result"] == 42
|
||||||
|
assert result_dict["execution_time"] == 0.5
|
||||||
|
assert result_dict["exit_code"] == 0
|
||||||
|
assert result_dict["error"] is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_prepare_code_wraps_code_properly(executor):
|
||||||
|
"""Test _prepare_code wraps code to capture result."""
|
||||||
|
code = "x = 2 + 2\nx"
|
||||||
|
wrapped = executor._prepare_code(code)
|
||||||
|
|
||||||
|
# Wrapped code should be executable Python
|
||||||
|
assert "import" in wrapped
|
||||||
|
assert "json" in wrapped
|
||||||
|
assert code in wrapped or "2 + 2" in wrapped
|
||||||
|
|
||||||
|
|
||||||
|
def test_parse_output_extracts_result_and_error(executor):
|
||||||
|
"""Test _parse_output correctly extracts result and error from JSON."""
|
||||||
|
# Test successful result
|
||||||
|
stdout = json.dumps({"result": 42, "error": None})
|
||||||
|
result, error = executor._parse_output(stdout)
|
||||||
|
assert result == 42
|
||||||
|
assert error is None
|
||||||
|
|
||||||
|
# Test error
|
||||||
|
stdout = json.dumps({"result": None, "error": "ValueError: invalid"})
|
||||||
|
result, error = executor._parse_output(stdout)
|
||||||
|
assert result is None
|
||||||
|
assert error == "ValueError: invalid"
|
||||||
|
|
||||||
|
|
||||||
|
def test_cleanup_happens_even_on_create_failure(executor, mock_container_manager):
|
||||||
|
"""Test container cleanup happens even if create fails."""
|
||||||
|
mock_container_manager.create_container.side_effect = Exception("Create failed")
|
||||||
|
|
||||||
|
with pytest.raises(Exception, match="Create failed"):
|
||||||
|
executor.execute("print('test')")
|
||||||
|
|
||||||
|
# No container to remove since create failed
|
||||||
|
mock_container_manager.remove_container.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
def test_cleanup_happens_even_on_start_failure(executor, mock_container_manager):
|
||||||
|
"""Test container cleanup happens even if start fails."""
|
||||||
|
mock_container_manager.start_container.side_effect = Exception("Start failed")
|
||||||
|
|
||||||
|
with pytest.raises(Exception, match="Start failed"):
|
||||||
|
executor.execute("print('test')")
|
||||||
|
|
||||||
|
# Container should still be removed
|
||||||
|
mock_container_manager.remove_container.assert_called_once_with("test-container-123")
|
||||||
|
|
||||||
|
|
||||||
|
def test_execution_time_tracking(executor, mock_container_manager):
|
||||||
|
"""Test execution time is tracked accurately."""
|
||||||
|
mock_container_manager.get_container_logs.return_value = (
|
||||||
|
json.dumps({"result": None, "error": None}),
|
||||||
|
""
|
||||||
|
)
|
||||||
|
|
||||||
|
result = executor.execute("pass")
|
||||||
|
|
||||||
|
assert result.execution_time >= 0
|
||||||
|
assert isinstance(result.execution_time, float)
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_with_custom_timeout(executor, mock_container_manager):
|
||||||
|
"""Test execute respects custom timeout parameter."""
|
||||||
|
mock_container_manager.get_container_logs.return_value = (
|
||||||
|
json.dumps({"result": None, "error": None}),
|
||||||
|
""
|
||||||
|
)
|
||||||
|
|
||||||
|
executor.execute("pass", timeout=60)
|
||||||
|
|
||||||
|
# Verify wait was called with custom timeout
|
||||||
|
mock_container_manager.wait_for_container.assert_called_with("test-container-123", timeout=60)
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_uses_default_timeout_from_resource_limits(executor, mock_container_manager):
|
||||||
|
"""Test execute uses default timeout from resource limits when not specified."""
|
||||||
|
mock_container_manager.get_container_logs.return_value = (
|
||||||
|
json.dumps({"result": None, "error": None}),
|
||||||
|
""
|
||||||
|
)
|
||||||
|
|
||||||
|
executor.execute("pass") # No timeout specified
|
||||||
|
|
||||||
|
# Should use resource_limits.timeout (30)
|
||||||
|
mock_container_manager.wait_for_container.assert_called_with("test-container-123", timeout=30)
|
||||||
|
|
@ -15,10 +15,10 @@ from pathlib import Path
|
||||||
def test_create_container_with_valid_params_succeeds():
|
def test_create_container_with_valid_params_succeeds():
|
||||||
"""Test that container creation with valid params succeeds."""
|
"""Test that container creation with valid params succeeds."""
|
||||||
from mcp_forge.podman.containers import ContainerConfig, SecureContainerManager
|
from mcp_forge.podman.containers import ContainerConfig, SecureContainerManager
|
||||||
from mcp_forge.podman.client import PodmanClient
|
from pod_executor.containers.client import PodmanClient
|
||||||
from mcp_forge.security.allowlist import OperationValidator
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
from mcp_forge.security.audit import AuditLogger, AuditEventType
|
from mcp_forge.security.audit import AuditLogger, AuditEventType
|
||||||
from mcp_forge.security.resource_limits import ResourceLimits
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
from mcp_forge.config.schema import SecurityConfig
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
||||||
# Setup mocks
|
# Setup mocks
|
||||||
|
|
@ -51,7 +51,7 @@ def test_create_container_with_valid_params_succeeds():
|
||||||
def test_create_container_with_forbidden_params_raises_security_error():
|
def test_create_container_with_forbidden_params_raises_security_error():
|
||||||
"""Test that forbidden parameters raise SecurityError."""
|
"""Test that forbidden parameters raise SecurityError."""
|
||||||
from mcp_forge.podman.containers import ContainerConfig, SecureContainerManager
|
from mcp_forge.podman.containers import ContainerConfig, SecureContainerManager
|
||||||
from mcp_forge.podman.client import PodmanClient
|
from pod_executor.containers.client import PodmanClient
|
||||||
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
from mcp_forge.config.schema import SecurityConfig
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
@ -80,7 +80,7 @@ def test_create_container_with_forbidden_params_raises_security_error():
|
||||||
def test_create_container_with_invalid_image_raises_security_error():
|
def test_create_container_with_invalid_image_raises_security_error():
|
||||||
"""Test that invalid/disallowed images raise SecurityError."""
|
"""Test that invalid/disallowed images raise SecurityError."""
|
||||||
from mcp_forge.podman.containers import ContainerConfig, SecureContainerManager
|
from mcp_forge.podman.containers import ContainerConfig, SecureContainerManager
|
||||||
from mcp_forge.podman.client import PodmanClient
|
from pod_executor.containers.client import PodmanClient
|
||||||
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
from mcp_forge.config.schema import SecurityConfig
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
@ -126,7 +126,7 @@ def test_create_container_enforces_required_parameters():
|
||||||
def test_resource_limits_are_applied_correctly():
|
def test_resource_limits_are_applied_correctly():
|
||||||
"""Test that resource limits are correctly applied."""
|
"""Test that resource limits are correctly applied."""
|
||||||
from mcp_forge.podman.containers import ContainerConfig
|
from mcp_forge.podman.containers import ContainerConfig
|
||||||
from mcp_forge.security.resource_limits import ResourceLimits
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
limits = ResourceLimits(
|
limits = ResourceLimits(
|
||||||
memory="1g",
|
memory="1g",
|
||||||
|
|
@ -149,7 +149,7 @@ def test_resource_limits_are_applied_correctly():
|
||||||
def test_volume_mounts_are_validated():
|
def test_volume_mounts_are_validated():
|
||||||
"""Test that volume mounts are validated against allowlist."""
|
"""Test that volume mounts are validated against allowlist."""
|
||||||
from mcp_forge.podman.containers import ContainerConfig, SecureContainerManager
|
from mcp_forge.podman.containers import ContainerConfig, SecureContainerManager
|
||||||
from mcp_forge.podman.client import PodmanClient
|
from pod_executor.containers.client import PodmanClient
|
||||||
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
from mcp_forge.config.schema import SecurityConfig
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
@ -190,7 +190,7 @@ def test_volume_mounts_are_validated():
|
||||||
def test_start_container_on_session_container_succeeds():
|
def test_start_container_on_session_container_succeeds():
|
||||||
"""Test that starting a session container succeeds."""
|
"""Test that starting a session container succeeds."""
|
||||||
from mcp_forge.podman.containers import SecureContainerManager
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
from mcp_forge.podman.client import PodmanClient
|
from pod_executor.containers.client import PodmanClient
|
||||||
from mcp_forge.security.allowlist import OperationValidator
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
from mcp_forge.config.schema import SecurityConfig
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
@ -218,7 +218,7 @@ def test_start_container_on_session_container_succeeds():
|
||||||
def test_start_container_on_non_session_container_raises_security_error():
|
def test_start_container_on_non_session_container_raises_security_error():
|
||||||
"""Test that starting a non-session container raises SecurityError."""
|
"""Test that starting a non-session container raises SecurityError."""
|
||||||
from mcp_forge.podman.containers import SecureContainerManager
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
from mcp_forge.podman.client import PodmanClient
|
from pod_executor.containers.client import PodmanClient
|
||||||
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
from mcp_forge.config.schema import SecurityConfig
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
@ -241,7 +241,7 @@ def test_start_container_on_non_session_container_raises_security_error():
|
||||||
def test_stop_container_works():
|
def test_stop_container_works():
|
||||||
"""Test that stopping a container works."""
|
"""Test that stopping a container works."""
|
||||||
from mcp_forge.podman.containers import SecureContainerManager
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
from mcp_forge.podman.client import PodmanClient
|
from pod_executor.containers.client import PodmanClient
|
||||||
from mcp_forge.security.allowlist import OperationValidator
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
from mcp_forge.config.schema import SecurityConfig
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
@ -269,7 +269,7 @@ def test_stop_container_works():
|
||||||
def test_remove_container_works():
|
def test_remove_container_works():
|
||||||
"""Test that removing a container works."""
|
"""Test that removing a container works."""
|
||||||
from mcp_forge.podman.containers import SecureContainerManager
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
from mcp_forge.podman.client import PodmanClient
|
from pod_executor.containers.client import PodmanClient
|
||||||
from mcp_forge.security.allowlist import OperationValidator
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
from mcp_forge.config.schema import SecurityConfig
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
@ -297,7 +297,7 @@ def test_remove_container_works():
|
||||||
def test_cleanup_old_containers():
|
def test_cleanup_old_containers():
|
||||||
"""Test cleanup of old containers."""
|
"""Test cleanup of old containers."""
|
||||||
from mcp_forge.podman.containers import SecureContainerManager
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
from mcp_forge.podman.client import PodmanClient
|
from pod_executor.containers.client import PodmanClient
|
||||||
from mcp_forge.security.allowlist import OperationValidator
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
from mcp_forge.config.schema import SecurityConfig
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
@ -340,7 +340,7 @@ def test_cleanup_old_containers():
|
||||||
def test_get_container_logs():
|
def test_get_container_logs():
|
||||||
"""Test getting container logs."""
|
"""Test getting container logs."""
|
||||||
from mcp_forge.podman.containers import SecureContainerManager
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
from mcp_forge.podman.client import PodmanClient
|
from pod_executor.containers.client import PodmanClient
|
||||||
from mcp_forge.security.allowlist import OperationValidator
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
from mcp_forge.config.schema import SecurityConfig
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
@ -371,7 +371,7 @@ def test_get_container_logs():
|
||||||
def test_wait_for_container():
|
def test_wait_for_container():
|
||||||
"""Test waiting for container to exit."""
|
"""Test waiting for container to exit."""
|
||||||
from mcp_forge.podman.containers import SecureContainerManager
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
from mcp_forge.podman.client import PodmanClient
|
from pod_executor.containers.client import PodmanClient
|
||||||
from mcp_forge.security.allowlist import OperationValidator
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
from mcp_forge.config.schema import SecurityConfig
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
@ -402,7 +402,7 @@ def test_wait_for_container():
|
||||||
def test_container_config_to_podman_params_includes_all_security_settings():
|
def test_container_config_to_podman_params_includes_all_security_settings():
|
||||||
"""Test that ContainerConfig.to_podman_params includes all required settings."""
|
"""Test that ContainerConfig.to_podman_params includes all required settings."""
|
||||||
from mcp_forge.podman.containers import ContainerConfig
|
from mcp_forge.podman.containers import ContainerConfig
|
||||||
from mcp_forge.security.resource_limits import ResourceLimits
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
config = ContainerConfig(
|
config = ContainerConfig(
|
||||||
image="mcp-forge/python:3.11",
|
image="mcp-forge/python:3.11",
|
||||||
|
|
|
||||||
|
|
@ -13,7 +13,7 @@ from unittest.mock import Mock, MagicMock, patch
|
||||||
|
|
||||||
def test_connection_to_podman_socket_succeeds(tmp_path):
|
def test_connection_to_podman_socket_succeeds(tmp_path):
|
||||||
"""Test that connection to Podman socket succeeds."""
|
"""Test that connection to Podman socket succeeds."""
|
||||||
from mcp_forge.podman.client import PodmanClient
|
from pod_executor.containers.client import PodmanClient
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
from mcp_forge.security.allowlist import OperationValidator
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
from mcp_forge.config.schema import SecurityConfig
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
@ -42,7 +42,7 @@ def test_connection_to_podman_socket_succeeds(tmp_path):
|
||||||
|
|
||||||
def test_connection_failure_raises_clear_error(tmp_path):
|
def test_connection_failure_raises_clear_error(tmp_path):
|
||||||
"""Test that connection failure raises clear error."""
|
"""Test that connection failure raises clear error."""
|
||||||
from mcp_forge.podman.client import PodmanClient, PodmanConnectionError
|
from pod_executor.containers.client import PodmanClient, PodmanConnectionError
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
from mcp_forge.security.allowlist import OperationValidator
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
from mcp_forge.config.schema import SecurityConfig
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
@ -66,7 +66,7 @@ def test_connection_failure_raises_clear_error(tmp_path):
|
||||||
|
|
||||||
def test_socket_path_validation(tmp_path):
|
def test_socket_path_validation(tmp_path):
|
||||||
"""Test that socket path is validated before connecting."""
|
"""Test that socket path is validated before connecting."""
|
||||||
from mcp_forge.podman.client import PodmanClient, PodmanConnectionError
|
from pod_executor.containers.client import PodmanClient, PodmanConnectionError
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
from mcp_forge.security.allowlist import OperationValidator
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
from mcp_forge.config.schema import SecurityConfig
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
@ -89,7 +89,7 @@ def test_socket_path_validation(tmp_path):
|
||||||
|
|
||||||
def test_socket_permissions_check(tmp_path):
|
def test_socket_permissions_check(tmp_path):
|
||||||
"""Test that socket permissions are checked."""
|
"""Test that socket permissions are checked."""
|
||||||
from mcp_forge.podman.client import PodmanClient
|
from pod_executor.containers.client import PodmanClient
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
from mcp_forge.security.allowlist import OperationValidator
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
from mcp_forge.config.schema import SecurityConfig
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
@ -120,7 +120,7 @@ def test_socket_permissions_check(tmp_path):
|
||||||
|
|
||||||
def test_api_version_compatibility_check(tmp_path):
|
def test_api_version_compatibility_check(tmp_path):
|
||||||
"""Test that API version is checked."""
|
"""Test that API version is checked."""
|
||||||
from mcp_forge.podman.client import PodmanClient
|
from pod_executor.containers.client import PodmanClient
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
from mcp_forge.security.allowlist import OperationValidator
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
from mcp_forge.config.schema import SecurityConfig
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
@ -154,7 +154,7 @@ def test_api_version_compatibility_check(tmp_path):
|
||||||
|
|
||||||
def test_ping_health_check(tmp_path):
|
def test_ping_health_check(tmp_path):
|
||||||
"""Test that ping/health check works."""
|
"""Test that ping/health check works."""
|
||||||
from mcp_forge.podman.client import PodmanClient
|
from pod_executor.containers.client import PodmanClient
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
from mcp_forge.security.allowlist import OperationValidator
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
from mcp_forge.config.schema import SecurityConfig
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
@ -184,7 +184,7 @@ def test_ping_health_check(tmp_path):
|
||||||
|
|
||||||
def test_lazy_connection(tmp_path):
|
def test_lazy_connection(tmp_path):
|
||||||
"""Test that connection is lazy (only connects when needed)."""
|
"""Test that connection is lazy (only connects when needed)."""
|
||||||
from mcp_forge.podman.client import PodmanClient
|
from pod_executor.containers.client import PodmanClient
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
from mcp_forge.security.allowlist import OperationValidator
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
from mcp_forge.config.schema import SecurityConfig
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
@ -215,7 +215,7 @@ def test_lazy_connection(tmp_path):
|
||||||
|
|
||||||
def test_disconnect_cleanup(tmp_path):
|
def test_disconnect_cleanup(tmp_path):
|
||||||
"""Test that disconnect cleans up properly."""
|
"""Test that disconnect cleans up properly."""
|
||||||
from mcp_forge.podman.client import PodmanClient
|
from pod_executor.containers.client import PodmanClient
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
from mcp_forge.security.allowlist import OperationValidator
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
from mcp_forge.config.schema import SecurityConfig
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
@ -246,7 +246,7 @@ def test_disconnect_cleanup(tmp_path):
|
||||||
|
|
||||||
def test_connection_error_includes_socket_path(tmp_path):
|
def test_connection_error_includes_socket_path(tmp_path):
|
||||||
"""Test that connection errors include the socket path for debugging."""
|
"""Test that connection errors include the socket path for debugging."""
|
||||||
from mcp_forge.podman.client import PodmanClient, PodmanConnectionError
|
from pod_executor.containers.client import PodmanClient, PodmanConnectionError
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
from mcp_forge.security.allowlist import OperationValidator
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
from mcp_forge.config.schema import SecurityConfig
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
@ -270,7 +270,7 @@ def test_connection_error_includes_socket_path(tmp_path):
|
||||||
|
|
||||||
def test_client_property_auto_connects(tmp_path):
|
def test_client_property_auto_connects(tmp_path):
|
||||||
"""Test that accessing client property auto-connects if not connected."""
|
"""Test that accessing client property auto-connects if not connected."""
|
||||||
from mcp_forge.podman.client import PodmanClient
|
from pod_executor.containers.client import PodmanClient
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
from mcp_forge.security.allowlist import OperationValidator
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
from mcp_forge.config.schema import SecurityConfig
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
@ -304,7 +304,7 @@ def test_client_property_auto_connects(tmp_path):
|
||||||
|
|
||||||
def test_validator_and_audit_logger_stored(tmp_path):
|
def test_validator_and_audit_logger_stored(tmp_path):
|
||||||
"""Test that validator and audit logger are stored for later use."""
|
"""Test that validator and audit logger are stored for later use."""
|
||||||
from mcp_forge.podman.client import PodmanClient
|
from pod_executor.containers.client import PodmanClient
|
||||||
from mcp_forge.security.audit import AuditLogger
|
from mcp_forge.security.audit import AuditLogger
|
||||||
from mcp_forge.security.allowlist import OperationValidator
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
from mcp_forge.config.schema import SecurityConfig
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
|
||||||
|
|
@ -10,7 +10,7 @@ import pytest
|
||||||
|
|
||||||
def test_parse_memory_string_megabytes():
|
def test_parse_memory_string_megabytes():
|
||||||
"""Test parsing memory string with megabytes suffix."""
|
"""Test parsing memory string with megabytes suffix."""
|
||||||
from mcp_forge.security.resource_limits import parse_memory_string
|
from pod_executor.security.resource_limits import parse_memory_string
|
||||||
|
|
||||||
result = parse_memory_string("512m")
|
result = parse_memory_string("512m")
|
||||||
assert result == 536870912 # 512 * 1024 * 1024
|
assert result == 536870912 # 512 * 1024 * 1024
|
||||||
|
|
@ -18,7 +18,7 @@ def test_parse_memory_string_megabytes():
|
||||||
|
|
||||||
def test_parse_memory_string_gigabytes():
|
def test_parse_memory_string_gigabytes():
|
||||||
"""Test parsing memory string with gigabytes suffix."""
|
"""Test parsing memory string with gigabytes suffix."""
|
||||||
from mcp_forge.security.resource_limits import parse_memory_string
|
from pod_executor.security.resource_limits import parse_memory_string
|
||||||
|
|
||||||
result = parse_memory_string("2g")
|
result = parse_memory_string("2g")
|
||||||
assert result == 2147483648 # 2 * 1024 * 1024 * 1024
|
assert result == 2147483648 # 2 * 1024 * 1024 * 1024
|
||||||
|
|
@ -26,7 +26,7 @@ def test_parse_memory_string_gigabytes():
|
||||||
|
|
||||||
def test_parse_memory_string_kilobytes():
|
def test_parse_memory_string_kilobytes():
|
||||||
"""Test parsing memory string with kilobytes suffix."""
|
"""Test parsing memory string with kilobytes suffix."""
|
||||||
from mcp_forge.security.resource_limits import parse_memory_string
|
from pod_executor.security.resource_limits import parse_memory_string
|
||||||
|
|
||||||
result = parse_memory_string("1024k")
|
result = parse_memory_string("1024k")
|
||||||
assert result == 1048576 # 1024 * 1024
|
assert result == 1048576 # 1024 * 1024
|
||||||
|
|
@ -34,7 +34,7 @@ def test_parse_memory_string_kilobytes():
|
||||||
|
|
||||||
def test_parse_memory_string_case_insensitive():
|
def test_parse_memory_string_case_insensitive():
|
||||||
"""Test that memory string parsing is case-insensitive."""
|
"""Test that memory string parsing is case-insensitive."""
|
||||||
from mcp_forge.security.resource_limits import parse_memory_string
|
from pod_executor.security.resource_limits import parse_memory_string
|
||||||
|
|
||||||
assert parse_memory_string("512M") == 536870912
|
assert parse_memory_string("512M") == 536870912
|
||||||
assert parse_memory_string("2G") == 2147483648
|
assert parse_memory_string("2G") == 2147483648
|
||||||
|
|
@ -43,7 +43,7 @@ def test_parse_memory_string_case_insensitive():
|
||||||
|
|
||||||
def test_parse_memory_string_invalid_format_raises_value_error():
|
def test_parse_memory_string_invalid_format_raises_value_error():
|
||||||
"""Test that invalid format raises ValueError."""
|
"""Test that invalid format raises ValueError."""
|
||||||
from mcp_forge.security.resource_limits import parse_memory_string
|
from pod_executor.security.resource_limits import parse_memory_string
|
||||||
|
|
||||||
with pytest.raises(ValueError) as exc_info:
|
with pytest.raises(ValueError) as exc_info:
|
||||||
parse_memory_string("invalid")
|
parse_memory_string("invalid")
|
||||||
|
|
@ -58,7 +58,7 @@ def test_parse_memory_string_invalid_format_raises_value_error():
|
||||||
|
|
||||||
def test_parse_memory_string_negative_value_raises_value_error():
|
def test_parse_memory_string_negative_value_raises_value_error():
|
||||||
"""Test that negative values raise ValueError."""
|
"""Test that negative values raise ValueError."""
|
||||||
from mcp_forge.security.resource_limits import parse_memory_string
|
from pod_executor.security.resource_limits import parse_memory_string
|
||||||
|
|
||||||
with pytest.raises(ValueError) as exc_info:
|
with pytest.raises(ValueError) as exc_info:
|
||||||
parse_memory_string("-512m")
|
parse_memory_string("-512m")
|
||||||
|
|
@ -67,7 +67,7 @@ def test_parse_memory_string_negative_value_raises_value_error():
|
||||||
|
|
||||||
def test_parse_memory_string_zero_value_raises_value_error():
|
def test_parse_memory_string_zero_value_raises_value_error():
|
||||||
"""Test that zero value raises ValueError."""
|
"""Test that zero value raises ValueError."""
|
||||||
from mcp_forge.security.resource_limits import parse_memory_string
|
from pod_executor.security.resource_limits import parse_memory_string
|
||||||
|
|
||||||
with pytest.raises(ValueError) as exc_info:
|
with pytest.raises(ValueError) as exc_info:
|
||||||
parse_memory_string("0m")
|
parse_memory_string("0m")
|
||||||
|
|
@ -76,7 +76,7 @@ def test_parse_memory_string_zero_value_raises_value_error():
|
||||||
|
|
||||||
def test_parse_cpu_quota_valid_value():
|
def test_parse_cpu_quota_valid_value():
|
||||||
"""Test that valid CPU quota values are accepted."""
|
"""Test that valid CPU quota values are accepted."""
|
||||||
from mcp_forge.security.resource_limits import parse_cpu_quota
|
from pod_executor.security.resource_limits import parse_cpu_quota
|
||||||
|
|
||||||
result = parse_cpu_quota(50000)
|
result = parse_cpu_quota(50000)
|
||||||
assert result == 50000
|
assert result == 50000
|
||||||
|
|
@ -87,7 +87,7 @@ def test_parse_cpu_quota_valid_value():
|
||||||
|
|
||||||
def test_parse_cpu_quota_max_limit():
|
def test_parse_cpu_quota_max_limit():
|
||||||
"""Test that CPU quota has a reasonable maximum (10 cores)."""
|
"""Test that CPU quota has a reasonable maximum (10 cores)."""
|
||||||
from mcp_forge.security.resource_limits import parse_cpu_quota
|
from pod_executor.security.resource_limits import parse_cpu_quota
|
||||||
|
|
||||||
# Should accept up to 1000000 (10 cores)
|
# Should accept up to 1000000 (10 cores)
|
||||||
result = parse_cpu_quota(1000000)
|
result = parse_cpu_quota(1000000)
|
||||||
|
|
@ -101,7 +101,7 @@ def test_parse_cpu_quota_max_limit():
|
||||||
|
|
||||||
def test_parse_cpu_quota_negative_raises_value_error():
|
def test_parse_cpu_quota_negative_raises_value_error():
|
||||||
"""Test that negative CPU quota raises ValueError."""
|
"""Test that negative CPU quota raises ValueError."""
|
||||||
from mcp_forge.security.resource_limits import parse_cpu_quota
|
from pod_executor.security.resource_limits import parse_cpu_quota
|
||||||
|
|
||||||
with pytest.raises(ValueError) as exc_info:
|
with pytest.raises(ValueError) as exc_info:
|
||||||
parse_cpu_quota(-1)
|
parse_cpu_quota(-1)
|
||||||
|
|
@ -110,7 +110,7 @@ def test_parse_cpu_quota_negative_raises_value_error():
|
||||||
|
|
||||||
def test_parse_cpu_quota_zero_raises_value_error():
|
def test_parse_cpu_quota_zero_raises_value_error():
|
||||||
"""Test that zero CPU quota raises ValueError."""
|
"""Test that zero CPU quota raises ValueError."""
|
||||||
from mcp_forge.security.resource_limits import parse_cpu_quota
|
from pod_executor.security.resource_limits import parse_cpu_quota
|
||||||
|
|
||||||
with pytest.raises(ValueError) as exc_info:
|
with pytest.raises(ValueError) as exc_info:
|
||||||
parse_cpu_quota(0)
|
parse_cpu_quota(0)
|
||||||
|
|
@ -119,7 +119,7 @@ def test_parse_cpu_quota_zero_raises_value_error():
|
||||||
|
|
||||||
def test_parse_storage_string_same_as_memory():
|
def test_parse_storage_string_same_as_memory():
|
||||||
"""Test that storage parsing works the same as memory parsing."""
|
"""Test that storage parsing works the same as memory parsing."""
|
||||||
from mcp_forge.security.resource_limits import parse_storage_string
|
from pod_executor.security.resource_limits import parse_storage_string
|
||||||
|
|
||||||
assert parse_storage_string("1g") == 1073741824
|
assert parse_storage_string("1g") == 1073741824
|
||||||
assert parse_storage_string("512m") == 536870912
|
assert parse_storage_string("512m") == 536870912
|
||||||
|
|
@ -128,7 +128,7 @@ def test_parse_storage_string_same_as_memory():
|
||||||
|
|
||||||
def test_resource_limits_class_initialization():
|
def test_resource_limits_class_initialization():
|
||||||
"""Test ResourceLimits class initializes correctly."""
|
"""Test ResourceLimits class initializes correctly."""
|
||||||
from mcp_forge.security.resource_limits import ResourceLimits
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
limits = ResourceLimits(
|
limits = ResourceLimits(
|
||||||
memory="512m",
|
memory="512m",
|
||||||
|
|
@ -145,7 +145,7 @@ def test_resource_limits_class_initialization():
|
||||||
|
|
||||||
def test_resource_limits_validates_memory():
|
def test_resource_limits_validates_memory():
|
||||||
"""Test that ResourceLimits validates memory string."""
|
"""Test that ResourceLimits validates memory string."""
|
||||||
from mcp_forge.security.resource_limits import ResourceLimits
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
with pytest.raises(ValueError):
|
with pytest.raises(ValueError):
|
||||||
ResourceLimits(
|
ResourceLimits(
|
||||||
|
|
@ -158,7 +158,7 @@ def test_resource_limits_validates_memory():
|
||||||
|
|
||||||
def test_resource_limits_validates_storage():
|
def test_resource_limits_validates_storage():
|
||||||
"""Test that ResourceLimits validates storage string."""
|
"""Test that ResourceLimits validates storage string."""
|
||||||
from mcp_forge.security.resource_limits import ResourceLimits
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
with pytest.raises(ValueError):
|
with pytest.raises(ValueError):
|
||||||
ResourceLimits(
|
ResourceLimits(
|
||||||
|
|
@ -171,7 +171,7 @@ def test_resource_limits_validates_storage():
|
||||||
|
|
||||||
def test_resource_limits_validates_cpu_quota():
|
def test_resource_limits_validates_cpu_quota():
|
||||||
"""Test that ResourceLimits validates CPU quota."""
|
"""Test that ResourceLimits validates CPU quota."""
|
||||||
from mcp_forge.security.resource_limits import ResourceLimits
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
with pytest.raises(ValueError):
|
with pytest.raises(ValueError):
|
||||||
ResourceLimits(
|
ResourceLimits(
|
||||||
|
|
@ -184,7 +184,7 @@ def test_resource_limits_validates_cpu_quota():
|
||||||
|
|
||||||
def test_resource_limits_to_podman_params():
|
def test_resource_limits_to_podman_params():
|
||||||
"""Test conversion to Podman container parameters."""
|
"""Test conversion to Podman container parameters."""
|
||||||
from mcp_forge.security.resource_limits import ResourceLimits
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
limits = ResourceLimits(
|
limits = ResourceLimits(
|
||||||
memory="512m",
|
memory="512m",
|
||||||
|
|
@ -205,7 +205,7 @@ def test_resource_limits_to_podman_params():
|
||||||
|
|
||||||
def test_resource_limits_default_timeout():
|
def test_resource_limits_default_timeout():
|
||||||
"""Test that ResourceLimits has a default timeout."""
|
"""Test that ResourceLimits has a default timeout."""
|
||||||
from mcp_forge.security.resource_limits import ResourceLimits
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
limits = ResourceLimits(
|
limits = ResourceLimits(
|
||||||
memory="512m",
|
memory="512m",
|
||||||
|
|
@ -218,7 +218,7 @@ def test_resource_limits_default_timeout():
|
||||||
|
|
||||||
def test_parse_memory_string_with_spaces():
|
def test_parse_memory_string_with_spaces():
|
||||||
"""Test parsing memory strings that have spaces."""
|
"""Test parsing memory strings that have spaces."""
|
||||||
from mcp_forge.security.resource_limits import parse_memory_string
|
from pod_executor.security.resource_limits import parse_memory_string
|
||||||
|
|
||||||
# Should handle spaces gracefully (strip them)
|
# Should handle spaces gracefully (strip them)
|
||||||
result = parse_memory_string(" 512m ")
|
result = parse_memory_string(" 512m ")
|
||||||
|
|
@ -227,7 +227,7 @@ def test_parse_memory_string_with_spaces():
|
||||||
|
|
||||||
def test_parse_memory_string_bytes_suffix():
|
def test_parse_memory_string_bytes_suffix():
|
||||||
"""Test parsing memory string with bytes suffix (no multiplier)."""
|
"""Test parsing memory string with bytes suffix (no multiplier)."""
|
||||||
from mcp_forge.security.resource_limits import parse_memory_string
|
from pod_executor.security.resource_limits import parse_memory_string
|
||||||
|
|
||||||
# Just a number (bytes) - should this be supported?
|
# Just a number (bytes) - should this be supported?
|
||||||
# Based on architecture, we support k, m, g suffixes
|
# Based on architecture, we support k, m, g suffixes
|
||||||
|
|
@ -238,7 +238,7 @@ def test_parse_memory_string_bytes_suffix():
|
||||||
|
|
||||||
def test_resource_limits_storage_quota_in_podman_params():
|
def test_resource_limits_storage_quota_in_podman_params():
|
||||||
"""Test that storage limits are included in Podman params."""
|
"""Test that storage limits are included in Podman params."""
|
||||||
from mcp_forge.security.resource_limits import ResourceLimits
|
from pod_executor.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
limits = ResourceLimits(
|
limits = ResourceLimits(
|
||||||
memory="512m",
|
memory="512m",
|
||||||
|
|
@ -256,7 +256,7 @@ def test_resource_limits_storage_quota_in_podman_params():
|
||||||
|
|
||||||
def test_cpu_quota_explanation():
|
def test_cpu_quota_explanation():
|
||||||
"""Test that CPU quota values have clear meaning."""
|
"""Test that CPU quota values have clear meaning."""
|
||||||
from mcp_forge.security.resource_limits import parse_cpu_quota
|
from pod_executor.security.resource_limits import parse_cpu_quota
|
||||||
|
|
||||||
# 100000 = 100% of one CPU core
|
# 100000 = 100% of one CPU core
|
||||||
# 50000 = 50% of one CPU core
|
# 50000 = 50% of one CPU core
|
||||||
|
|
@ -269,7 +269,7 @@ def test_cpu_quota_explanation():
|
||||||
|
|
||||||
def test_parse_memory_with_decimal():
|
def test_parse_memory_with_decimal():
|
||||||
"""Test parsing memory strings with decimal values."""
|
"""Test parsing memory strings with decimal values."""
|
||||||
from mcp_forge.security.resource_limits import parse_memory_string
|
from pod_executor.security.resource_limits import parse_memory_string
|
||||||
|
|
||||||
# Should handle decimals
|
# Should handle decimals
|
||||||
result = parse_memory_string("1.5g")
|
result = parse_memory_string("1.5g")
|
||||||
|
|
|
||||||
|
|
@ -6,7 +6,7 @@ from mcp.types import Tool, TextContent
|
||||||
import json
|
import json
|
||||||
|
|
||||||
from mcp_forge.server.tools.execute_python import ExecutePythonTool
|
from mcp_forge.server.tools.execute_python import ExecutePythonTool
|
||||||
from mcp_forge.execution.simple.backend import ExecutionResult
|
from pod_executor.simple.executor import ExecutionResult
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
|
|
|
||||||
132
verify_containers.py
Executable file
132
verify_containers.py
Executable file
|
|
@ -0,0 +1,132 @@
|
||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
Minimal test script to verify container execution works.
|
||||||
|
|
||||||
|
This tests the basic container creation and execution without
|
||||||
|
using the full MCP-Forge stack.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import sys
|
||||||
|
import subprocess
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
# ANSI colors
|
||||||
|
GREEN = "\033[92m"
|
||||||
|
RED = "\033[91m"
|
||||||
|
GRAY = "\033[90m"
|
||||||
|
RESET = "\033[0m"
|
||||||
|
|
||||||
|
|
||||||
|
def test_simple_execution():
|
||||||
|
"""Test simple container execution."""
|
||||||
|
print(f"\n{GRAY}Testing simple (stateless) execution...{RESET}")
|
||||||
|
|
||||||
|
code = "print('Hello from MCP-Forge!')\nprint(2 + 2)"
|
||||||
|
|
||||||
|
cmd = [
|
||||||
|
"podman", "run", "--rm",
|
||||||
|
"--network=none",
|
||||||
|
"mcp-forge/python:3.12",
|
||||||
|
"python3", "-c", code
|
||||||
|
]
|
||||||
|
|
||||||
|
try:
|
||||||
|
result = subprocess.run(
|
||||||
|
cmd,
|
||||||
|
capture_output=True,
|
||||||
|
text=True,
|
||||||
|
timeout=10
|
||||||
|
)
|
||||||
|
|
||||||
|
if result.returncode == 0:
|
||||||
|
print(f"{GREEN}✓ Simple execution works!{RESET}")
|
||||||
|
print(f"{GRAY}Output:{RESET}")
|
||||||
|
print(result.stdout)
|
||||||
|
return True
|
||||||
|
else:
|
||||||
|
print(f"{RED}✗ Simple execution failed{RESET}")
|
||||||
|
print(f"{RED}Error:{RESET}")
|
||||||
|
print(result.stderr)
|
||||||
|
return False
|
||||||
|
|
||||||
|
except subprocess.TimeoutExpired:
|
||||||
|
print(f"{RED}✗ Execution timed out{RESET}")
|
||||||
|
return False
|
||||||
|
except Exception as e:
|
||||||
|
print(f"{RED}✗ Error: {e}{RESET}")
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
def test_jupyter_kernel():
|
||||||
|
"""Test Jupyter kernel container."""
|
||||||
|
print(f"\n{GRAY}Testing Jupyter kernel availability...{RESET}")
|
||||||
|
|
||||||
|
cmd = [
|
||||||
|
"podman", "run", "--rm",
|
||||||
|
"--network=host",
|
||||||
|
"mcp-forge/jupyter:latest",
|
||||||
|
"python3", "-c", "import ipykernel; print(ipykernel.__version__)"
|
||||||
|
]
|
||||||
|
|
||||||
|
try:
|
||||||
|
result = subprocess.run(
|
||||||
|
cmd,
|
||||||
|
capture_output=True,
|
||||||
|
text=True,
|
||||||
|
timeout=10
|
||||||
|
)
|
||||||
|
|
||||||
|
if result.returncode == 0:
|
||||||
|
version = result.stdout.strip()
|
||||||
|
print(f"{GREEN}✓ Jupyter kernel available (ipykernel {version})!{RESET}")
|
||||||
|
return True
|
||||||
|
else:
|
||||||
|
print(f"{RED}✗ Jupyter kernel check failed{RESET}")
|
||||||
|
print(f"{RED}Error:{RESET}")
|
||||||
|
print(result.stderr)
|
||||||
|
return False
|
||||||
|
|
||||||
|
except subprocess.TimeoutExpired:
|
||||||
|
print(f"{RED}✗ Check timed out{RESET}")
|
||||||
|
return False
|
||||||
|
except Exception as e:
|
||||||
|
print(f"{RED}✗ Error: {e}{RESET}")
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
"""Main entry point."""
|
||||||
|
print(f"{GREEN}MCP-Forge Container Tests{RESET}")
|
||||||
|
print(f"{GRAY}{'='*40}{RESET}")
|
||||||
|
|
||||||
|
results = []
|
||||||
|
|
||||||
|
# Test simple execution
|
||||||
|
results.append(("Simple Execution", test_simple_execution()))
|
||||||
|
|
||||||
|
# Test Jupyter kernel
|
||||||
|
results.append(("Jupyter Kernel", test_jupyter_kernel()))
|
||||||
|
|
||||||
|
# Summary
|
||||||
|
print(f"\n{GRAY}{'='*40}{RESET}")
|
||||||
|
print(f"{GREEN}Test Summary:{RESET}\n")
|
||||||
|
|
||||||
|
passed = sum(1 for _, result in results if result)
|
||||||
|
total = len(results)
|
||||||
|
|
||||||
|
for name, result in results:
|
||||||
|
status = f"{GREEN}✓ PASS{RESET}" if result else f"{RED}✗ FAIL{RESET}"
|
||||||
|
print(f" {name}: {status}")
|
||||||
|
|
||||||
|
print(f"\n{GRAY}Passed: {passed}/{total}{RESET}")
|
||||||
|
|
||||||
|
if passed == total:
|
||||||
|
print(f"\n{GREEN}All tests passed!{RESET}")
|
||||||
|
return 0
|
||||||
|
else:
|
||||||
|
print(f"\n{RED}Some tests failed.{RESET}")
|
||||||
|
return 1
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
sys.exit(main())
|
||||||
Loading…
Add table
Add a link
Reference in a new issue