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

massgen/filesystem_manager/_docker_manager.py

@@ -0,0 +1,477 @@
+# -*- coding: utf-8 -*-
+"""
+Docker Container Manager for MassGen
+
+Manages Docker containers for isolated command execution.
+Provides strong filesystem isolation by executing commands inside containers
+while keeping MCP servers on the host.
+"""
+
+import logging
+import threading
+import time
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+# Check if docker is available
+try:
+    import docker
+    from docker.errors import DockerException, ImageNotFound, NotFound
+    from docker.models.containers import Container
+
+    DOCKER_AVAILABLE = True
+except ImportError:
+    DOCKER_AVAILABLE = False
+    logger.warning("Docker Python library not available. Install with: pip install docker")
+
+
+class DockerManager:
+    """
+    Manages Docker containers for isolated command execution.
+
+    Each agent gets a persistent container for the orchestration session:
+    - Volume mounts for workspace and context paths
+    - Network isolation (configurable)
+    - Resource limits (CPU, memory)
+    - Commands executed via docker exec
+    - State persists across turns (packages stay installed)
+    """
+
+    def __init__(
+        self,
+        image: str = "massgen/mcp-runtime:latest",
+        network_mode: str = "none",
+        memory_limit: Optional[str] = None,
+        cpu_limit: Optional[float] = None,
+    ):
+        """
+        Initialize Docker manager.
+
+        Args:
+            image: Docker image to use for containers
+            network_mode: Network mode (none/bridge/host)
+            memory_limit: Memory limit (e.g., "2g", "512m")
+            cpu_limit: CPU limit (e.g., 2.0 for 2 CPUs)
+
+        Raises:
+            RuntimeError: If Docker is not available or cannot connect
+        """
+        if not DOCKER_AVAILABLE:
+            raise RuntimeError("Docker Python library not available. Install with: pip install docker")
+
+        self.image = image
+        self.network_mode = network_mode
+        self.memory_limit = memory_limit
+        self.cpu_limit = cpu_limit
+
+        try:
+            self.client = docker.from_env()
+            # Test connection
+            self.client.ping()
+
+            # Get Docker version info for logging
+            version_info = self.client.version()
+            docker_version = version_info.get("Version", "unknown")
+            api_version = version_info.get("ApiVersion", "unknown")
+
+            logger.info("🐳 [Docker] Client initialized successfully")
+            logger.info(f"    Docker version: {docker_version}")
+            logger.info(f"    API version: {api_version}")
+        except DockerException as e:
+            logger.error(f"❌ [Docker] Failed to connect to Docker daemon: {e}")
+            raise RuntimeError(f"Failed to connect to Docker: {e}")
+
+        self.containers: Dict[str, Container] = {}  # agent_id -> container
+
+    def ensure_image_exists(self) -> None:
+        """
+        Ensure the Docker image exists locally.
+
+        Pulls the image if not found locally.
+
+        Raises:
+            RuntimeError: If image cannot be pulled
+        """
+        try:
+            self.client.images.get(self.image)
+            logger.info(f"✅ [Docker] Image '{self.image}' found locally")
+        except ImageNotFound:
+            logger.info(f"📥 [Docker] Image '{self.image}' not found locally, pulling...")
+            try:
+                self.client.images.pull(self.image)
+                logger.info(f"✅ [Docker] Successfully pulled image '{self.image}'")
+            except DockerException as e:
+                raise RuntimeError(f"Failed to pull Docker image '{self.image}': {e}")
+
+    def create_container(
+        self,
+        agent_id: str,
+        workspace_path: Path,
+        temp_workspace_path: Optional[Path] = None,
+        context_paths: Optional[List[Dict[str, Any]]] = None,
+    ) -> str:
+        """
+        Create and start a persistent Docker container for an agent.
+
+        The container runs for the entire orchestration session and maintains state
+        across command executions (installed packages, generated files, etc.).
+
+        IMPORTANT: Paths are mounted at the SAME location as on the host to maintain
+        path transparency. The LLM sees identical paths whether in Docker or local mode.
+
+        Args:
+            agent_id: Unique identifier for the agent
+            workspace_path: Path to agent's workspace (mounted at same path, read-write)
+            temp_workspace_path: Path to shared temp workspace (mounted at same path, read-only)
+            context_paths: List of context path dicts with 'path', 'permission', and optional 'name' keys
+                          (each mounted at its host path)
+
+        Returns:
+            Container ID
+
+        Raises:
+            RuntimeError: If container creation fails
+        """
+        if agent_id in self.containers:
+            logger.warning(f"⚠️ [Docker] Container for agent {agent_id} already exists")
+            return self.containers[agent_id].id
+
+        # Ensure image exists
+        self.ensure_image_exists()
+
+        # Check for and remove any existing container with the same name
+        container_name = f"massgen-{agent_id}"
+        try:
+            existing = self.client.containers.get(container_name)
+            logger.warning(
+                f"🔄 [Docker] Found existing container '{container_name}' (id: {existing.short_id}), removing it",
+            )
+            existing.remove(force=True)
+        except NotFound:
+            # No existing container, this is expected
+            pass
+        except DockerException as e:
+            logger.warning(f"⚠️ [Docker] Error checking for existing container '{container_name}': {e}")
+
+        logger.info(f"🐳 [Docker] Creating container for agent '{agent_id}'")
+        logger.info(f"    Image: {self.image}")
+        logger.info(f"    Network: {self.network_mode}")
+        if self.memory_limit:
+            logger.info(f"    Memory limit: {self.memory_limit}")
+        if self.cpu_limit:
+            logger.info(f"    CPU limit: {self.cpu_limit} cores")
+
+        # Build volume mounts
+        # IMPORTANT: Mount paths at the SAME location as on host to avoid path confusion
+        # This makes Docker completely transparent to the LLM - it sees identical paths
+        volumes = {}
+        mount_info = []
+
+        # Mount agent workspace (read-write) at the SAME path as host
+        workspace_path = workspace_path.resolve()
+        volumes[str(workspace_path)] = {"bind": str(workspace_path), "mode": "rw"}
+        mount_info.append(f"      {workspace_path} ← {workspace_path} (rw)")
+
+        # Mount temp workspace (read-only) at the SAME path as host
+        if temp_workspace_path:
+            temp_workspace_path = temp_workspace_path.resolve()
+            volumes[str(temp_workspace_path)] = {"bind": str(temp_workspace_path), "mode": "ro"}
+            mount_info.append(f"      {temp_workspace_path} ← {temp_workspace_path} (ro)")
+
+        # Mount context paths at the SAME paths as host
+        if context_paths:
+            for ctx_path_config in context_paths:
+                ctx_path = Path(ctx_path_config["path"]).resolve()
+                permission = ctx_path_config.get("permission", "read")
+                mode = "rw" if permission == "write" else "ro"
+
+                volumes[str(ctx_path)] = {"bind": str(ctx_path), "mode": mode}
+                mount_info.append(f"      {ctx_path} ← {ctx_path} ({mode})")
+
+        # Log volume mounts
+        if mount_info:
+            logger.info("    Volume mounts:")
+            for mount_line in mount_info:
+                logger.info(mount_line)
+
+        # Build resource limits
+        resource_config = {}
+        if self.memory_limit:
+            resource_config["mem_limit"] = self.memory_limit
+        if self.cpu_limit:
+            resource_config["nano_cpus"] = int(self.cpu_limit * 1e9)
+
+        # Container configuration
+        container_config = {
+            "image": self.image,
+            "name": container_name,
+            "command": ["tail", "-f", "/dev/null"],  # Keep container running
+            "detach": True,
+            "volumes": volumes,
+            "working_dir": str(workspace_path),  # Use host workspace path
+            "network_mode": self.network_mode,
+            "auto_remove": False,  # Manual cleanup for better control
+            "stdin_open": True,
+            "tty": True,
+            **resource_config,
+        }
+
+        try:
+            # Create and start container
+            container = self.client.containers.run(**container_config)
+            self.containers[agent_id] = container
+
+            # Get container info for logging
+            container.reload()  # Refresh container state
+            status = container.status
+
+            logger.info("✅ [Docker] Container created successfully")
+            logger.info(f"    Container ID: {container.short_id}")
+            logger.info(f"    Container name: {container_name}")
+            logger.info(f"    Status: {status}")
+
+            # Show how to inspect the container
+            logger.debug(f"💡 [Docker] Inspect container: docker inspect {container.short_id}")
+            logger.debug(f"💡 [Docker] View logs: docker logs {container.short_id}")
+            logger.debug(f"💡 [Docker] Execute commands: docker exec -it {container.short_id} /bin/bash")
+
+            return container.id
+
+        except DockerException as e:
+            logger.error(f"❌ [Docker] Failed to create container for agent {agent_id}: {e}")
+            raise RuntimeError(f"Failed to create Docker container for agent {agent_id}: {e}")
+
+    def get_container(self, agent_id: str) -> Optional[Container]:
+        """
+        Get container for an agent.
+
+        Args:
+            agent_id: Agent identifier
+
+        Returns:
+            Container object or None if not found
+        """
+        return self.containers.get(agent_id)
+
+    def exec_command(
+        self,
+        agent_id: str,
+        command: str,
+        workdir: Optional[str] = None,
+        timeout: Optional[int] = None,
+    ) -> Dict[str, Any]:
+        """
+        Execute a command inside the agent's container.
+
+        Args:
+            agent_id: Agent identifier
+            command: Command to execute (as string, will be run in shell)
+            workdir: Working directory (uses host path - same path is mounted in container)
+            timeout: Command timeout in seconds (implemented using threading)
+
+        Returns:
+            Dictionary with:
+            - success: bool (True if exit_code == 0)
+            - exit_code: int
+            - stdout: str
+            - stderr: str (combined with stdout in Docker exec)
+            - execution_time: float
+            - command: str
+            - work_dir: str
+
+        Raises:
+            ValueError: If container not found
+            RuntimeError: If execution fails
+        """
+        container = self.containers.get(agent_id)
+        if not container:
+            raise ValueError(f"No container found for agent {agent_id}")
+
+        # Default workdir is the container's default working dir (set to workspace_path at creation)
+        effective_workdir = workdir if workdir else None
+
+        try:
+            # Run command through shell to support pipes, redirects, etc.
+            exec_config = {
+                "cmd": ["/bin/sh", "-c", command],
+                "stdout": True,
+                "stderr": True,
+            }
+
+            if effective_workdir:
+                exec_config["workdir"] = effective_workdir
+
+            logger.debug(f"🔧 [Docker] Executing in container {container.short_id}: {command}")
+
+            start_time = time.time()
+
+            # Handle timeout using threading
+            if timeout:
+                result_container = {}
+                exception_container = {}
+
+                def run_exec():
+                    try:
+                        result_container["data"] = container.exec_run(**exec_config)
+                    except Exception as e:
+                        exception_container["error"] = e
+
+                thread = threading.Thread(target=run_exec)
+                thread.daemon = True
+                thread.start()
+                thread.join(timeout=timeout)
+
+                execution_time = time.time() - start_time
+
+                if thread.is_alive():
+                    # Timeout occurred
+                    logger.warning(f"⚠️ [Docker] Command timed out after {timeout}s: {command}")
+                    return {
+                        "success": False,
+                        "exit_code": -1,
+                        "stdout": "",
+                        "stderr": f"Command timed out after {timeout} seconds",
+                        "execution_time": execution_time,
+                        "command": command,
+                        "work_dir": effective_workdir or "(container default)",
+                    }
+
+                if "error" in exception_container:
+                    raise exception_container["error"]
+
+                exit_code, output = result_container["data"]
+            else:
+                # No timeout - execute directly
+                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
+
+            if exit_code != 0:
+                logger.debug(f"⚠️ [Docker] Command exited with code {exit_code}")
+
+            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": effective_workdir or "(container default)",
+            }
+
+        except DockerException as e:
+            logger.error(f"❌ [Docker] Failed to execute command in container: {e}")
+            raise RuntimeError(f"Failed to execute command in container: {e}")
+
+    def stop_container(self, agent_id: str, timeout: int = 10) -> None:
+        """
+        Stop a container gracefully.
+
+        Args:
+            agent_id: Agent identifier
+            timeout: Seconds to wait before killing
+
+        Raises:
+            ValueError: If container not found
+        """
+        container = self.containers.get(agent_id)
+        if not container:
+            raise ValueError(f"No container found for agent {agent_id}")
+
+        try:
+            logger.info(f"🛑 [Docker] Stopping container {container.short_id} for agent {agent_id}")
+            container.stop(timeout=timeout)
+            logger.info("✅ [Docker] Container stopped successfully")
+        except DockerException as e:
+            logger.error(f"❌ [Docker] Failed to stop container for agent {agent_id}: {e}")
+
+    def remove_container(self, agent_id: str, force: bool = False) -> None:
+        """
+        Remove a container.
+
+        Args:
+            agent_id: Agent identifier
+            force: Force removal even if running
+
+        Raises:
+            ValueError: If container not found
+        """
+        container = self.containers.get(agent_id)
+        if not container:
+            raise ValueError(f"No container found for agent {agent_id}")
+
+        try:
+            container_id = container.short_id
+            logger.info(f"🗑️  [Docker] Removing container {container_id} for agent {agent_id}")
+            container.remove(force=force)
+            del self.containers[agent_id]
+            logger.info("✅ [Docker] Container removed successfully")
+        except DockerException as e:
+            logger.error(f"❌ [Docker] Failed to remove container for agent {agent_id}: {e}")
+
+    def cleanup(self, agent_id: Optional[str] = None) -> None:
+        """
+        Clean up containers.
+
+        Args:
+            agent_id: If provided, cleanup specific agent. Otherwise cleanup all.
+        """
+        if agent_id:
+            # Cleanup specific agent
+            if agent_id in self.containers:
+                logger.info(f"🧹 [Docker] Cleaning up container for agent {agent_id}")
+                try:
+                    self.stop_container(agent_id)
+                    self.remove_container(agent_id, force=True)
+                except Exception as e:
+                    logger.error(f"❌ [Docker] Error cleaning up container for agent {agent_id}: {e}")
+        else:
+            # Cleanup all containers
+            if self.containers:
+                logger.info(f"🧹 [Docker] Cleaning up {len(self.containers)} container(s)")
+            for aid in list(self.containers.keys()):
+                try:
+                    self.stop_container(aid)
+                    self.remove_container(aid, force=True)
+                except Exception as e:
+                    logger.error(f"❌ [Docker] Error cleaning up container for agent {aid}: {e}")
+
+    def log_container_info(self, agent_id: str) -> None:
+        """
+        Log detailed container information (useful for debugging).
+
+        Args:
+            agent_id: Agent identifier
+        """
+        container = self.containers.get(agent_id)
+        if not container:
+            logger.warning(f"⚠️ [Docker] No container found for agent {agent_id}")
+            return
+
+        try:
+            container.reload()  # Refresh state
+
+            logger.info(f"📊 [Docker] Container information for agent '{agent_id}':")
+            logger.info(f"    ID: {container.short_id}")
+            logger.info(f"    Name: {container.name}")
+            logger.info(f"    Status: {container.status}")
+            logger.info(f"    Network: {self.network_mode}")
+            if self.memory_limit:
+                logger.info(f"    Memory limit: {self.memory_limit}")
+            if self.cpu_limit:
+                logger.info(f"    CPU limit: {self.cpu_limit} cores")
+        except Exception as e:
+            logger.warning(f"⚠️ [Docker] Could not log container info: {e}")
+
+    def __del__(self):
+        """Cleanup all containers on deletion."""
+        try:
+            if hasattr(self, "containers") and self.containers:
+                self.cleanup()
+        except Exception:
+            # Silently fail during cleanup - already logged in cleanup()
+            pass