ms-enclave 0.0.3__py3-none-any.whl

ms_enclave/__init__.py

@@ -0,0 +1,3 @@
+from .version import __release_date__, __version__
+
+__all__ = ['__version__', '__release_date__']

ms_enclave/cli/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.

ms_enclave/cli/base.py

@@ -0,0 +1,20 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+from abc import ABC, abstractmethod
+from argparse import ArgumentParser
+
+
+class CLICommand(ABC):
+    """
+    Base class for command line tool.
+
+    """
+
+    @staticmethod
+    @abstractmethod
+    def define_args(parsers: ArgumentParser):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def execute(self):
+        raise NotImplementedError()

ms_enclave/cli/cli.py

@@ -0,0 +1,27 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import argparse
+
+from ms_enclave import __version__
+from ms_enclave.cli.start_server import ServerCMD
+
+
+def run_cmd():
+    parser = argparse.ArgumentParser('MS-Enclave Command Line tool', usage='ms-enclave <command> [<args>]')
+    parser.add_argument('-v', '--version', action='version', version=f'ms-enclave {__version__}')
+    subparsers = parser.add_subparsers(help='MS-Enclave command line helper.')
+
+    ServerCMD.define_args(subparsers)
+
+    args = parser.parse_args()
+
+    if not hasattr(args, 'func'):
+        parser.print_help()
+        exit(1)
+
+    cmd = args.func(args)
+    cmd.execute()
+
+
+if __name__ == '__main__':
+    run_cmd()

ms_enclave/cli/start_server.py

@@ -0,0 +1,84 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from argparse import ArgumentParser
+from typing import Optional
+
+from ms_enclave.cli.base import CLICommand
+from ms_enclave.sandbox import create_server
+from ms_enclave.utils import get_logger
+
+logger = get_logger()
+
+
+def subparser_func(args):
+    """ Function which will be called for a specific sub parser.
+    """
+    return ServerCMD(args)
+
+
+class ServerCMD(CLICommand):
+    name = 'server'
+
+    def __init__(self, args):
+        self.args = args
+
+    @staticmethod
+    def define_args(parsers: ArgumentParser):
+        """Define args for the server command.
+        """
+        parser = parsers.add_parser(ServerCMD.name, help='Start the MS-Enclave sandbox HTTP server')
+        add_argument(parser)
+        parser.set_defaults(func=subparser_func)
+
+    def execute(self):
+        """Start the sandbox server using provided CLI arguments."""
+        cleanup_interval: int = getattr(self.args, 'cleanup_interval', 300)
+        host: str = getattr(self.args, 'host', '0.0.0.0')
+        port: int = getattr(self.args, 'port', 8000)
+        log_level: str = getattr(self.args, 'log_level', 'info')
+        api_key: Optional[str] = getattr(self.args, 'api_key', None)
+
+        server = create_server(cleanup_interval=cleanup_interval, api_key=api_key)
+
+        logger.info('Starting Sandbox Server...')
+        logger.info('API docs: http://%s:%d/docs', host, port)
+        logger.info('Health check: http://%s:%d/health', host, port)
+
+        try:
+            server.run(host=host, port=port, log_level=log_level)
+        except KeyboardInterrupt:
+            logger.info('Server interrupted by user, shutting down...')
+        except Exception as e:  # pragma: no cover - runtime dependent
+            logger.error('Failed to start server: %s', e)
+            raise
+
+
+def add_argument(parser: ArgumentParser) -> None:
+    """Register command line arguments for the server command.
+
+    Args:
+        parser: The argparse parser to add arguments to.
+    """
+    parser.add_argument(
+        '--host', type=str, default='0.0.0.0', help='Host interface to bind the server (default: 0.0.0.0)'
+    )
+    parser.add_argument('--port', type=int, default=8000, help='Port for the HTTP server (default: 8000)')
+    parser.add_argument(
+        '--log-level',
+        type=str,
+        choices=['critical', 'error', 'warning', 'info', 'debug'],
+        default='info',
+        help='Log level for the server (default: info)'
+    )
+    parser.add_argument(
+        '--cleanup-interval',
+        type=int,
+        default=300,
+        metavar='SECONDS',
+        help='Background cleanup interval in seconds (default: 300)'
+    )
+    parser.add_argument(
+        '--api-key',
+        type=str,
+        default=None,
+        help='Optional API key to protect endpoints. If omitted, no authentication is enforced.'
+    )

