Add test CLI tools for executor testing
- simple_test_cli.py: Working CLI for simple (stateless) backend - Supports one-shot execution and interactive REPL - Uses PassthroughValidator and wrappers to bypass security for testing - Skips resource limits to avoid cgroupv2 issues in rootless Podman - test_containers.py: Container verification script (all tests passing) - simple_test_cli_README.md: Documentation for test tools Note: Jupyter backend has connection file timing issues (future work)
This commit is contained in:
Hans Aschauer
parent
db677ae537
commit
8b6b237be9
3 changed files with 648 additions and 0 deletions
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 mcp_forge.execution.simple.executor import ExecutionResult
|
||||
from mcp_forge.execution.jupyter.kernel import JupyterKernelManager
|
||||
from mcp_forge.podman.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 mcp_forge.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()
|
||||
Loading…
Add table
Add a link
Reference in a new issue