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