ms_enclave/sandbox/__init__.py

@@ -0,0 +1,27 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+"""Modern agent sandbox system.
+
+A modular, extensible sandbox system for safe code execution with Docker isolation,
+FastAPI-based client/server architecture, and comprehensive tool support.
+"""
+
+from .boxes import DockerSandbox, Sandbox, SandboxFactory
+from .manager import HttpSandboxManager, LocalSandboxManager
+
+# Import main components
+from .model import (
+    DockerSandboxConfig,
+    ExecuteCodeRequest,
+    ExecuteCommandRequest,
+    ExecutionStatus,
+    HealthCheckResult,
+    ReadFileRequest,
+    SandboxConfig,
+    SandboxInfo,
+    SandboxStatus,
+    ToolExecutionRequest,
+    ToolResult,
+    WriteFileRequest,
+)
+from .server.server import SandboxServer, create_server
+from .tools import PythonExecutor, Tool, ToolFactory

ms_enclave/sandbox/boxes/__init__.py

@@ -0,0 +1,16 @@
+"""Sandbox implementations."""
+
+from .base import Sandbox, SandboxFactory, register_sandbox
+from .docker_notebook import DockerNotebookSandbox
+from .docker_sandbox import DockerSandbox
+
+__all__ = [
+    # Base interfaces
+    'Sandbox',
+    'SandboxFactory',
+    'register_sandbox',
+
+    # Implementations
+    'DockerSandbox',
+    'DockerNotebookSandbox',
+]

ms_enclave/sandbox/boxes/base.py

