mcp-forge/test_cli.py

#!/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()