ms-enclave 0.0.3__py3-none-any.whl

ms_enclave/sandbox/boxes/docker_sandbox.py

@@ -0,0 +1,317 @@
+"""Docker-based sandbox implementation."""
+
+import asyncio
+import time
+from typing import Any, Dict, List, Optional, Union
+
+import docker
+from docker import DockerClient
+from docker.errors import APIError, ContainerError, ImageNotFound, NotFound
+from docker.models.containers import Container
+
+from ms_enclave.utils import get_logger
+
+from ..model import CommandResult, DockerSandboxConfig, ExecutionStatus, SandboxStatus, SandboxType
+from .base import Sandbox, register_sandbox
+
+logger = get_logger()
+
+
+@register_sandbox(SandboxType.DOCKER)
+class DockerSandbox(Sandbox):
+    """Docker-based sandbox implementation."""
+
+    def __init__(self, config: DockerSandboxConfig, sandbox_id: Optional[str] = None):
+        """Initialize Docker sandbox.
+
+        Args:
+            config: Docker sandbox configuration
+            sandbox_id: Optional sandbox ID
+        """
+        super().__init__(config, sandbox_id)
+        self.config: DockerSandboxConfig = config
+        self.client: Optional[DockerClient] = None
+        self.container: Optional[Container] = None
+
+    @property
+    def sandbox_type(self) -> SandboxType:
+        """Return sandbox type."""
+        return SandboxType.DOCKER
+
+    async def start(self) -> None:
+        """Start the Docker container."""
+        try:
+            self.update_status(SandboxStatus.INITIALIZING)
+
+            # Initialize Docker client
+            self.client = docker.from_env()
+
+            # Ensure image exists
+            await self._ensure_image_exists()
+
+            # Create and start container
+            await self._create_container()
+            await self._start_container()
+
+            # Initialize tools
+            await self.initialize_tools()
+
+            self.update_status(SandboxStatus.RUNNING)
+
+        except Exception as e:
+            self.update_status(SandboxStatus.ERROR)
+            self.metadata['error'] = str(e)
+            logger.error(f'Failed to start Docker sandbox: {e}')
+            raise RuntimeError(f'Failed to start Docker sandbox: {e}')
+
+    async def stop(self) -> None:
+        """Stop the Docker container without removing it unless configured.
+
+        When remove_on_exit is False, this method stops the container but keeps
+        the container reference so get_execution_context() can return it.
+        """
+        if not self.container:
+            self.update_status(SandboxStatus.STOPPED)
+            return
+
+        try:
+            self.update_status(SandboxStatus.STOPPING)
+            await self.stop_container()
+
+            # If configured to remove on exit, perform full cleanup (removes container and closes client)
+            if self.config.remove_on_exit:
+                await self.cleanup()
+
+            self.update_status(SandboxStatus.STOPPED)
+        except Exception as e:
+            logger.error(f'Error stopping container: {e}')
+            self.update_status(SandboxStatus.ERROR)
+            raise
+
+    async def cleanup(self) -> None:
+        """Clean up Docker resources.
+
+        - Always stops the container if it is running.
+        - Removes the container only when remove_on_exit is True.
+        - Preserves container reference and client when remove_on_exit is False.
+        """
+        if self.container:
+            try:
+                self.container.remove(force=True)
+                logger.debug(f'Container {self.container.id} removed')
+            except Exception as e:
+                logger.error(f'Error cleaning up container: {e}')
+            finally:
+                # Only drop the reference when we actually removed it
+                self.container = None
+
+        # Close Docker client only if we dropped the container reference
+        if self.client:
+            try:
+                self.client.close()
+            except Exception as e:
+                logger.warning(f'Error closing Docker client: {e}')
+            finally:
+                self.client = None
+
+    async def stop_container(self) -> None:
+        """Stop the container if it is running."""
+        if not self.container:
+            return
+        try:
+            self.container.reload()
+            if self.container.status == SandboxStatus.RUNNING:
+                self.container.stop(timeout=10)
+        except NotFound:
+            logger.warning('Container not found while stopping')
+        except Exception as e:
+            logger.error(f'Error stopping container: {e}')
+            raise
+
+    async def get_execution_context(self) -> Any:
+        """Return the container for tool execution."""
+        return self.container
+
+    def _run_streaming(self, command: Union[str, List[str]]) -> tuple[int, str, str]:
+        """Execute command with streaming logs using low-level API.
+
+        Returns:
+            A tuple of (exit_code, stdout, stderr)
+        """
+        if not self.client or not self.container:
+            raise RuntimeError('Container is not running')
+
+        # Use low-level API for precise control over streaming and exit code.
+        exec_id = self.client.api.exec_create(
+            container=self.container.id,
+            cmd=command,
+            tty=False,
+        )['Id']
+
+        stdout_parts: List[str] = []
+        stderr_parts: List[str] = []
+
+        try:
+            for chunk in self.client.api.exec_start(exec_id, stream=True, demux=True):
+                if not chunk:
+                    continue
+                out, err = chunk  # each is Optional[bytes]
+                if out:
+                    text = out.decode('utf-8', errors='replace')
+                    stdout_parts.append(text)
+                    for line in text.splitlines():
+                        logger.info(f'[📦 {self.id}] {line}')
+                if err:
+                    text = err.decode('utf-8', errors='replace')
+                    stderr_parts.append(text)
+                    for line in text.splitlines():
+                        logger.error(f'[📦 {self.id}] {line}')
+        finally:
+            inspect = self.client.api.exec_inspect(exec_id)
+            exit_code = inspect.get('ExitCode')
+            if exit_code is None:
+                exit_code = -1
+
+        return exit_code, ''.join(stdout_parts), ''.join(stderr_parts)
+
+    def _run_buffered(self, command: Union[str, List[str]]) -> tuple[int, str, str]:
+        """Execute command and return buffered output using high-level API.
+
+        Returns:
+            A tuple of (exit_code, stdout, stderr)
+        """
+        if not self.container:
+            raise RuntimeError('Container is not running')
+
+        res = self.container.exec_run(command, tty=False, stream=False, demux=True)
+        out_tuple = res.output
+        if isinstance(out_tuple, tuple):
+            out_bytes, err_bytes = out_tuple
+        else:
+            # Fallback: when demux was not honored, treat all as stdout
+            out_bytes, err_bytes = out_tuple, b''
+
+        stdout = out_bytes.decode('utf-8', errors='replace') if out_bytes else ''
+        stderr = err_bytes.decode('utf-8', errors='replace') if err_bytes else ''
+        return res.exit_code, stdout, stderr
+
+    async def execute_command(
+        self, command: Union[str, List[str]], timeout: Optional[int] = None, stream: bool = True
+    ) -> CommandResult:
+        """Execute a command in the container.
+
+        When stream=True (default), logs are printed in real-time through the logger,
+        while stdout/stderr are still accumulated and returned in the result.
+        When stream=False, the command is executed and buffered, returning the full output at once.
+
+        Args:
+            command: Command to run (str or list)
+            timeout: Optional timeout in seconds
+            stream: Whether to stream logs in real time
+
+        Returns:
+            CommandResult with status, exit_code, stdout and stderr
+        """
+        if not self.container or not self.client:
+            raise RuntimeError('Container is not running')
+
+        loop = asyncio.get_running_loop()
+
+        run_func = self._run_streaming if stream else self._run_buffered
+        try:
+            exit_code, stdout, stderr = await asyncio.wait_for(
+                loop.run_in_executor(None, lambda: run_func(command)), timeout=timeout
+            )
+            status = ExecutionStatus.SUCCESS if exit_code == 0 else ExecutionStatus.ERROR
+            return CommandResult(command=command, status=status, exit_code=exit_code, stdout=stdout, stderr=stderr)
+        except asyncio.TimeoutError:
+            return CommandResult(
+                command=command,
+                status=ExecutionStatus.TIMEOUT,
+                exit_code=-1,
+                stdout='',
+                stderr=f'Command timed out after {timeout} seconds'
+            )
+        except Exception as e:
+            return CommandResult(command=command, status=ExecutionStatus.ERROR, exit_code=-1, stdout='', stderr=str(e))
+
+    async def _ensure_image_exists(self) -> None:
+        """Ensure Docker image exists."""
+        try:
+            self.client.images.get(self.config.image)
+        except ImageNotFound:
+            # Try to pull the image
+            try:
+                self.client.images.pull(self.config.image)
+            except Exception as e:
+                raise RuntimeError(f'Failed to pull image {self.config.image}: {e}')
+
+    async def _create_container(self) -> None:
+        """Create Docker container."""
+        try:
+            # Prepare container configuration
+            container_config = {
+                'image': self.config.image,
+                'name': f'sandbox-{self.id}',
+                'working_dir': self.config.working_dir,
+                'environment': self.config.env_vars,
+                'detach': True,
+                'tty': True,
+                'stdin_open': True,
+            }
+
+            # Add command if specified
+            if self.config.command:
+                container_config['command'] = self.config.command
+
+            # Add resource limits
+            if self.config.memory_limit:
+                container_config['mem_limit'] = self.config.memory_limit
+
+            if self.config.cpu_limit:
+                container_config['cpu_quota'] = int(self.config.cpu_limit * 100000)
+                container_config['cpu_period'] = 100000
+
+            # Add volumes
+            if self.config.volumes:
+                container_config['volumes'] = self.config.volumes
+
+            # Add ports
+            if self.config.ports:
+                container_config['ports'] = self.config.ports
+
+            # Network configuration
+            if not self.config.network_enabled:
+                container_config['network_mode'] = 'none'
+            elif self.config.network:
+                container_config['network'] = self.config.network
+
+            # Privileged mode
+            container_config['privileged'] = self.config.privileged
+
+            # Create container
+            self.container = self.client.containers.create(**container_config)
+            self.metadata['container_id'] = self.container.id
+
+        except Exception as e:
+            raise RuntimeError(f'Failed to create container: {e}')
+
+    async def _start_container(self) -> None:
+        """Start Docker container."""
+        try:
+            self.container.start()
+
+            # Wait for container to be ready
+            timeout = 30
+            start_time = time.time()
+
+            while time.time() - start_time < timeout:
+                self.container.reload()
+                if self.container.status == 'running':
+                    break
+                await asyncio.sleep(0.5)
+            else:
+                raise RuntimeError('Container failed to start within timeout')
+
+        except Exception as e:
+            raise RuntimeError(f'Failed to start container: {e}')