@@ -0,0 +1,274 @@
+"""Base sandbox interface and factory."""
+
+import abc
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Type, Union
+
+import shortuuid as uuid
+
+from ms_enclave.utils import get_logger
+
+from ..model import (
+    CommandResult,
+    DockerNotebookConfig,
+    DockerSandboxConfig,
+    SandboxConfig,
+    SandboxInfo,
+    SandboxStatus,
+    SandboxType,
+    ToolResult,
+)
+from ..tools import Tool, ToolFactory
+
+logger = get_logger()
+
+
+class Sandbox(abc.ABC):
+    """Abstract base class for all sandbox implementations."""
+
+    def __init__(self, config: SandboxConfig, sandbox_id: Optional[str] = None):
+        """Initialize sandbox.
+
+        Args:
+            config: Sandbox configuration
+            sandbox_id: Optional sandbox ID (will be generated if not provided)
+        """
+        self.id = sandbox_id or uuid.ShortUUID(alphabet='23456789abcdefghijkmnopqrstuvwxyz').random(length=8)
+        self.config = config
+        self.status = SandboxStatus.INITIALIZING
+        self.created_at = datetime.now()
+        self.updated_at = datetime.now()
+        self.metadata: Dict[str, Any] = {}
+        self._tools: Dict[str, Tool] = {}
+
+    @property
+    @abc.abstractmethod
+    def sandbox_type(self) -> SandboxType:
+        """Return the sandbox type identifier."""
+        pass
+
+    @abc.abstractmethod
+    async def start(self) -> None:
+        """Start the sandbox environment."""
+        pass
+
+    @abc.abstractmethod
+    async def stop(self) -> None:
+        """Stop the sandbox environment."""
+        pass
+
+    @abc.abstractmethod
+    async def cleanup(self) -> None:
+        """Clean up sandbox resources."""
+        pass
+
+    async def initialize_tools(self) -> None:
+        """Initialize sandbox tools."""
+        for tool_name, config in self.config.tools_config.items():
+            try:
+                tool = ToolFactory.create_tool(tool_name, **config)
+                self.add_tool(tool)
+            except Exception as e:
+                logger.error(f'Failed to initialize tool {tool_name}: {e}')
+
+    def get_available_tools(self) -> Dict[str, Any]:
+        """Get list of available tools."""
+        return {tool.name: tool.schema for tool in self._tools.values() if tool.enabled}
+
+    def get_tool(self, tool_name: str) -> Optional[Tool]:
+        """Get tool instance by type.
+
+        Args:
+            tool_name: Tool name
+
+        Returns:
+            Tool instance or None if not available
+        """
+        return self._tools.get(tool_name)
+
+    def add_tool(self, tool: Tool) -> None:
+        """Add a tool to the sandbox.
+
+        Args:
+            tool: Tool instance to add
+        """
+        if tool.name in self._tools:
+            logger.warning(f'Tool {tool.name} is already added to the sandbox')
+            return
+        if tool.enabled:
+            if tool.is_compatible_with_sandbox(self.sandbox_type):
+                self._tools[tool.name] = tool
+            else:
+                logger.warning(
+                    f"Tool '{tool.name}' requires sandbox type '{tool.required_sandbox_type}' "
+                    f"but this is a '{self.sandbox_type}' sandbox. "
+                    f'Compatible types: {SandboxType.get_compatible_types(self.sandbox_type)}'
+                )
+        else:
+            logger.warning(f'Tool {tool.name} is not enabled and cannot be added')
+
+    def list_tools(self) -> List[str]:
+        """
+        List all registered tools compatible with this sandbox.
+
+        Returns:
+            List of tool names
+        """
+        return list(self._tools.keys())
+
+    async def execute_tool(self, tool_name: str, parameters: Dict[str, Any]) -> ToolResult:
+        """Execute a tool with given parameters.
+
+        Args:
+            tool_name: Tool name
+            parameters: Tool parameters
+
+        Returns:
+            Tool execution result
+
+        Raises:
+            ValueError: If tool is not found or not enabled
+            TimeoutError: If tool execution exceeds timeout
+            Exception: For other execution errors
+        """
+        tool = self.get_tool(tool_name)
+        if not tool:
+            raise ValueError(f'Tool {tool_name} is not available')
+        if not tool.enabled:
+            raise ValueError(f'Tool {tool_name} is not enabled')
+
+        result = await tool.execute(sandbox_context=self, **parameters)
+        return result
+
+    async def execute_command(
+        self, command: Union[str, List[str]], timeout: Optional[int] = None, stream: bool = True
+    ) -> CommandResult:
+        """Execute a command in the sandbox environment.
+
+        Args:
+            command: Command to execute
+            timeout: Optional execution timeout in seconds
+            stream: Whether to stream output (if supported)
+        """
+        raise NotImplementedError('execute_command must be implemented by subclasses')
+
+    @abc.abstractmethod
+    async def get_execution_context(self) -> Any:
+        """Get the execution context for tools (e.g., container, process, etc.)."""
+        pass
+
+    def update_status(self, status: SandboxStatus) -> None:
+        """Update sandbox status.
+
+        Args:
+            status: New status
+        """
+        self.status = status
+        self.updated_at = datetime.now()
+
+    def get_info(self) -> SandboxInfo:
+        """Get sandbox information.
+
+        Returns:
+            Sandbox information
+        """
+        return SandboxInfo(
+            id=self.id,
+            status=self.status,
+            type=self.sandbox_type,
+            config=self.config.model_dump(exclude_none=True),
+            created_at=self.created_at,
+            updated_at=self.updated_at,
+            metadata=self.metadata,
+            available_tools=self.get_available_tools()
+        )
+
+    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()
+
+
+class SandboxFactory:
+    """Factory for creating sandbox instances."""
+
+    _sandboxes: Dict[SandboxType, Type[Sandbox]] = {}
+
+    @classmethod
+    def register_sandbox(cls, sandbox_type: SandboxType, sandbox_class: Type[Sandbox]):
+        """Register a sandbox class.
+
+        Args:
+            sandbox_type: Sandbox type identifier
+            sandbox_class: Sandbox class
+        """
+        cls._sandboxes[sandbox_type] = sandbox_class
+
+    @classmethod
+    def create_sandbox(
+        cls,
+        sandbox_type: SandboxType,
+        config: Optional[Union[SandboxConfig, Dict]] = None,
+        sandbox_id: Optional[str] = None
+    ) -> Sandbox:
+        """Create a sandbox instance.
+
+        Args:
+            sandbox_type: Sandbox type
+            config: Sandbox configuration
+            sandbox_id: Optional sandbox ID
+
+        Returns:
+            Sandbox instance
+
+        Raises:
+            ValueError: If sandbox type is not registered
+        """
+        if sandbox_type not in cls._sandboxes:
+            raise ValueError(f'Sandbox type {sandbox_type} is not registered')
+
+        # Parse config based on sandbox type
+        if not config:
+            if sandbox_type == SandboxType.DOCKER:
+                config = DockerSandboxConfig()
+            elif sandbox_type == SandboxType.DOCKER_NOTEBOOK:
+                config = DockerNotebookConfig()
+            else:
+                config = SandboxConfig()
+        elif isinstance(config, dict):
+            if sandbox_type == SandboxType.DOCKER:
+                config = DockerSandboxConfig(**config)
+            elif sandbox_type == SandboxType.DOCKER_NOTEBOOK:
+                config = DockerNotebookConfig(**config)
+            else:
+                config = SandboxConfig(**config)
+
+        sandbox_class = cls._sandboxes[sandbox_type]
+        return sandbox_class(config, sandbox_id)
+
+    @classmethod
+    def get_available_types(cls) -> List[SandboxType]:
+        """Get list of available sandbox types.
+
+        Returns:
+            List of available sandbox types
+        """
+        return list(cls._sandboxes.keys())
+
+
+def register_sandbox(sandbox_type: SandboxType):
+    """Decorator for registering sandboxes.
+
+    Args:
+        sandbox_type: Sandbox type identifier
+    """
+
+    def decorator(sandbox_class: Type[Sandbox]):
+        SandboxFactory.register_sandbox(sandbox_type, sandbox_class)
+        return sandbox_class
+
+    return decorator

