massgen 0.0.3__py3-none-any.whl → 0.1.0__py3-none-any.whl

massgen/filesystem_manager/_code_execution_server.py

@@ -0,0 +1,545 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Code Execution MCP Server for MassGen
+
+This MCP server provides command line execution capabilities for agents, allowing
+them to run tests, execute scripts, and perform other command-line operations.
+
+Tools provided:
+- execute_command: Execute any command line command with timeout and working directory control
+
+Inspired by AG2's LocalCommandLineCodeExecutor sanitization patterns.
+"""
+
+import argparse
+import os
+import re
+import subprocess
+import sys
+import time
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import fastmcp
+
+# Platform detection
+WIN32 = sys.platform == "win32"
+
+# Docker integration (optional)
+try:
+    import docker
+    from docker.errors import DockerException
+
+    DOCKER_AVAILABLE = True
+except ImportError:
+    DOCKER_AVAILABLE = False
+    docker = None  # type: ignore
+    DockerException = Exception  # type: ignore
+
+
+def _validate_path_access(path: Path, allowed_paths: List[Path]) -> None:
+    """
+    Validate that a path is within allowed directories.
+
+    Args:
+        path: Path to validate
+        allowed_paths: List of allowed base paths
+
+    Raises:
+        ValueError: If path is not within allowed directories
+    """
+    if not allowed_paths:
+        return  # No restrictions
+
+    for allowed_path in allowed_paths:
+        try:
+            path.relative_to(allowed_path)
+            return  # Path is within this allowed directory
+        except ValueError:
+            continue
+
+    raise ValueError(f"Path not in allowed directories: {path}")
+
+
+def _sanitize_command(command: str) -> None:
+    """
+    Sanitize the command to prevent dangerous operations.
+
+    Adapted from AG2's LocalCommandLineCodeExecutor.sanitize_command().
+    This provides basic protection for users running commands outside Docker.
+
+    Args:
+        command: The command to sanitize
+
+    Raises:
+        ValueError: If dangerous command is detected
+    """
+    dangerous_patterns = [
+        # AG2 original patterns
+        (r"\brm\s+-rf\s+/", "Use of 'rm -rf /' is not allowed"),
+        (r"\bmv\b.*?\s+/dev/null", "Moving files to /dev/null is not allowed"),
+        (r"\bdd\b", "Use of 'dd' command is not allowed"),
+        (r">\s*/dev/sd[a-z][1-9]?", "Overwriting disk blocks directly is not allowed"),
+        (r":\(\)\{\s*:\|\:&\s*\};:", "Fork bombs are not allowed"),
+        # Additional safety patterns
+        (r"\bsudo\b", "Use of 'sudo' is not allowed"),
+        (r"\bsu\b", "Use of 'su' is not allowed"),
+        (r"\bchown\b", "Use of 'chown' is not allowed"),
+        (r"\bchmod\b", "Use of 'chmod' is not allowed"),
+    ]
+
+    for pattern, message in dangerous_patterns:
+        if re.search(pattern, command):
+            raise ValueError(f"Potentially dangerous command detected: {message}")
+
+
+def _check_command_filters(command: str, allowed_patterns: Optional[List[str]], blocked_patterns: Optional[List[str]]) -> None:
+    """
+    Check command against whitelist/blacklist filters.
+
+    Args:
+        command: The command to check
+        allowed_patterns: Whitelist regex patterns (if provided, command MUST match one)
+        blocked_patterns: Blacklist regex patterns (command must NOT match any)
+
+    Raises:
+        ValueError: If command doesn't match whitelist or matches blacklist
+    """
+    # Check whitelist (if provided, command MUST match at least one pattern)
+    if allowed_patterns:
+        if not any(re.match(pattern, command) for pattern in allowed_patterns):
+            raise ValueError(
+                f"Command not in allowed list. Allowed patterns: {', '.join(allowed_patterns)}",
+            )
+
+    # Check blacklist (command must NOT match any blocked pattern)
+    if blocked_patterns:
+        for pattern in blocked_patterns:
+            if re.match(pattern, command):
+                raise ValueError(
+                    f"Command matches blocked pattern: '{pattern}'",
+                )
+
+
+def _prepare_environment(work_dir: Path) -> Dict[str, str]:
+    """
+    Prepare environment by auto-detecting .venv in work_dir.
+
+    This function checks for a .venv directory in the working directory and
+    automatically modifies PATH to use it if found. Each workspace manages
+    its own virtual environment independently.
+
+    Args:
+        work_dir: Working directory to check for .venv
+
+    Returns:
+        Environment variables dict with PATH modified if .venv exists
+    """
+    env = os.environ.copy()
+
+    # Auto-detect .venv in work_dir
+    venv_dir = work_dir / ".venv"
+    if venv_dir.exists():
+        # Determine bin directory based on platform
+        venv_bin = venv_dir / ("Scripts" if WIN32 else "bin")
+        if venv_bin.exists():
+            # Prepend venv bin to PATH
+            env["PATH"] = f"{venv_bin}{os.pathsep}{env['PATH']}"
+            # Set VIRTUAL_ENV for tools that check it
+            env["VIRTUAL_ENV"] = str(venv_dir)
+
+    return env
+
+
+async def create_server() -> fastmcp.FastMCP:
+    """Factory function to create and configure the code execution server."""
+
+    parser = argparse.ArgumentParser(description="Code Execution MCP Server")
+    parser.add_argument(
+        "--allowed-paths",
+        type=str,
+        nargs="*",
+        default=[],
+        help="List of allowed base paths for execution (default: no restrictions)",
+    )
+    parser.add_argument(
+        "--timeout",
+        type=int,
+        default=60,
+        help="Default timeout in seconds (default: 60)",
+    )
+    parser.add_argument(
+        "--max-output-size",
+        type=int,
+        default=1024 * 1024,  # 1MB
+        help="Maximum output size in bytes (default: 1MB)",
+    )
+    parser.add_argument(
+        "--allowed-commands",
+        type=str,
+        nargs="*",
+        default=None,
+        help="Whitelist: Only allow commands matching these regex patterns (e.g., 'python .*', 'pytest .*')",
+    )
+    parser.add_argument(
+        "--blocked-commands",
+        type=str,
+        nargs="*",
+        default=None,
+        help="Blacklist: Block commands matching these regex patterns (e.g., 'rm .*', 'sudo .*')",
+    )
+    parser.add_argument(
+        "--execution-mode",
+        type=str,
+        default="local",
+        choices=["local", "docker"],
+        help="Execution mode: local (subprocess) or docker (container isolation)",
+    )
+    parser.add_argument(
+        "--agent-id",
+        type=str,
+        default=None,
+        help="Agent ID (required for Docker mode to identify container)",
+    )
+    args = parser.parse_args()
+
+    # Create the FastMCP server
+    mcp = fastmcp.FastMCP("Command Execution")
+
+    # Store configuration
+    mcp.allowed_paths = [Path(p).resolve() for p in args.allowed_paths]
+    mcp.default_timeout = args.timeout
+    mcp.max_output_size = args.max_output_size
+    mcp.allowed_commands = args.allowed_commands  # Whitelist patterns
+    mcp.blocked_commands = args.blocked_commands  # Blacklist patterns
+    mcp.execution_mode = args.execution_mode
+    mcp.agent_id = args.agent_id
+
+    # Initialize Docker client if Docker mode
+    mcp.docker_client = None
+    if args.execution_mode == "docker":
+        if not DOCKER_AVAILABLE:
+            raise RuntimeError("Docker mode requested but docker library not available. Install with: pip install docker")
+        # Note: agent_id validation is deferred to first command execution
+        # This allows MCP server to start before agent_id is set by orchestrator
+
+        try:
+            mcp.docker_client = docker.from_env()
+            mcp.docker_client.ping()  # Test connection
+            print("✅ [Docker] Connected to Docker daemon")
+        except DockerException as e:
+            raise RuntimeError(f"Failed to connect to Docker: {e}")
+
+    @mcp.tool()
+    def execute_command(
+        command: str,
+        timeout: Optional[int] = None,
+        work_dir: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Execute a command line command.
+
+        This tool allows executing any command line program including:
+        - Python: execute_command("python script.py")
+        - Node.js: execute_command("node app.js")
+        - Tests: execute_command("pytest tests/")
+        - Build tools: execute_command("npm run build")
+        - Shell commands: execute_command("ls -la")
+
+        The command is executed in a shell environment, so you can use shell features
+        like pipes, redirection, and environment variables. On Windows, this uses
+        cmd.exe; on Unix/Mac, this uses the default shell (typically bash).
+
+        Args:
+            command: The command to execute (required)
+            timeout: Maximum execution time in seconds (default: 60)
+                    Set to None for no timeout (use with caution)
+            work_dir: Working directory for execution (relative to workspace)
+                     If not specified, uses the current workspace directory
+
+        Returns:
+            Dictionary containing:
+            - success: bool - True if exit code was 0
+            - exit_code: int - Process exit code
+            - stdout: str - Standard output from the command
+            - stderr: str - Standard error from the command
+            - execution_time: float - Time taken to execute in seconds
+            - command: str - The command that was executed
+            - work_dir: str - The working directory used
+
+        Security:
+            - Execution is confined to allowed paths
+            - Timeout enforced to prevent infinite loops
+            - Output size limited to prevent memory exhaustion
+            - Basic sanitization against dangerous commands
+
+        Examples:
+            # Run Python script
+            execute_command("python test.py")
+
+            # Run tests with pytest
+            execute_command("pytest tests/ -v")
+
+            # Install package and run script
+            execute_command("pip install requests && python scraper.py")
+
+            # Check Python version
+            execute_command("python --version")
+
+            # List files
+            execute_command("ls -la")  # Unix/Mac
+            execute_command("dir")      # Windows
+        """
+        try:
+            # Basic command sanitization (dangerous patterns)
+            try:
+                _sanitize_command(command)
+            except ValueError as e:
+                return {
+                    "success": False,
+                    "exit_code": -1,
+                    "stdout": "",
+                    "stderr": str(e),
+                    "execution_time": 0.0,
+                    "command": command,
+                    "work_dir": work_dir or str(Path.cwd()),
+                }
+
+            # Check whitelist/blacklist filters
+            try:
+                _check_command_filters(command, mcp.allowed_commands, mcp.blocked_commands)
+            except ValueError as e:
+                return {
+                    "success": False,
+                    "exit_code": -1,
+                    "stdout": "",
+                    "stderr": str(e),
+                    "execution_time": 0.0,
+                    "command": command,
+                    "work_dir": work_dir or str(Path.cwd()),
+                }
+
+            # Use default timeout if not specified
+            if timeout is None:
+                timeout = mcp.default_timeout
+
+            # Resolve working directory
+            if work_dir:
+                if Path(work_dir).is_absolute():
+                    work_path = Path(work_dir).resolve()
+                else:
+                    # Relative path - resolve relative to current working directory
+                    work_path = (Path.cwd() / work_dir).resolve()
+            else:
+                work_path = Path.cwd()
+
+            # Validate working directory is within allowed paths
+            _validate_path_access(work_path, mcp.allowed_paths)
+
+            # Verify working directory exists
+            if not work_path.exists():
+                return {
+                    "success": False,
+                    "exit_code": -1,
+                    "stdout": "",
+                    "stderr": f"Working directory does not exist: {work_path}",
+                    "execution_time": 0.0,
+                    "command": command,
+                    "work_dir": str(work_path),
+                }
+
+            if not work_path.is_dir():
+                return {
+                    "success": False,
+                    "exit_code": -1,
+                    "stdout": "",
+                    "stderr": f"Working directory is not a directory: {work_path}",
+                    "execution_time": 0.0,
+                    "command": command,
+                    "work_dir": str(work_path),
+                }
+
+            # Execute command based on execution mode
+            if mcp.execution_mode == "docker":
+                # Docker mode: execute in container via Docker client
+                if not mcp.docker_client:
+                    return {
+                        "success": False,
+                        "exit_code": -1,
+                        "stdout": "",
+                        "stderr": "Docker mode enabled but docker_client not initialized",
+                        "execution_time": 0.0,
+                        "command": command,
+                        "work_dir": str(work_path),
+                    }
+
+                # Validate agent_id is set before executing
+                if not mcp.agent_id:
+                    return {
+                        "success": False,
+                        "exit_code": -1,
+                        "stdout": "",
+                        "stderr": "Docker mode requires agent_id to be set. This should be configured by the orchestrator.",
+                        "execution_time": 0.0,
+                        "command": command,
+                        "work_dir": str(work_path),
+                    }
+
+                try:
+                    # Get container by name
+                    container_name = f"massgen-{mcp.agent_id}"
+                    container = mcp.docker_client.containers.get(container_name)
+
+                    # IMPORTANT: Use host paths directly in container
+                    # Container mounts are configured to use the SAME paths as host
+                    # This makes Docker completely transparent to the LLM
+
+                    # Execute command via docker exec
+                    exec_config = {
+                        "cmd": ["/bin/sh", "-c", command],
+                        "workdir": str(work_path),  # Use host path directly
+                        "stdout": True,
+                        "stderr": True,
+                    }
+
+                    start_time = time.time()
+                    exit_code, output = container.exec_run(**exec_config)
+                    execution_time = time.time() - start_time
+
+                    # Docker exec_run combines stdout and stderr
+                    output_str = output.decode("utf-8") if isinstance(output, bytes) else output
+
+                    # Truncate output if too large
+                    if len(output_str) > mcp.max_output_size:
+                        output_str = output_str[: mcp.max_output_size] + f"\n... (truncated, exceeded {mcp.max_output_size} bytes)"
+
+                    return {
+                        "success": exit_code == 0,
+                        "exit_code": exit_code,
+                        "stdout": output_str,
+                        "stderr": "",  # Docker exec_run combines stdout/stderr
+                        "execution_time": execution_time,
+                        "command": command,
+                        "work_dir": str(work_path),  # Return host path
+                    }
+
+                except DockerException as e:
+                    return {
+                        "success": False,
+                        "exit_code": -1,
+                        "stdout": "",
+                        "stderr": f"Docker container error: {str(e)}",
+                        "execution_time": 0.0,
+                        "command": command,
+                        "work_dir": str(work_path),
+                    }
+                except Exception as e:
+                    return {
+                        "success": False,
+                        "exit_code": -1,
+                        "stdout": "",
+                        "stderr": f"Docker execution error: {str(e)}",
+                        "execution_time": 0.0,
+                        "command": command,
+                        "work_dir": str(work_path),
+                    }
+
+            else:
+                # Local mode: execute using subprocess (existing logic)
+                # Prepare environment (auto-detects .venv in work_dir)
+                env = _prepare_environment(work_path)
+
+                # Execute command
+                start_time = time.time()
+
+                try:
+                    result = subprocess.run(
+                        command,
+                        shell=True,
+                        cwd=str(work_path),
+                        timeout=timeout,
+                        capture_output=True,
+                        text=True,
+                        env=env,
+                    )
+
+                    execution_time = time.time() - start_time
+
+                    # Truncate output if too large
+                    stdout = result.stdout
+                    stderr = result.stderr
+
+                    if len(stdout) > mcp.max_output_size:
+                        stdout = stdout[: mcp.max_output_size] + f"\n... (truncated, exceeded {mcp.max_output_size} bytes)"
+
+                    if len(stderr) > mcp.max_output_size:
+                        stderr = stderr[: mcp.max_output_size] + f"\n... (truncated, exceeded {mcp.max_output_size} bytes)"
+
+                    return {
+                        "success": result.returncode == 0,
+                        "exit_code": result.returncode,
+                        "stdout": stdout,
+                        "stderr": stderr,
+                        "execution_time": execution_time,
+                        "command": command,
+                        "work_dir": str(work_path),
+                    }
+
+                except subprocess.TimeoutExpired:
+                    execution_time = time.time() - start_time
+                    return {
+                        "success": False,
+                        "exit_code": -1,
+                        "stdout": "",
+                        "stderr": f"Command timed out after {timeout} seconds",
+                        "execution_time": execution_time,
+                        "command": command,
+                        "work_dir": str(work_path),
+                    }
+
+                except Exception as e:
+                    execution_time = time.time() - start_time
+                    return {
+                        "success": False,
+                        "exit_code": -1,
+                        "stdout": "",
+                        "stderr": f"Execution error: {str(e)}",
+                        "execution_time": execution_time,
+                        "command": command,
+                        "work_dir": str(work_path),
+                    }
+
+        except ValueError as e:
+            # Path validation error
+            return {
+                "success": False,
+                "exit_code": -1,
+                "stdout": "",
+                "stderr": f"Path validation error: {str(e)}",
+                "execution_time": 0.0,
+                "command": command,
+                "work_dir": work_dir or str(Path.cwd()),
+            }
+
+        except Exception as e:
+            # Unexpected error
+            return {
+                "success": False,
+                "exit_code": -1,
+                "stdout": "",
+                "stderr": f"Unexpected error: {str(e)}",
+                "execution_time": 0.0,
+                "command": command,
+                "work_dir": work_dir or str(Path.cwd()),
+            }
+
+    print("🚀 Command Execution MCP Server started and ready")
+    print(f"Execution mode: {mcp.execution_mode}")
+    if mcp.execution_mode == "docker":
+        print(f"Agent ID: {mcp.agent_id}")
+    print(f"Default timeout: {mcp.default_timeout}s")
+    print(f"Max output size: {mcp.max_output_size} bytes")
+    print(f"Allowed paths: {[str(p) for p in mcp.allowed_paths]}")
+
+    return mcp