ms_enclave/sandbox/manager/__init__.py

@@ -0,0 +1,11 @@
+"""Sandbox manager implementations."""
+
+from .base import SandboxManager
+from .http_manager import HttpSandboxManager
+from .local_manager import LocalSandboxManager
+
+__all__ = [
+    'SandboxManager',
+    'LocalSandboxManager',
+    'HttpSandboxManager',
+]

ms_enclave/sandbox/manager/base.py

@@ -0,0 +1,155 @@
+"""Base sandbox manager interface."""
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+from ..model import SandboxConfig, SandboxInfo, SandboxStatus, SandboxType, ToolResult
+
+if TYPE_CHECKING:
+    from ..boxes import Sandbox
+
+
+class SandboxManager(ABC):
+    """Abstract base class for sandbox managers."""
+
+    def __init__(self):
+        """Initialize the sandbox manager."""
+        self._running = False
+        self._sandboxes: Dict[str, 'Sandbox'] = {}
+
+    @abstractmethod
+    async def start(self) -> None:
+        """Start the sandbox manager."""
+        pass
+
+    @abstractmethod
+    async def stop(self) -> None:
+        """Stop the sandbox manager."""
+        pass
+
+    @abstractmethod
+    async def create_sandbox(
+        self,
+        sandbox_type: SandboxType,
+        config: Optional[Union[SandboxConfig, Dict]] = None,
+        sandbox_id: Optional[str] = None
+    ) -> str:
+        """Create a new sandbox.
+
+        Args:
+            sandbox_type: Type of sandbox to create
+            config: Sandbox configuration
+            sandbox_id: Optional sandbox ID
+
+        Returns:
+            Sandbox ID
+
+        Raises:
+            ValueError: If sandbox type is not supported
+            RuntimeError: If sandbox creation fails
+        """
+        pass
+
+    @abstractmethod
+    async def get_sandbox_info(self, sandbox_id: str) -> Optional[SandboxInfo]:
+        """Get sandbox information.
+
+        Args:
+            sandbox_id: Sandbox ID
+
+        Returns:
+            Sandbox information or None if not found
+        """
+        pass
+
+    @abstractmethod
+    async def list_sandboxes(self, status_filter: Optional[SandboxStatus] = None) -> List[SandboxInfo]:
+        """List all sandboxes.
+
+        Args:
+            status_filter: Optional status filter
+
+        Returns:
+            List of sandbox information
+        """
+        pass
+
+    @abstractmethod
+    async def stop_sandbox(self, sandbox_id: str) -> bool:
+        """Stop a sandbox.
+
+        Args:
+            sandbox_id: Sandbox ID
+
+        Returns:
+            True if stopped successfully, False if not found
+        """
+        pass
+
+    @abstractmethod
+    async def delete_sandbox(self, sandbox_id: str) -> bool:
+        """Delete a sandbox.
+
+        Args:
+            sandbox_id: Sandbox ID
+
+        Returns:
+            True if deleted successfully, False if not found
+        """
+        pass
+
+    @abstractmethod
+    async def execute_tool(self, sandbox_id: str, tool_name: str, parameters: Dict[str, Any]) -> ToolResult:
+        """Execute tool in sandbox.
+
+        Args:
+            sandbox_id: Sandbox ID
+            tool_name: Tool name to execute
+            parameters: Tool parameters
+
+        Returns:
+            Tool execution result
+
+        Raises:
+            ValueError: If sandbox or tool not found
+        """
+        pass
+
+    @abstractmethod
+    async def get_sandbox_tools(self, sandbox_id: str) -> Dict[str, Any]:
+        """Get available tools for a sandbox.
+
+        Args:
+            sandbox_id: Sandbox ID
+
+        Returns:
+            Dictionary of available tool types, e.g., {"tool_name": tool_schema}
+
+        Raises:
+            ValueError: If sandbox not found
+        """
+        pass
+
+    @abstractmethod
+    async def get_stats(self) -> Dict[str, Any]:
+        """Get manager statistics.
+
+        Returns:
+            Statistics dictionary
+        """
+        pass
+
+    @abstractmethod
+    async def cleanup_all_sandboxes(self) -> None:
+        """Clean up all sandboxes."""
+        pass
+
+    # Context manager support
+    async def __aenter__(self):
+        """Async context manager entry."""
+        await self.start()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Async context manager exit."""
+        await self.stop()