ms_enclave/sandbox/boxes/docker_notebook.py

@@ -0,0 +1,224 @@
+# flake8: noqa E501
+import asyncio
+import json
+import tempfile
+from pathlib import Path
+from textwrap import dedent
+from typing import Optional
+
+from ms_enclave.utils import get_logger
+
+from ..model import DockerNotebookConfig, SandboxStatus, SandboxType
+from .base import register_sandbox
+from .docker_sandbox import DockerSandbox
+
+logger = get_logger()
+
+
+@register_sandbox(SandboxType.DOCKER_NOTEBOOK)
+class DockerNotebookSandbox(DockerSandbox):
+    """
+    Docker sandbox that executes Python code using Jupyter Kernel Gateway.
+    """
+
+    def __init__(
+        self,
+        config: DockerNotebookConfig,
+        sandbox_id: Optional[str] = None,
+    ):
+        """
+        Initialize the Docker-based Jupyter Kernel Gateway executor.
+
+        Args:
+            config: Docker sandbox configuration
+            sandbox_id: Optional sandbox ID
+            host: Host to bind to.
+            port: Port to bind to.
+        """
+        super().__init__(config, sandbox_id)
+
+        self.config: DockerNotebookConfig = config
+        self.host = self.config.host
+        self.port = self.config.port
+        self.kernel_id = None
+        self.ws = None
+        self.base_url = None
+        self.config.ports['8888/tcp'] = self.port
+        self.config.network_enabled = True  # Ensure network is enabled for Jupyter
+
+    @property
+    def sandbox_type(self) -> SandboxType:
+        """Return sandbox type."""
+        return SandboxType.DOCKER_NOTEBOOK
+
+    async def start(self) -> None:
+        """Start the Docker container with Jupyter Kernel Gateway."""
+        try:
+            self.update_status(SandboxStatus.INITIALIZING)
+
+            # Initialize Docker client first
+            import docker
+            self.client = docker.from_env()
+
+            # Build Jupyter image if needed before creating container
+            await self._build_jupyter_image()
+
+            # Now start the base container with the Jupyter image
+            await super().start()
+
+            # Setup Jupyter kernel gateway services
+            await self._setup_jupyter()
+
+            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 Jupyter Docker sandbox: {e}')
+            raise RuntimeError(f'Failed to start Jupyter Docker sandbox: {e}')
+
+    async def _setup_jupyter(self) -> None:
+        """Setup Jupyter Kernel Gateway services in the container."""
+        try:
+            # Wait for Jupyter Kernel Gateway to be ready
+            await self._wait_for_jupyter_ready()
+
+            # Create kernel and establish websocket connection
+            await self._create_kernel()
+
+        except Exception as e:
+            logger.error(f'Failed to setup Jupyter: {e}')
+            raise
+
+    async def _wait_for_jupyter_ready(self) -> None:
+        """Wait for Jupyter Kernel Gateway to be ready."""
+        import requests
+
+        self.base_url = f'http://{self.host}:{self.port}'
+        max_retries = 10  # Wait up to 30 seconds
+        retry_interval = 3  # Check every 3 second
+
+        for attempt in range(max_retries):
+            try:
+                # Try to get the API status
+                response = requests.get(f'{self.base_url}/api', timeout=5)
+                if response.status_code == 200:
+                    logger.info(f'Jupyter Kernel Gateway is ready at {self.base_url}')
+                    return
+            except requests.exceptions.RequestException:
+                # Connection failed, Jupyter not ready yet
+                pass
+
+            if attempt < max_retries - 1:
+                logger.info(f'Waiting for Jupyter Kernel Gateway to be ready... (attempt {attempt + 1}/{max_retries})')
+                await asyncio.sleep(retry_interval)
+
+        raise RuntimeError(f'Jupyter Kernel Gateway failed to become ready within {max_retries} seconds')
+
+    async def _build_jupyter_image(self) -> None:
+        """Build or ensure Jupyter image exists."""
+        # Step 1: Try to get the image directly
+        try:
+            self.client.images.get(self.config.image)
+            logger.info(f'Using existing Docker image: {self.config.image}')
+            return
+        except Exception as e:
+            logger.debug(f'Direct image get failed: {e}, trying list method...')
+
+        # Step 2: Try to find image in list
+        image_exists = any(self.config.image in img.tags for img in self.client.images.list() if img.tags)
+        if image_exists:
+            logger.info(f'Using existing Docker image: {self.config.image}')
+            return
+
+        # Step 3: Image not found, build it
+        logger.info(f'Building Docker image {self.config.image}...')
+
+        # Create Dockerfile
+        dockerfile_content = dedent(
+            """\
+            FROM python:3.12-slim
+
+            RUN pip install jupyter_kernel_gateway jupyter_client ipykernel
+
+            # Install and register the Python kernel
+            RUN python -m ipykernel install --sys-prefix --name python3 --display-name "Python 3"
+
+            EXPOSE 8888
+            CMD ["jupyter", "kernelgateway", "--KernelGatewayApp.ip=0.0.0.0", "--KernelGatewayApp.port=8888", "--KernelGatewayApp.allow_origin=*"]
+            """
+        )
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            dockerfile_path = Path(tmpdir) / 'Dockerfile'
+            dockerfile_path.write_text(dockerfile_content)
+
+            # Build image with output
+            def build_image():
+                build_logs = self.client.images.build(
+                    path=tmpdir, dockerfile='Dockerfile', tag=self.config.image, rm=True
+                )
+                # Process and log build output
+                for log in build_logs[1]:  # build_logs[1] contains the build log generator
+                    if 'stream' in log:
+                        logger.info(f"[📦 {self.id}] {log['stream'].strip()}")
+                    elif 'error' in log:
+                        logger.error(f"[📦 {self.id}] {log['error']}")
+                return build_logs[0]  # Return the built image
+
+            await asyncio.get_event_loop().run_in_executor(None, build_image)
+
+    async def _create_kernel(self) -> None:
+        """Create a new kernel and establish websocket connection."""
+        import requests
+
+        # Create new kernel via HTTP
+        response = requests.post(f'{self.base_url}/api/kernels')
+        if response.status_code != 201:
+            error_details = {
+                'status_code': response.status_code,
+                'headers': dict(response.headers),
+                'url': response.url,
+                'body': response.text,
+                'request_method': response.request.method,
+                'request_headers': dict(response.request.headers),
+                'request_body': response.request.body,
+            }
+            raise RuntimeError(f'Failed to create kernel: {json.dumps(error_details, indent=2)}')
+
+        self.kernel_id = response.json()['id']
+
+        # Establish websocket connection
+        try:
+            from websocket import create_connection
+            ws_url = f'ws://{self.host}:{self.port}/api/kernels/{self.kernel_id}/channels'
+            self.ws = create_connection(ws_url)
+            logger.info(f'Kernel {self.kernel_id} created and connected')
+        except ImportError:
+            raise RuntimeError('websocket-client package is required. Install with: pip install websocket-client')
+
+    async def cleanup(self) -> None:
+        """Clean up Jupyter resources and Docker container."""
+        try:
+            # Close websocket connection
+            if self.ws:
+                try:
+                    self.ws.close()
+                except Exception:
+                    pass
+                self.ws = None
+
+            # Delete kernel
+            if self.kernel_id and self.base_url:
+                try:
+                    import requests
+                    requests.delete(f'{self.base_url}/api/kernels/{self.kernel_id}')
+                except Exception:
+                    pass
+                self.kernel_id = None
+
+        except Exception as e:
+            logger.error(f'Error during Jupyter cleanup: {e}')
+
+        # Call parent cleanup
+        await super().cleanup()