diff --git a/ms_enclave/sandbox/boxes/__init__.py b/ms_enclave/sandbox/boxes/__init__.py index d0c83d5..1a722c7 100644 --- a/ms_enclave/sandbox/boxes/__init__.py +++ b/ms_enclave/sandbox/boxes/__init__.py @@ -3,6 +3,7 @@ from .base import Sandbox, SandboxFactory, register_sandbox from .docker_notebook import DockerNotebookSandbox from .docker_sandbox import DockerSandbox +from .e2b_sandbox import E2BSandbox from .stateless_sandbox import StatelessSandbox from .volcengine import VolcengineSandbox @@ -14,6 +15,7 @@ 'StatelessSandbox', # Implementations 'DockerSandbox', + 'E2BSandbox', 'DockerNotebookSandbox', 'VolcengineSandbox', ] diff --git a/ms_enclave/sandbox/boxes/base.py b/ms_enclave/sandbox/boxes/base.py index ec7e67c..a03994f 100644 --- a/ms_enclave/sandbox/boxes/base.py +++ b/ms_enclave/sandbox/boxes/base.py @@ -15,6 +15,7 @@ CommandResult, DockerNotebookConfig, DockerSandboxConfig, + E2BSandboxConfig, SandboxConfig, SandboxInfo, SandboxStatus, @@ -265,6 +266,8 @@ def create_sandbox( config = DockerNotebookConfig() elif sandbox_type == SandboxType.VOLCENGINE: config = VolcengineSandboxConfig() + elif sandbox_type == SandboxType.E2B: + config = E2BSandboxConfig() else: config = SandboxConfig() elif isinstance(config, dict): @@ -274,6 +277,8 @@ def create_sandbox( config = DockerNotebookConfig(**config) elif sandbox_type == SandboxType.VOLCENGINE: config = VolcengineSandboxConfig(**config) + elif sandbox_type == SandboxType.E2B: + config = E2BSandboxConfig(**config) else: config = SandboxConfig(**config) diff --git a/ms_enclave/sandbox/boxes/docker_sandbox.py b/ms_enclave/sandbox/boxes/docker_sandbox.py index 2c1b8f6..8432b7f 100644 --- a/ms_enclave/sandbox/boxes/docker_sandbox.py +++ b/ms_enclave/sandbox/boxes/docker_sandbox.py @@ -536,6 +536,8 @@ async def _create_container(self) -> None: # Privileged mode container_config['privileged'] = self.config.privileged + if self.config.cap_add: + container_config['cap_add'] = list(self.config.cap_add) # Platform (e.g. linux/amd64 for amd64 images on Apple Silicon) if self.config.platform: diff --git a/ms_enclave/sandbox/boxes/e2b_connection.py b/ms_enclave/sandbox/boxes/e2b_connection.py new file mode 100644 index 0000000..dbe2793 --- /dev/null +++ b/ms_enclave/sandbox/boxes/e2b_connection.py @@ -0,0 +1,96 @@ +"""Optional SDK boundary and connection settings for E2B services.""" + +import asyncio +from dataclasses import dataclass +from typing import Any, Awaitable, Callable, Dict, Optional, TypeVar + +import httpx +from pydantic import SecretStr + +from ms_enclave.utils import get_logger + +from ..model import E2BSandboxManagerConfig +from ..utils.e2b_helpers import BuildMode, E2BProvider, redact + + +class _MissingE2BError(Exception): + """Placeholder exception when the optional SDK is unavailable.""" + + +try: + from e2b import AsyncSandbox, AsyncTemplate, CommandExitException, RateLimitException + from e2b.api.client_async import get_api_client + from e2b.connection_config import ConnectionConfig + + _E2B_IMPORT_ERROR: Optional[ImportError] = None +except ImportError as exc: + AsyncSandbox = AsyncTemplate = get_api_client = ConnectionConfig = None + CommandExitException = RateLimitException = _MissingE2BError + _E2B_IMPORT_ERROR = exc + +logger = get_logger() +E2B_REQUEST_TIMEOUT_S = 60.0 +_DISPATCH_ATTEMPTS = 3 +T = TypeVar('T') + + +def require_e2b() -> None: + """Fail with an installation hint when E2B was not installed.""" + if _E2B_IMPORT_ERROR is not None: + raise ImportError( + "The E2B sandbox requires the e2b SDK (Python >= 3.11): pip install 'ms-enclave[e2b]'" + ) from _E2B_IMPORT_ERROR + + +@dataclass(frozen=True) +class E2BConnection: + """Resolved settings shared by the manager, sandboxes and template resolver.""" + + api_key: SecretStr + api_url: Optional[str] + domain: Optional[str] + provider: E2BProvider + region: Optional[str] + build_mode: BuildMode + + @classmethod + def from_config(cls, config: E2BSandboxManagerConfig) -> 'E2BConnection': + """Resolve a validated manager configuration.""" + return cls( + config.api_key, + config.api_url, + config.domain, + config.provider, + config.target_region, + config.resolved_build_mode, + ) + + def api_params(self, **extra: Any) -> Dict[str, Any]: + """Build SDK options, omitting unset endpoints to preserve SDK defaults.""" + params = {'api_key': self.api_key.get_secret_value(), 'request_timeout': E2B_REQUEST_TIMEOUT_S} + if self.api_url: + params['api_url'] = self.api_url + if self.domain: + params['domain'] = self.domain + return {**params, **extra} + + def redact(self, text: str) -> str: + """Remove the connection key from diagnostic text.""" + return redact(text, self.api_key.get_secret_value()) + + +async def retry_dispatch( + op_name: str, call: Callable[[], Awaitable[T]], *, run_id: str, redactor: Callable[[str], str] +) -> T: + """Retry only connection establishment failures and explicit rate limits.""" + retryable = (ConnectionError, httpx.ConnectError, httpx.ConnectTimeout, httpx.PoolTimeout, RateLimitException) + for attempt in range(1, _DISPATCH_ATTEMPTS + 1): + try: + return await call() + except Exception as exc: + reason = redactor(str(exc)) + rate_limited = isinstance(exc, httpx.HTTPStatusError) and exc.response.status_code == 429 + if (not isinstance(exc, retryable) and not rate_limited) or attempt == _DISPATCH_ATTEMPTS: + raise RuntimeError(f'E2B {op_name} failed (run_id={run_id}): {reason}') from None + logger.warning(f'E2B {op_name} attempt={attempt} run_id={run_id}: {reason}; retrying') + await asyncio.sleep(2 ** (attempt - 1)) diff --git a/ms_enclave/sandbox/boxes/e2b_sandbox.py b/ms_enclave/sandbox/boxes/e2b_sandbox.py new file mode 100644 index 0000000..f2fb154 --- /dev/null +++ b/ms_enclave/sandbox/boxes/e2b_sandbox.py @@ -0,0 +1,201 @@ +"""Stateful remote sandbox backed by the optional E2B SDK.""" + +import asyncio +import posixpath +import shlex +import time +import uuid +from typing import Any, Dict, List, Optional, Union + +from ms_enclave.utils import get_logger + +from ..model import ( + CommandResult, + E2BSandboxConfig, + E2BSandboxManagerConfig, + ExecutionStatus, + SandboxStatus, + SandboxType, +) +from .base import Sandbox, register_sandbox +from .e2b_connection import AsyncSandbox, CommandExitException, E2BConnection, require_e2b, retry_dispatch +from .e2b_template import E2BTemplateResolver + +logger = get_logger() + + +@register_sandbox(SandboxType.E2B) +class E2BSandbox(Sandbox): + """Execute commands and transfer archives without a local Docker daemon.""" + + def __init__( + self, + config: E2BSandboxConfig, + sandbox_id: Optional[str] = None, + *, + connection: Optional[E2BConnection] = None, + templates: Optional[E2BTemplateResolver] = None, + run_id: Optional[str] = None, + ) -> None: + require_e2b() + super().__init__(config, sandbox_id) + self.config: E2BSandboxConfig = config + self._conn = connection + self._templates = templates + self._run_id = run_id or uuid.uuid4().hex[:12] + self._sbx: Any = None + self._started = 0.0 + + @property + def sandbox_type(self) -> SandboxType: + """Return the registered E2B type.""" + return SandboxType.E2B + + async def start(self) -> None: + """Create the remote sandbox with a fixed lifetime and run metadata.""" + if self._conn is None: + cfg = E2BSandboxManagerConfig() + self._conn = E2BConnection.from_config(cfg) + if self._templates is None: + self._templates = E2BTemplateResolver(self._conn, {}, 4) + self.update_status(SandboxStatus.INITIALIZING) + self._started = time.monotonic() + template = await self._templates.resolve(self.config, run_id=self._run_id) + try: + self._sbx = await retry_dispatch( + 'create sandbox', + lambda: AsyncSandbox.create( + template=template, + timeout=self.config.sandbox_timeout, + metadata={'created_by': 'ms-enclave', 'run_id': self._run_id}, + allow_internet_access=self.config.allow_internet_access, + **self._conn.api_params(retries=0), + ), + run_id=self._run_id, + redactor=self._conn.redact, + ) + except Exception as exc: + self.update_status(SandboxStatus.ERROR) + message = ( + f'Failed to create E2B sandbox (template={template}, run_id={self._run_id}, ' + f'sandbox_timeout={self.config.sandbox_timeout}): {exc}' + ) + if 'cannot be greater than' in str(exc): + message += ( + '; sandbox_timeout may exceed the plan limit (e2b official Hobby: 3600s); ' + 'lower sandbox.default_config.sandbox_timeout or upgrade the plan.' + ) + raise RuntimeError(self._conn.redact(message)) from None + self.metadata.update(remote_id=self._sbx.sandbox_id, template=template, run_id=self._run_id) + await self.initialize_tools() + self.update_status(SandboxStatus.RUNNING) + logger.info( + f'E2B sandbox {self.id} remote={self._sbx.sandbox_id} template={template} ' + f'run_id={self._run_id} elapsed={time.monotonic() - self._started:.1f}s' + ) + + async def stop(self) -> None: + """Kill once, removing the local reference even if the service rejects kill.""" + remote, self._sbx = self._sbx, None + if remote is not None: + try: + await remote.kill(**self._conn.api_params(retries=0)) + except Exception as exc: + logger.warning( + self._conn.redact(f'E2B sandbox {self.id} remote={remote.sandbox_id}: kill failed: {exc}') + ) + logger.info( + f'E2B sandbox {self.id} stopped run_id={self._run_id} lifetime={time.monotonic() - self._started:.1f}s' + ) + self.update_status(SandboxStatus.STOPPED) + + async def cleanup(self) -> None: + """Destroy the sandbox.""" + await self.stop() + + async def get_execution_context(self) -> Any: + """Return the underlying SDK sandbox.""" + return self._sbx + + def _result( + self, + command: Union[str, List[str]], + status: ExecutionStatus, + code: int, + stdout: str, + stderr: str, + started: float, + ) -> CommandResult: + return CommandResult( + command=command, + status=status, + exit_code=code, + stdout=stdout, + stderr=stderr, + execution_time=time.monotonic() - started, + ) + + async def execute_command( + self, + command: Union[str, List[str]], + timeout: Optional[int] = None, + stream: bool = True, + *, + user: Optional[str] = None, + cwd: Optional[str] = None, + env: Optional[Dict[str, str]] = None, + ) -> CommandResult: + """Run once, preserve command output and kill the process on local timeout.""" + if self._sbx is None: + raise RuntimeError('E2B sandbox is not running') + cmd = command if isinstance(command, str) else shlex.join(command) + started = time.monotonic() + try: + handle = await self._sbx.commands.run( + cmd, + background=True, + user=user if user is not None else self.config.user, + cwd=cwd if cwd is not None else self.config.working_dir, + envs={**self.config.env_vars, **(env or {})} or None, + timeout=0, + ) + except Exception as exc: + return self._result(command, ExecutionStatus.ERROR, -1, '', self._conn.redact(str(exc)), started) + try: + result = await asyncio.wait_for(handle.wait(), timeout=timeout if timeout and timeout > 0 else None) + return self._result(command, ExecutionStatus.SUCCESS, 0, result.stdout, result.stderr, started) + except CommandExitException as exc: + return self._result(command, ExecutionStatus.ERROR, exc.exit_code, exc.stdout, exc.stderr, started) + except asyncio.TimeoutError: + try: + await handle.kill() + except Exception as exc: + logger.warning(self._conn.redact(f'E2B sandbox {self.id}: process kill failed: {exc}')) + stderr = handle.stderr or f'Command timed out after {timeout} seconds' + return self._result(command, ExecutionStatus.TIMEOUT, -1, handle.stdout, stderr, started) + except Exception as exc: + return self._result(command, ExecutionStatus.ERROR, -1, '', self._conn.redact(str(exc)), started) + + async def put_archive(self, target_dir: str, data: bytes) -> bool: + """Upload a tar archive, extract as the configured user, and remove it.""" + if not posixpath.isabs(target_dir): + raise ValueError(f'target_dir must be an absolute path: {target_dir}') + if self._sbx is None: + raise RuntimeError('E2B sandbox is not running') + tmp = f'/tmp/ms-enclave-upload-{uuid.uuid4().hex}.tar' + try: + await self._sbx.files.write(tmp, data, user=self.config.user) + except Exception as exc: + logger.error(self._conn.redact(f'E2B sandbox {self.id}: upload to {target_dir} failed: {exc}')) + return False + q = shlex.quote + result = await self.execute_command( + f'mkdir -p {q(target_dir)} && tar -xf {q(tmp)} -C {q(target_dir)}; rc=$?; rm -f {q(tmp)}; exit $rc', + timeout=300, + ) + if result.exit_code != 0: + logger.error( + self._conn.redact(f'E2B sandbox {self.id}: extracting into {target_dir} failed: {result.stderr}') + ) + return False + return True diff --git a/ms_enclave/sandbox/boxes/e2b_template.py b/ms_enclave/sandbox/boxes/e2b_template.py new file mode 100644 index 0000000..265bf14 --- /dev/null +++ b/ms_enclave/sandbox/boxes/e2b_template.py @@ -0,0 +1,156 @@ +"""Resolve and build reusable E2B templates from container images.""" + +import asyncio +import time +from dataclasses import dataclass +from typing import Dict, Literal, Mapping, Optional, Set + +from ms_enclave.utils import get_logger + +from ..model import E2BSandboxConfig +from ..utils.e2b_helpers import build_template_name, is_acr_image, normalize_image_ref, rewrite_acr_region +from .e2b_connection import AsyncTemplate, ConnectionConfig, E2BConnection, get_api_client, retry_dispatch + +logger = get_logger() +TemplateSource = Literal['explicit', 'reused', 'built', 'rebuilt'] +_POLL_INTERVAL_S = 10.0 +_POLL_TIMEOUT_S = 3600.0 +_BUSY = {'building', 'waiting'} + + +class E2BTemplateError(RuntimeError): + """A remote template could not be resolved or built.""" + + +@dataclass(frozen=True) +class _RemoteTemplate: + template_id: str + status: str + + +class E2BTemplateResolver: + """Share in-flight builds and cache terminal outcomes within one manager.""" + + def __init__(self, connection: E2BConnection, image_mapping: Mapping[str, str], max_concurrent_builds: int) -> None: + self._conn = connection + self._mapping = {normalize_image_ref(k): v for k, v in image_mapping.items()} + self._build_slots = asyncio.Semaphore(max_concurrent_builds) + self._locks: Dict[str, asyncio.Lock] = {} + self._ready: Set[str] = set() + self._failed: Dict[str, str] = {} + + async def resolve(self, config: E2BSandboxConfig, *, run_id: str) -> str: + """Prefer an explicit template, otherwise reuse or build the resolved image.""" + if config.template: + logger.info(f'E2B template {config.template} source=explicit run_id={run_id}') + return config.template + if not config.image: + raise ValueError('E2B sandbox config requires either `template` or `image`.') + image = self.final_image(config.image) + name = build_template_name( + image, + build_mode=self._conn.build_mode, + cpu_count=config.cpu_count, + memory_mb=config.memory_mb, + disk_size_mb=config.disk_size_mb, + ) + async with self._locks.setdefault(name, asyncio.Lock()): + if name in self._ready: + return name + if name in self._failed: + raise E2BTemplateError(self._failed[name]) + started = time.monotonic() + try: + source = await self._ensure(name, image, config, run_id=run_id) + except Exception as exc: + message = self._conn.redact(f'Failed to build E2B template {name!r} from image {image!r}: {exc}') + self._failed[name] = message + raise E2BTemplateError(message) from None + self._ready.add(name) + logger.info( + self._conn.redact( + f'E2B template {name} source={source} image={config.image} -> {image} ' + f'build_mode={self._conn.build_mode} elapsed={time.monotonic() - started:.1f}s run_id={run_id}' + ) + ) + return name + + def final_image(self, image: str) -> str: + """Map official images on Aliyun and rewrite only an ACR registry region.""" + if self._conn.provider == 'aliyun' and not is_acr_image(image): + mapped = self._mapping.get(normalize_image_ref(image)) + if mapped is None: + raise E2BTemplateError( + self._conn.redact( + f'No converted ACR image for {image!r}; Aliyun Agent Sandbox cannot build templates from ' + 'docker.io images. Add the image to the evalscope mapping or set sandbox.default_config.template.' + ) + ) + image = mapped + return rewrite_acr_region(image, self._conn.region) + + async def _ensure(self, name: str, image: str, config: E2BSandboxConfig, *, run_id: str) -> TemplateSource: + remote = await self._find(name, run_id=run_id) + if remote is not None and remote.status in _BUSY: + remote = await self._wait(name, run_id=run_id) + if remote is not None and remote.status == 'ready': + return 'reused' + source: TemplateSource = 'built' + if remote is not None and remote.status == 'error': + await self._delete(remote.template_id) + source = 'rebuilt' + async with self._build_slots: + await self._build(name, image, config, run_id=run_id) + remote = await self._wait(name, run_id=run_id) + if remote is None or remote.status != 'ready': + raise E2BTemplateError(f'template {name!r} finished with status {remote.status if remote else "missing"}') + return source + + async def _find(self, name: str, *, run_id: str) -> Optional[_RemoteTemplate]: + async def fetch() -> Optional[_RemoteTemplate]: + # Read raw JSON: compatible providers may omit generated-model fields. + client = get_api_client(ConnectionConfig(**self._conn.api_params(retries=0))) + async with client: + response = await client.get_async_httpx_client().get('/templates') + response.raise_for_status() + for item in response.json(): + names = (item.get('aliases') or []) + (item.get('names') or []) + if any(n == name or n.endswith(f'/{name}') for n in names): + return _RemoteTemplate(item['templateID'], item['buildStatus']) + return None + + return await retry_dispatch('template status', fetch, run_id=run_id, redactor=self._conn.redact) + + async def _delete(self, template_id: str) -> None: + client = get_api_client(ConnectionConfig(**self._conn.api_params(retries=0))) + async with client: + response = await client.get_async_httpx_client().delete(f'/templates/{template_id}') + response.raise_for_status() + + async def _build(self, name: str, image: str, config: E2BSandboxConfig, *, run_id: str) -> None: + headers = {} + if self._conn.build_mode == 'micro': + headers = {'X-E2B-Template-Build-Mode': 'micro', 'X-E2B-Template-Alpha-Disk-Size': str(config.disk_size_mb)} + await retry_dispatch( + 'template build', + lambda: AsyncTemplate.build_in_background( + AsyncTemplate().from_image(image), + name, + cpu_count=config.cpu_count, + memory_mb=config.memory_mb, + **self._conn.api_params(retries=0, api_headers=headers), + ), + run_id=run_id, + redactor=self._conn.redact, + ) + + async def _wait(self, name: str, *, run_id: str) -> Optional[_RemoteTemplate]: + deadline = time.monotonic() + _POLL_TIMEOUT_S + while True: + remote = await self._find(name, run_id=run_id) + if remote is not None and remote.status not in _BUSY: + return remote + if time.monotonic() >= deadline: + status = remote.status if remote is not None else 'missing' + raise E2BTemplateError(f'Timed out waiting for E2B template {name!r}; last status={status}') + await asyncio.sleep(_POLL_INTERVAL_S) diff --git a/ms_enclave/sandbox/manager/__init__.py b/ms_enclave/sandbox/manager/__init__.py index 4fd67f7..23aa1fd 100644 --- a/ms_enclave/sandbox/manager/__init__.py +++ b/ms_enclave/sandbox/manager/__init__.py @@ -1,6 +1,7 @@ """Sandbox manager implementations.""" from .base import SandboxManager, SandboxManagerFactory +from .e2b_manager import E2BSandboxManager from .http_manager import HttpSandboxManager from .local_manager import LocalSandboxManager from .volcengine import VolcengineSandboxManager @@ -9,6 +10,7 @@ 'SandboxManager', 'SandboxManagerFactory', 'LocalSandboxManager', + 'E2BSandboxManager', 'HttpSandboxManager', 'VolcengineSandboxManager', ] diff --git a/ms_enclave/sandbox/manager/e2b_manager.py b/ms_enclave/sandbox/manager/e2b_manager.py new file mode 100644 index 0000000..7d3997f --- /dev/null +++ b/ms_enclave/sandbox/manager/e2b_manager.py @@ -0,0 +1,103 @@ +"""E2B resource lifecycle using the standard local manager registry and tools.""" + +import asyncio +import uuid +from typing import Any, Dict, List, Optional, Union + +from ms_enclave.utils import get_logger + +from ..boxes.e2b_connection import E2BConnection, require_e2b +from ..boxes.e2b_sandbox import E2BSandbox +from ..boxes.e2b_template import E2BTemplateResolver +from ..model import ( + E2BSandboxConfig, + E2BSandboxManagerConfig, + SandboxConfig, + SandboxManagerConfig, + SandboxManagerType, + SandboxType, +) +from .base import register_manager +from .local_manager import LocalSandboxManager + +logger = get_logger() + + +@register_manager(SandboxManagerType.E2B) +class E2BSandboxManager(LocalSandboxManager): + """Limit live sandboxes and share template resolution within one run.""" + + def __init__(self, config: Optional[Union[SandboxManagerConfig, Dict[str, Any]]] = None, **kwargs: Any) -> None: + require_e2b() + cfg = ( + config + if isinstance(config, E2BSandboxManagerConfig) + else E2BSandboxManagerConfig.model_validate( + config.model_dump(exclude_unset=True) if isinstance(config, SandboxManagerConfig) else (config or {}) + ) + ) + super().__init__(cfg, **kwargs) + self.config: E2BSandboxManagerConfig = cfg + self.run_id = cfg.run_id or uuid.uuid4().hex[:12] + self._conn = E2BConnection.from_config(cfg) + self._templates = E2BTemplateResolver(self._conn, cfg.image_mapping, cfg.max_concurrent_builds) + self._slots = asyncio.Semaphore(cfg.max_sandboxes) + + async def start(self) -> None: + """Start standard lifecycle management and log the run's cleanup identifier.""" + await super().start() + logger.info( + f'E2B manager started run_id={self.run_id} provider={self._conn.provider} ' + f'region={self._conn.region} build_mode={self._conn.build_mode}' + ) + + async def create_sandbox( + self, + sandbox_type: SandboxType, + config: Optional[Union[SandboxConfig, Dict[str, Any]]] = None, + sandbox_id: Optional[str] = None, + ) -> str: + """Resolve templates before reserving a live sandbox slot.""" + if sandbox_type != SandboxType.E2B: + raise ValueError(f'E2BSandboxManager only supports SandboxType.E2B, got {sandbox_type}') + cfg = ( + config + if isinstance(config, E2BSandboxConfig) + else E2BSandboxConfig.model_validate( + config.model_dump(exclude_unset=True) if isinstance(config, SandboxConfig) else (config or {}) + ) + ) + await self._templates.resolve(cfg, run_id=self.run_id) + await self._slots.acquire() + sandbox = E2BSandbox(cfg, sandbox_id, connection=self._conn, templates=self._templates, run_id=self.run_id) + try: + await sandbox.start() + except BaseException: + self._slots.release() + raise + self._sandboxes[sandbox.id] = sandbox + return sandbox.id + + async def delete_sandbox(self, sandbox_id: str) -> bool: + """Remove registration before kill so concurrent deletes are idempotent.""" + sandbox = self._sandboxes.pop(sandbox_id, None) + if sandbox is None: + return False + try: + await sandbox.stop() + finally: + self._sandbox_locks.pop(sandbox_id, None) + self._slots.release() + return True + + async def initialize_pool( + self, + pool_size: Optional[int] = None, + sandbox_type: Optional[SandboxType] = None, + config: Optional[Union[SandboxConfig, Dict[str, Any]]] = None, + ) -> List[str]: + """Reject pools that cannot fit within the live sandbox limit.""" + size = pool_size or self.config.pool_size + if size > self.config.max_sandboxes: + raise ValueError(f'pool_size={size} exceeds E2B max_sandboxes={self.config.max_sandboxes}') + return await super().initialize_pool(pool_size, sandbox_type or SandboxType.E2B, config) diff --git a/ms_enclave/sandbox/model/__init__.py b/ms_enclave/sandbox/model/__init__.py index b9c2d44..b7b0edf 100644 --- a/ms_enclave/sandbox/model/__init__.py +++ b/ms_enclave/sandbox/model/__init__.py @@ -4,6 +4,8 @@ from .config import ( DockerNotebookConfig, DockerSandboxConfig, + E2BSandboxConfig, + E2BSandboxManagerConfig, FileOperationConfig, PythonExecutorConfig, SandboxConfig, diff --git a/ms_enclave/sandbox/model/base.py b/ms_enclave/sandbox/model/base.py index c419a68..51b038a 100644 --- a/ms_enclave/sandbox/model/base.py +++ b/ms_enclave/sandbox/model/base.py @@ -23,6 +23,7 @@ class SandboxType(str, Enum): DOCKER_NOTEBOOK = 'docker_notebook' VOLCENGINE = 'volcengine' DUMMY = 'dummy' + E2B = 'e2b' @classmethod def get_compatible_types(cls, sandbox_type: 'SandboxType') -> Set['SandboxType']: @@ -69,6 +70,7 @@ class SandboxManagerType(str, Enum): """Sandbox manager type enumeration.""" LOCAL = 'local' + E2B = 'e2b' HTTP = 'http' VOLCENGINE = 'volcengine' diff --git a/ms_enclave/sandbox/model/config.py b/ms_enclave/sandbox/model/config.py index 3cd4e84..dc544e3 100644 --- a/ms_enclave/sandbox/model/config.py +++ b/ms_enclave/sandbox/model/config.py @@ -1,8 +1,11 @@ """Configuration data models.""" +import os from typing import Any, Dict, List, Optional, Tuple, Union -from pydantic import BaseModel, Field, field_validator, model_validator +from pydantic import BaseModel, ConfigDict, Field, SecretStr, field_validator, model_validator + +from ..utils.e2b_helpers import BuildMode, E2BProvider, detect_provider, infer_region class SandboxManagerConfig(BaseModel): @@ -60,6 +63,7 @@ class DockerSandboxConfig(SandboxConfig): """Docker-specific sandbox configuration.""" image: str = Field('python:3.11-slim', description='Docker image name') + cap_add: List[str] = Field(default_factory=list, description='Additional Linux capabilities') entrypoint: Optional[Union[str, List[str]]] = Field( None, description='Container entrypoint override. Use an empty list to clear the image entrypoint.', @@ -86,8 +90,8 @@ class DockerSandboxConfig(SandboxConfig): description='Docker platform (e.g. "linux/amd64"). Required when the image arch differs from the host ' '(e.g. amd64 image on Apple Silicon).', ) - memory_limit: str = Field(default='4g', description='Memory limit') - cpu_limit: float = Field(default=2.0, description='CPU limit') + memory_limit: Optional[str] = Field(default='4g', description='Memory limit; None means unrestricted') + cpu_limit: Optional[float] = Field(default=2.0, description='CPU limit; None means unrestricted') network_enabled: bool = Field(default=True, description='Enable network access') privileged: bool = Field(default=False, description='Run in privileged mode') remove_on_exit: bool = Field(default=True, description='Remove container on exit') @@ -106,6 +110,8 @@ class DockerSandboxConfig(SandboxConfig): @field_validator('memory_limit') def validate_memory_limit(cls, v): """Validate memory limit format.""" + if v is None: + return v if not isinstance(v, str): raise ValueError('Memory limit must be a string') # Basic validation for memory format (e.g., '512m', '1g', '2G') @@ -118,11 +124,76 @@ def validate_memory_limit(cls, v): @field_validator('cpu_limit') def validate_cpu_limit(cls, v): """Validate CPU limit.""" - if v <= 0: + if v is not None and v <= 0: raise ValueError('CPU limit must be positive') return v +class E2BSandboxManagerConfig(SandboxManagerConfig): + """Connection and concurrency settings for an E2B-compatible service.""" + + model_config = ConfigDict(extra='forbid', hide_input_in_errors=True) + + api_key: Optional[SecretStr] = None + api_url: Optional[str] = None + domain: Optional[str] = None + max_sandboxes: int = Field(default=16, gt=0) + max_concurrent_builds: int = Field(default=4, gt=0) + build_mode: Optional[BuildMode] = None + image_region: Optional[str] = None + image_mapping: Dict[str, str] = Field(default_factory=dict, repr=False) + run_id: Optional[str] = None + + @model_validator(mode='after') + def _validate_connection(self) -> 'E2BSandboxManagerConfig': + if self.api_key is None or not self.api_key.get_secret_value(): + self.api_key = SecretStr(os.environ.get('E2B_API_KEY', '')) + if not self.api_key.get_secret_value(): + raise ValueError( + 'E2B api_key is missing: set sandbox.manager_config.api_key or the E2B_API_KEY environment variable.' + ) + if self.provider != 'aliyun' and self.build_mode == 'micro': + raise ValueError("build_mode='micro' is only supported by Aliyun Agent Sandbox.") + return self + + @property + def provider(self) -> E2BProvider: + """Return the provider inferred from the endpoint.""" + return detect_provider(self.domain, self.api_url) + + @property + def target_region(self) -> Optional[str]: + """Return an explicit or endpoint-derived image region.""" + return self.image_region or infer_region(self.domain, self.api_url) + + @property + def resolved_build_mode(self) -> BuildMode: + """Return the provider's default build mode unless overridden.""" + return self.build_mode or ('micro' if self.provider == 'aliyun' else 'container') + + +class E2BSandboxConfig(SandboxConfig): + """Image or template and fixed lifetime for one remote sandbox.""" + + model_config = ConfigDict(extra='forbid', hide_input_in_errors=True) + + image: Optional[str] = None + template: Optional[str] = None + user: str = 'user' + working_dir: str = '/home/user' + sandbox_timeout: int = Field(default=3600, gt=0) + cpu_count: int = Field(default=2, gt=0) + memory_mb: int = Field(default=4096, gt=0) + disk_size_mb: int = Field(default=15360, gt=0) + allow_internet_access: bool = True + + @model_validator(mode='after') + def _validate_source(self) -> 'E2BSandboxConfig': + if not self.image and not self.template: + raise ValueError('E2B sandbox config requires either `template` or `image`.') + return self + + class DockerNotebookConfig(DockerSandboxConfig): """Docker Notebook-specific sandbox configuration.""" diff --git a/ms_enclave/sandbox/tools/sandbox_tools/python_executor.py b/ms_enclave/sandbox/tools/sandbox_tools/python_executor.py index 5eb1a4e..b2eaacf 100644 --- a/ms_enclave/sandbox/tools/sandbox_tools/python_executor.py +++ b/ms_enclave/sandbox/tools/sandbox_tools/python_executor.py @@ -9,13 +9,13 @@ from ms_enclave.sandbox.tools.tool_info import ToolParams if TYPE_CHECKING: - from ms_enclave.sandbox.boxes import DockerSandbox, VolcengineSandbox + from ms_enclave.sandbox.boxes import Sandbox, VolcengineSandbox @register_tool('python_executor') class PythonExecutor(SandboxTool): _name = 'python_executor' - _sandbox_types = [SandboxType.DOCKER, SandboxType.VOLCENGINE] + _sandbox_types = [SandboxType.DOCKER, SandboxType.VOLCENGINE, SandboxType.E2B] _description = 'Execute Python code in an isolated environment using IPython' _parameters = ToolParams( type='object', @@ -26,7 +26,7 @@ class PythonExecutor(SandboxTool): required=['code'], ) - async def execute(self, sandbox_context: 'DockerSandbox', code: str, timeout: Optional[int] = 30) -> ToolResult: + async def execute(self, sandbox_context: 'Sandbox', code: str, timeout: Optional[int] = 30) -> ToolResult: """Execute Python code by writing to a temporary file and executing it.""" if not code.strip(): diff --git a/ms_enclave/sandbox/tools/sandbox_tools/shell_executor.py b/ms_enclave/sandbox/tools/sandbox_tools/shell_executor.py index 2be02bb..ef97d8a 100644 --- a/ms_enclave/sandbox/tools/sandbox_tools/shell_executor.py +++ b/ms_enclave/sandbox/tools/sandbox_tools/shell_executor.py @@ -15,7 +15,7 @@ @register_tool('shell_executor') class ShellExecutor(SandboxTool): _name = 'shell_executor' - _sandbox_types = [SandboxType.DOCKER, SandboxType.VOLCENGINE] + _sandbox_types = [SandboxType.DOCKER, SandboxType.VOLCENGINE, SandboxType.E2B] _description = 'Execute shell commands in an isolated environment' _parameters = ToolParams( type='object', @@ -38,7 +38,7 @@ class ShellExecutor(SandboxTool): async def execute( self, sandbox_context: 'Sandbox', command: Union[str, List[str]], timeout: Optional[int] = 30 ) -> ToolResult: - """Execute shell command in the Docker container.""" + """Execute a shell command in the sandbox.""" if not command or (isinstance(command, str) and not command.strip()): return ToolResult(tool_name=self.name, status=ExecutionStatus.ERROR, output='', error='No command provided') diff --git a/ms_enclave/sandbox/utils/e2b_helpers.py b/ms_enclave/sandbox/utils/e2b_helpers.py new file mode 100644 index 0000000..5dd02da --- /dev/null +++ b/ms_enclave/sandbox/utils/e2b_helpers.py @@ -0,0 +1,66 @@ +"""Pure image, endpoint and secret handling for E2B-compatible services.""" + +import hashlib +import re +from typing import Literal, Optional +from urllib.parse import urlparse + +E2BProvider = Literal['e2b_official', 'aliyun'] +BuildMode = Literal['micro', 'container'] +TEMPLATE_NAME_PREFIX = 'evalscope' +TEMPLATE_NAME_MAX_LEN = 128 +ACR_HOST_PATTERN = re.compile(r'^fc-e2b-registry\.(?P[a-z0-9-]+)\.cr\.aliyuncs\.com/') +_REGION_LABEL = re.compile(r'[a-z]{2}-[a-z]+(?:-\d+)?') +REDACTED = '***' + + +def detect_provider(domain: Optional[str], api_url: Optional[str]) -> E2BProvider: + """Identify Aliyun endpoints by their host suffix.""" + hosts = (domain or '', (urlparse(api_url).hostname or '') if api_url else '') + return 'aliyun' if any(h.lower().endswith('aliyuncs.com') for h in hosts) else 'e2b_official' + + +def infer_region(domain: Optional[str], api_url: Optional[str]) -> Optional[str]: + """Find a region DNS label in a new or legacy Aliyun endpoint.""" + for host in (domain, urlparse(api_url).hostname if api_url else None): + if host and host.lower().endswith('aliyuncs.com'): + for label in host.lower().split('.'): + if _REGION_LABEL.fullmatch(label): + return label + return None + + +def is_acr_image(image: str) -> bool: + """Return whether the image uses the Agent Sandbox ACR registry.""" + return ACR_HOST_PATTERN.match(image) is not None + + +def rewrite_acr_region(image: str, region: Optional[str]) -> str: + """Replace only the region segment of an ACR registry host.""" + if not region: + return image + return ACR_HOST_PATTERN.sub(f'fc-e2b-registry.{region}.cr.aliyuncs.com/', image, count=1) + + +def normalize_image_ref(ref: str) -> str: + """Add the implicit docker.io registry without changing repository paths.""" + ref = ref.strip() + first, sep, _ = ref.partition('/') + if sep and ('.' in first or ':' in first or first == 'localhost'): + return ref + return f'docker.io/{ref}' + + +def build_template_name(image: str, *, build_mode: BuildMode, cpu_count: int, memory_mb: int, disk_size_mb: int) -> str: + """Build a deterministic template name including all build resource settings.""" + payload = '\n'.join([image, build_mode, str(cpu_count), str(memory_mb), str(disk_size_mb)]) + digest = hashlib.sha256(payload.encode('utf-8')).hexdigest()[:8] + base = image.rsplit('/', 1)[-1].split('@', 1)[0].split(':', 1)[0] + label = re.sub(r'[^a-z0-9_-]+', '-', base.lower()).strip('-') or 'image' + budget = TEMPLATE_NAME_MAX_LEN - len(TEMPLATE_NAME_PREFIX) - len(digest) - 2 + return f'{TEMPLATE_NAME_PREFIX}-{label[:budget]}-{digest}' + + +def redact(text: str, secret: Optional[str]) -> str: + """Mask every occurrence of a nonempty secret.""" + return text.replace(secret, REDACTED) if secret else text diff --git a/pyproject.toml b/pyproject.toml index 872e94c..161659f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,6 +62,7 @@ build-backend = "setuptools.build_meta" [project.optional-dependencies] dev = [ + "hypothesis==6.155.2", "pytest>=8.3.5", "pytest-asyncio>=0.23.0", "pre-commit>=4.2.0", @@ -70,6 +71,8 @@ dev = [ "ruff>=0.16.4", ] +e2b = ["e2b>=2.63,<3"] + docker = [ "docker>=7.1.0", "websocket-client" diff --git a/tests/test_e2b_helpers.py b/tests/test_e2b_helpers.py new file mode 100644 index 0000000..e99b146 --- /dev/null +++ b/tests/test_e2b_helpers.py @@ -0,0 +1,81 @@ +"""Property checks for image identity and endpoint handling.""" + +import hashlib +import re + +from hypothesis import given, settings +from hypothesis import strategies as st + +from ms_enclave.sandbox.model import E2BSandboxManagerConfig +from ms_enclave.sandbox.utils.e2b_helpers import ( + build_template_name, + detect_provider, + infer_region, + normalize_image_ref, + rewrite_acr_region, +) + +regions = st.from_regex(r'[a-z]{2}-[a-z]{1,12}(-[0-9]{1,2})?', fullmatch=True) +text = st.text(alphabet=st.characters(blacklist_categories=('Cs',))) + + +# Feature: e2b-sandbox, Property 1: Provider depends only on the host suffix. +@settings(max_examples=100) +@given(st.text(alphabet='abcdefghijklmnopqrstuvwxyz0123456789', min_size=1, max_size=15), + st.one_of(st.none(), st.booleans()), st.one_of(st.none(), st.booleans()), st.integers(0, 65535)) +def test_provider(prefix: str, domain_cloud: bool | None, url_cloud: bool | None, case_mask: int) -> None: + def mixed_case(host: str) -> str: + return ''.join(char.upper() if case_mask & (1 << (i % 16)) else char for i, char in enumerate(host)) + domain = mixed_case(prefix + ('.aliyuncs.com' if domain_cloud else '.example.com')) if domain_cloud is not None else None + url = 'https://' + mixed_case(prefix + ('.aliyuncs.com' if url_cloud else '.example.com')) + '/aliyuncs.com' if url_cloud is not None else None + expected = domain_cloud is True or url_cloud is True + assert detect_provider(domain, url) == ('aliyun' if expected else 'e2b_official') + + +# Feature: e2b-sandbox, Property 2: New and legacy endpoints yield the region; overrides win. +@settings(max_examples=100) +@given(regions, regions) +def test_region(region: str, override: str) -> None: + for kwargs in ({'domain': f'{region}.sandbox.aliyuncs.com'}, + {'domain': f'id.{region}.e2b.fc.aliyuncs.com'}, + {'api_url': f'https://api.{region}.sandbox.aliyuncs.com'}): + cfg = E2BSandboxManagerConfig(api_key='placeholder', **kwargs) + assert cfg.target_region == region + assert E2BSandboxManagerConfig(api_key='placeholder', image_region=override, **kwargs).target_region == override + assert infer_region('ap-southeast-1.sandbox.aliyuncs.com', None) == 'ap-southeast-1' + assert infer_region('e2b.app', None) is None + + +# Feature: e2b-sandbox, Property 3: Rewrite changes only the ACR host region and is idempotent. +@settings(max_examples=100) +@given(regions, regions, text) +def test_region_rewrite(source: str, target: str, path: str) -> None: + image = f'fc-e2b-registry.{source}.cr.aliyuncs.com/{path}' + expected = f'fc-e2b-registry.{target}.cr.aliyuncs.com/{path}' + assert rewrite_acr_region(image, target) == expected + assert rewrite_acr_region(expected, target) == expected + assert rewrite_acr_region(image, None) == image + assert rewrite_acr_region(f'docker.io/{path}', target) == f'docker.io/{path}' + + +# Feature: e2b-sandbox, Property 5: Template names are bounded, deterministic and sensitive to build inputs. +@settings(max_examples=100) +@given(text, st.sampled_from(['micro', 'container']), st.integers(1, 64), st.integers(1, 65536), st.integers(1, 99999)) +def test_template_name(image: str, mode: str, cpu: int, memory: int, disk: int) -> None: + kwargs = dict(build_mode=mode, cpu_count=cpu, memory_mb=memory, disk_size_mb=disk) + name = build_template_name(image, **kwargs) + assert re.fullmatch(r'evalscope-[a-z0-9_-]+-[0-9a-f]{8}', name) + assert len(name) <= 128 + payload = '\n'.join([image, mode, str(cpu), str(memory), str(disk)]) + assert name.endswith(hashlib.sha256(payload.encode()).hexdigest()[:8]) + assert build_template_name(image, **kwargs) == name + for key, value in [('build_mode', 'micro' if mode == 'container' else 'container'), + ('cpu_count', cpu + 1), ('memory_mb', memory + 1), ('disk_size_mb', disk + 1)]: + assert build_template_name(image, **(kwargs | {key: value}))[-8:] != name[-8:] + assert build_template_name(image + 'x', **kwargs)[-8:] != name[-8:] + + +def test_image_normalization() -> None: + assert normalize_image_ref(' swebench/task:latest ') == 'docker.io/swebench/task:latest' + for image in ['docker.io/swebench/task:latest', 'localhost:5000/repo:t', 'ghcr.io/org/image:t']: + assert normalize_image_ref(image) == image diff --git a/tests/test_e2b_sandbox.py b/tests/test_e2b_sandbox.py new file mode 100644 index 0000000..15c302b --- /dev/null +++ b/tests/test_e2b_sandbox.py @@ -0,0 +1,559 @@ +"""E2B backend tests use fakes and never contact a provider.""" + +import asyncio +import io +import re +import shlex +import subprocess +import sys +import tarfile +import tempfile +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest +from hypothesis import given, settings +from hypothesis import strategies as st +from pydantic import ValidationError + +from ms_enclave.sandbox.boxes import DockerSandbox, SandboxFactory +from ms_enclave.sandbox.boxes import e2b_connection as connection_module +from ms_enclave.sandbox.boxes import e2b_sandbox as sandbox_module +from ms_enclave.sandbox.boxes import e2b_template as template_module +from ms_enclave.sandbox.boxes.e2b_connection import CommandExitException, E2BConnection, retry_dispatch +from ms_enclave.sandbox.boxes.e2b_sandbox import E2BSandbox +from ms_enclave.sandbox.boxes.e2b_template import E2BTemplateError, E2BTemplateResolver +from ms_enclave.sandbox.manager import E2BSandboxManager +from ms_enclave.sandbox.model import ( + DockerSandboxConfig, + E2BSandboxConfig, + E2BSandboxManagerConfig, + ExecutionStatus, + SandboxType, +) +from ms_enclave.sandbox.utils.e2b_helpers import redact + + +class FakeHandle: + def __init__(self, code: int = 0, stdout: str = 'out', stderr: str = '') -> None: + self.stdout, self.stderr, self.code = stdout, stderr, code + self.hang = False + self.kill = AsyncMock() + + async def wait(self) -> SimpleNamespace: + if self.hang: + await asyncio.Event().wait() + if self.code: + raise CommandExitException(stderr=self.stderr, stdout=self.stdout, exit_code=self.code, error=None) + return SimpleNamespace(stdout=self.stdout, stderr=self.stderr) + + +class FakeCommands: + def __init__(self) -> None: + self.handle = FakeHandle() + self.run = AsyncMock(side_effect=lambda *a, **k: self.handle) + + +class FakeFiles: + def __init__(self) -> None: + self.write = AsyncMock() + + +class FakeAsyncSandbox: + def __init__(self) -> None: + self.sandbox_id = 'remote-id' + self.commands = FakeCommands() + self.files = FakeFiles() + self.kill = AsyncMock() + + +class FakeTemplateApi: + def __init__(self) -> None: + self.rows = [] + self.deleted = [] + self.error = None + + async def __aenter__(self) -> 'FakeTemplateApi': + return self + + async def __aexit__(self, *args) -> None: + pass + + def get_async_httpx_client(self) -> 'FakeTemplateApi': + return self + + async def get(self, path: str) -> httpx.Response: + if self.error: + raise self.error + return httpx.Response(200, json=self.rows, request=httpx.Request('GET', 'https://example.com/templates')) + + async def delete(self, path: str) -> httpx.Response: + self.deleted.append(path) + self.rows = [] + return httpx.Response(204, request=httpx.Request('DELETE', 'https://example.com' + path)) + + +class FakeAsyncTemplate: + def from_image(self, image: str) -> str: + return image + + +@pytest.fixture +def sdk(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace: + api = FakeTemplateApi() + remote = FakeAsyncSandbox() + create = AsyncMock(return_value=remote) + + async def build(image: str, name: str, **kwargs) -> None: + api.rows = [{'templateID': 'template-id', 'aliases': [name], 'buildStatus': 'ready'}] + + builder = AsyncMock(side_effect=build) + monkeypatch.setattr(sandbox_module, 'AsyncSandbox', SimpleNamespace(create=create)) + monkeypatch.setattr(template_module, 'get_api_client', lambda *a, **k: api) + monkeypatch.setattr(template_module, 'AsyncTemplate', FakeAsyncTemplate) + monkeypatch.setattr(FakeAsyncTemplate, 'build_in_background', builder, raising=False) + monkeypatch.setattr(template_module, '_POLL_INTERVAL_S', 0) + return SimpleNamespace(api=api, remote=remote, create=create, build=builder) + + +def connection(**kwargs) -> E2BConnection: + return E2BConnection.from_config(E2BSandboxManagerConfig(api_key='test-secret', **kwargs)) + + +def running_sandbox(handle: FakeHandle | None = None) -> E2BSandbox: + box = E2BSandbox(E2BSandboxConfig(template='explicit'), connection=connection()) + box._sbx = FakeAsyncSandbox() + if handle: + box._sbx.commands.handle = handle + return box + + +def test_manager_config(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv('E2B_API_KEY', raising=False) + with pytest.raises(ValidationError, match='api_key.*E2B_API_KEY'): + E2BSandboxManagerConfig() + monkeypatch.setenv('E2B_API_KEY', 'environment-key') + cfg = E2BSandboxManagerConfig() + assert cfg.api_key.get_secret_value() == 'environment-key' + assert cfg.api_url is cfg.domain is None + assert cfg.resolved_build_mode == 'container' + assert 'environment-key' not in repr(cfg) + cfg = E2BSandboxManagerConfig(domain='cn-hangzhou.sandbox.aliyuncs.com', image_mapping={'hidden': 'value'}) + assert cfg.resolved_build_mode == 'micro' + assert 'hidden' not in repr(cfg) + with pytest.raises(ValidationError, match='only supported by Aliyun'): + E2BSandboxManagerConfig(build_mode='micro') + for field in ['provider', 'request_timeout', 'metadata', 'template_label', 'unknown']: + with pytest.raises(ValidationError, match=field): + E2BSandboxManagerConfig(**{field: 'bad'}) + for field in ['max_sandboxes', 'max_concurrent_builds']: + for value in [0, -1]: + with pytest.raises(ValidationError, match=field): + E2BSandboxManagerConfig(**{field: value}) + + +def test_sandbox_config() -> None: + cfg = E2BSandboxConfig(image='example:image') + assert (cfg.sandbox_timeout, cfg.cpu_count, cfg.memory_mb, cfg.disk_size_mb) == (3600, 2, 4096, 15360) + assert cfg.user == 'user' and cfg.working_dir == '/home/user' + for field in ['sandbox_timeout', 'cpu_count', 'memory_mb', 'disk_size_mb']: + for value in [0, -1]: + with pytest.raises(ValidationError, match=field): + E2BSandboxConfig(image='image', **{field: value}) + with pytest.raises(ValidationError, match='template.*image'): + E2BSandboxConfig() + with pytest.raises(ValidationError, match='unknown'): + E2BSandboxConfig(template='existing', unknown=True) + + +def test_invalid_manager_config_masks_key() -> None: + with pytest.raises(ValidationError) as error: + E2BSandboxManagerConfig(api_key='plaintext-secret', build_mode='micro') + assert 'plaintext-secret' not in str(error.value) + + +@pytest.mark.parametrize('capabilities', [[], ['SYS_ADMIN']]) +async def test_docker_capabilities(capabilities: list[str]) -> None: + sandbox = DockerSandbox(DockerSandboxConfig(cap_add=capabilities)) + sandbox.client = MagicMock() + sandbox._run_blocking = AsyncMock(return_value=MagicMock()) + try: + await sandbox._create_container() + kwargs = sandbox._run_blocking.await_args.kwargs + if capabilities: + assert kwargs['cap_add'] == capabilities + else: + assert 'cap_add' not in kwargs + finally: + sandbox._executor.shutdown(wait=False, cancel_futures=True) + + +async def test_docker_unrestricted_resources() -> None: + defaults = DockerSandboxConfig() + assert defaults.memory_limit == '4g' and defaults.cpu_limit == 2 + sandbox = DockerSandbox(DockerSandboxConfig(memory_limit=None, cpu_limit=None)) + sandbox.client = MagicMock() + sandbox._run_blocking = AsyncMock(return_value=MagicMock()) + try: + await sandbox._create_container() + assert not {'mem_limit', 'cpu_quota', 'cpu_period'} & sandbox._run_blocking.await_args.kwargs.keys() + finally: + sandbox._executor.shutdown(wait=False, cancel_futures=True) + + +@pytest.mark.parametrize('failures', [0, 1, 2, 3]) +async def test_retry_connection_failures(failures: int, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(connection_module.asyncio, 'sleep', AsyncMock()) + call = AsyncMock(side_effect=[ConnectionError('test-secret')] * failures + ['ok']) + if failures == 3: + with pytest.raises(RuntimeError, match='run_id=run') as error: + await retry_dispatch('create', call, run_id='run', redactor=connection().redact) + assert 'test-secret' not in str(error.value) + else: + assert await retry_dispatch('create', call, run_id='run', redactor=connection().redact) == 'ok' + assert call.call_count == min(failures + 1, 3) + + +async def test_retry_status_and_read_timeout(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(connection_module.asyncio, 'sleep', AsyncMock()) + response = httpx.Response(429, request=httpx.Request('GET', 'https://example.com')) + call = AsyncMock(side_effect=[httpx.HTTPStatusError('limited', request=response.request, response=response), 'ok']) + assert await retry_dispatch('status', call, run_id='r', redactor=str) == 'ok' + call = AsyncMock(side_effect=httpx.ReadTimeout('unknown dispatch state')) + with pytest.raises(RuntimeError, match='unknown dispatch state'): + await retry_dispatch('create', call, run_id='r', redactor=str) + assert call.call_count == 1 + assert 'api_url' not in connection().api_params() + assert 'domain' not in connection().api_params() + + +# Feature: e2b-sandbox, Property 4: Final images match provider-specific mapping and region rules. +@settings(max_examples=100) +@given(st.booleans(), st.booleans(), st.booleans(), st.from_regex('[a-z]{1,12}', fullmatch=True)) +def test_final_image(aliyun: bool, mapped: bool, prefix: bool, label: str) -> None: + source = f'org/{label}:latest' + acr = f'fc-e2b-registry.cn-shanghai.cr.aliyuncs.com/repo/{label}:v1' + conn = connection(domain='cn-hangzhou.sandbox.aliyuncs.com' if aliyun else None) + resolver = E2BTemplateResolver(conn, {('docker.io/' if prefix else '') + source: acr} if mapped else {}, 2) + if aliyun and not mapped: + with pytest.raises(E2BTemplateError, match=source): + resolver.final_image(source) + else: + expected = acr.replace('cn-shanghai', 'cn-hangzhou') if aliyun else source + assert resolver.final_image(source) == expected + + +# Feature: e2b-sandbox, Property 6: Concurrent same-name resolution shares one terminal build outcome. +@settings(max_examples=100) +@given(st.integers(1, 16), st.sampled_from([None, 'ready', 'building', 'waiting', 'error']), st.booleans()) +def test_shared_template_resolution(count: int, initial: str | None, success: bool) -> None: + async def run() -> None: + resolver = E2BTemplateResolver(connection(), {}, 2) + remote = template_module._RemoteTemplate('id', initial) if initial else None + resolver._find = AsyncMock(return_value=remote) + resolver._delete = AsyncMock() + resolver._build = AsyncMock() + resolver._wait = AsyncMock(return_value=template_module._RemoteTemplate('id', 'ready' if success else 'error')) + cfg = E2BSandboxConfig(image='org/image') + results = await asyncio.gather(*(resolver.resolve(cfg, run_id='r') for _ in range(count)), return_exceptions=True) + failed = initial != 'ready' and not success + assert all(isinstance(r, E2BTemplateError) == failed for r in results) + assert len({str(r) for r in results}) == 1 + should_build = initial not in ['ready', 'building', 'waiting'] or (initial in ['building', 'waiting'] and not success) + assert resolver._build.call_count == int(should_build) + assert resolver._delete.call_count == int(initial == 'error' or (initial in ['building', 'waiting'] and not success)) + if failed: + with pytest.raises(E2BTemplateError): + await resolver.resolve(cfg, run_id='r') + assert resolver._build.call_count == int(should_build) + asyncio.run(run()) + + +async def test_template_api_and_headers(sdk: SimpleNamespace) -> None: + for domain, mode in [(None, 'container'), ('id.cn-shanghai.e2b.fc.aliyuncs.com', 'micro')]: + resolver = E2BTemplateResolver(connection(domain=domain), {}, 2) + cfg = E2BSandboxConfig(image='fc-e2b-registry.cn-shanghai.cr.aliyuncs.com/repo:i', disk_size_mb=12345) + name = await resolver.resolve(cfg, run_id='r') + headers = sdk.build.call_args.kwargs['api_headers'] + assert headers == ({'X-E2B-Template-Build-Mode': 'micro', 'X-E2B-Template-Alpha-Disk-Size': '12345'} if mode == 'micro' else {}) + assert sdk.build.call_args.kwargs['retries'] == 0 + fresh = E2BTemplateResolver(connection(domain=domain), {}, 2) + sdk.api.rows = [{'templateID': 'id', 'names': [f'team/{name}'], 'buildStatus': 'ready'}] + before = sdk.build.call_count + assert await fresh.resolve(cfg, run_id='r') == name + assert sdk.build.call_count == before + assert await resolver.resolve(E2BSandboxConfig(template='explicit'), run_id='r') == 'explicit' + + +async def test_template_poll_timeout(monkeypatch: pytest.MonkeyPatch) -> None: + resolver = E2BTemplateResolver(connection(), {}, 1) + resolver._find = AsyncMock(return_value=template_module._RemoteTemplate('id', 'waiting')) + monkeypatch.setattr(template_module, '_POLL_TIMEOUT_S', 0) + with pytest.raises(E2BTemplateError, match='name.*waiting'): + await resolver._wait('name', run_id='r') + + +async def test_template_eventual_visibility(monkeypatch: pytest.MonkeyPatch) -> None: + resolver = E2BTemplateResolver(connection(), {}, 1) + ready = template_module._RemoteTemplate('id', 'ready') + resolver._find = AsyncMock(side_effect=[None, template_module._RemoteTemplate('id', 'building'), ready]) + monkeypatch.setattr(template_module, '_POLL_INTERVAL_S', 0) + assert await resolver._wait('name', run_id='r') == ready + + +async def test_build_concurrency_limit() -> None: + resolver = E2BTemplateResolver(connection(), {}, 2) + resolver._find = AsyncMock(return_value=None) + active = peak = 0 + + async def build(*args, **kwargs) -> None: + nonlocal active, peak + active += 1 + peak = max(peak, active) + await asyncio.sleep(0) + + async def wait(*args, **kwargs): + nonlocal active + await asyncio.sleep(0) + active -= 1 + return template_module._RemoteTemplate('id', 'ready') + + resolver._build, resolver._wait = build, wait + await asyncio.gather(*(resolver.resolve(E2BSandboxConfig(image=f'org/{i}'), run_id='r') for i in range(8))) + assert peak == 2 and active == 0 + + +# Feature: e2b-sandbox, Property 7: POSIX list command quoting round-trips without argument loss. +@settings(max_examples=100) +@given(st.lists(st.text(alphabet=st.characters(blacklist_characters='\x00', blacklist_categories=('Cs',))), max_size=8)) +def test_command_quoting(argv: list[str]) -> None: + async def run() -> None: + box = running_sandbox() + await box.execute_command(argv) + assert shlex.split(box._sbx.commands.run.call_args.args[0]) == argv + asyncio.run(run()) + + +# Feature: e2b-sandbox, Property 8: Command output and exit codes are preserved without redispatch. +@settings(max_examples=100) +@given(st.integers(0, 255), st.text(), st.text()) +def test_command_result(code: int, stdout: str, stderr: str) -> None: + async def run() -> None: + box = running_sandbox(FakeHandle(code, stdout, stderr)) + result = await box.execute_command('command') + assert (result.exit_code, result.stdout, result.stderr) == (code, stdout, stderr) + assert result.status == (ExecutionStatus.ERROR if code else ExecutionStatus.SUCCESS) + assert box._sbx.commands.run.call_count == 1 + asyncio.run(run()) + + +# Feature: e2b-sandbox, Property 9: Uploaded directories preserve relative paths and bytes. +@settings(max_examples=100) +@given(st.dictionaries(st.lists(st.from_regex('[a-z]{1,5}', fullmatch=True), min_size=1, max_size=3).map('/'.join), + st.binary(max_size=100), max_size=8)) +def test_put_dir_roundtrip(files: dict[str, bytes]) -> None: + # Prefix file names so generated file paths cannot also be parent directories. + files = {str(Path(path).parent / ('file-' + Path(path).name)): data for path, data in files.items()} + async def run() -> None: + with tempfile.TemporaryDirectory() as root: + source, target = Path(root) / 'source', Path(root) / 'target' + source.mkdir() + for path, data in files.items(): + dest = source / path + dest.parent.mkdir(parents=True, exist_ok=True) + dest.write_bytes(data) + box = running_sandbox() + assert await box.put_dir(source, '/target') + archive = box._sbx.files.write.call_args.args[1] + with tarfile.open(fileobj=io.BytesIO(archive)) as tar: + tar.extractall(target, filter='data') + restored = {p.relative_to(target).as_posix(): p.read_bytes() for p in target.rglob('*') if p.is_file()} + assert restored == files + asyncio.run(run()) + + +async def test_command_timeout_and_overrides() -> None: + handle = FakeHandle(stdout='partial', stderr='partial error') + handle.hang = True + handle.kill.side_effect = RuntimeError('test-secret kill failed') + box = running_sandbox(handle) + box.config.env_vars = {'A': 'old', 'B': 'kept'} + result = await box.execute_command('sleep 100', timeout=0.001, user='root', cwd='/testbed', env={'A': 'new'}) + assert (result.status, result.stdout, result.stderr) == (ExecutionStatus.TIMEOUT, 'partial', 'partial error') + handle.kill.assert_awaited_once() + kwargs = box._sbx.commands.run.call_args.kwargs + assert kwargs == dict(background=True, user='root', cwd='/testbed', envs={'A': 'new', 'B': 'kept'}, timeout=0) + box._sbx.commands.run.side_effect = RuntimeError('test-secret dispatch failed') + result = await box.execute_command('command') + assert result.exit_code == -1 and result.status == ExecutionStatus.ERROR + assert 'test-secret' not in result.stderr + assert box._sbx.commands.run.call_count == 2 + + +async def test_archive_errors_and_cleanup() -> None: + box = running_sandbox() + with pytest.raises(ValueError, match='absolute'): + await box.put_archive('relative', b'archive') + box._sbx.files.write.assert_not_awaited() + assert await box.put_archive('/test path', b'archive') + assert box._sbx.files.write.call_args.kwargs == {'user': 'user'} + cmd = box._sbx.commands.run.call_args.args[0] + assert "-C '/test path'" in cmd and 'rm -f' in cmd + box._sbx.files.write.side_effect = RuntimeError('upload failed') + assert not await box.put_archive('/path', b'archive') + box._sbx.files.write.side_effect = None + box._sbx.commands.handle.code = 1 + assert not await box.put_archive('/path', b'archive') + remote = box._sbx + remote.kill.side_effect = RuntimeError('kill failed') + await box.stop() + await box.stop() + remote.kill.assert_awaited_once() + + +async def test_create_metadata_factory_and_failure(sdk: SimpleNamespace) -> None: + box = SandboxFactory.create_sandbox(SandboxType.E2B, {'template': 'explicit'}) + assert isinstance(box.config, E2BSandboxConfig) + box._conn = connection() + await box.start() + kwargs = sdk.create.call_args.kwargs + assert kwargs['metadata'] == {'created_by': 'ms-enclave', 'run_id': box._run_id} + assert kwargs['timeout'] == 3600 and kwargs['retries'] == 0 + assert 'lifecycle' not in kwargs and 'envs' not in kwargs + await box.stop() + sdk.create.side_effect = RuntimeError('test-secret Timeout cannot be greater than 1 hours') + with pytest.raises(RuntimeError, match='template=explicit.*run_id=.*sandbox_timeout=3600') as error: + await box.start() + assert '3600s' in str(error.value) and 'test-secret' not in str(error.value) + + +# Feature: e2b-sandbox, Property 10: Live sandbox limits and slots survive create failures and repeated deletion. +@settings(max_examples=100) +@given(st.integers(1, 8), st.lists(st.booleans(), min_size=1, max_size=16)) +def test_manager_slots(limit: int, outcomes: list[bool]) -> None: + async def run() -> None: + with pytest.MonkeyPatch.context() as patch: + active = peak = 0 + remotes = [] + async def create(**kwargs): + nonlocal active, peak + if not outcomes.pop(0): + raise RuntimeError('creation failed') + remote = FakeAsyncSandbox() + active += 1 + peak = max(peak, active) + async def kill(**kwargs): + nonlocal active + active -= 1 + remote.kill.side_effect = kill + remotes.append(remote) + return remote + patch.setattr(sandbox_module, 'AsyncSandbox', SimpleNamespace(create=create)) + manager = E2BSandboxManager({'api_key': 'placeholder', 'max_sandboxes': limit}) + await manager.start() + async def worker() -> None: + try: + sid = await manager.create_sandbox(SandboxType.E2B, {'template': 'explicit'}) + except RuntimeError: + return + await asyncio.sleep(0) + await asyncio.gather(manager.delete_sandbox(sid), manager.delete_sandbox(sid)) + await asyncio.gather(*(worker() for _ in range(len(outcomes)))) + await manager.stop() + assert peak <= limit and active == 0 + assert manager._slots._value == limit and not manager._sandboxes + assert all(remote.kill.call_count == 1 for remote in remotes) + outcomes = list(outcomes) + asyncio.run(run()) + + +# Feature: e2b-sandbox, Property 11: SDK failure diagnostics and cached errors never expose the API key. +@settings(max_examples=100) +@given(st.text(alphabet='abcdef0123456789', min_size=12, max_size=24)) +def test_secret_redaction(secret: str) -> None: + async def run() -> None: + conn = E2BConnection.from_config(E2BSandboxManagerConfig(api_key=secret)) + assert secret not in redact(f'{secret} twice {secret}', secret) + resolver = E2BTemplateResolver(conn, {}, 1) + resolver._find = AsyncMock(side_effect=RuntimeError(secret)) + with pytest.raises(E2BTemplateError) as error: + await resolver.resolve(E2BSandboxConfig(image='org/image'), run_id='r') + assert secret not in str(error.value) and secret not in str(resolver._failed) + box = running_sandbox() + box._conn = conn + box._sbx.commands.run.side_effect = RuntimeError(secret) + assert secret not in (await box.execute_command('cmd')).stderr + call = AsyncMock(side_effect=RuntimeError(secret)) + with pytest.raises(RuntimeError) as error: + await retry_dispatch('create', call, run_id='r', redactor=conn.redact) + assert secret not in str(error.value) + asyncio.run(run()) + + +async def test_manager_cleanup_and_invalid_type(sdk: SimpleNamespace) -> None: + manager = E2BSandboxManager({'api_key': 'placeholder'}) + assert re.fullmatch('[0-9a-f]{12}', manager.run_id) + await manager.start() + with pytest.raises(ValueError, match='only supports'): + await manager.create_sandbox(SandboxType.DOCKER, {}) + sid = await manager.create_sandbox(SandboxType.E2B, {'template': 'explicit'}) + assert sid in manager._sandboxes + await manager.stop() + assert not manager._sandboxes + sdk.remote.kill.assert_awaited_once() + + +@pytest.mark.parametrize('explicit_size', [None, 2]) +async def test_pool_rejects_impossible_size(sdk: SimpleNamespace, explicit_size: int | None) -> None: + manager = E2BSandboxManager({'api_key': 'placeholder', 'max_sandboxes': 1, 'pool_size': 2}) + await manager.start() + try: + with pytest.raises(ValueError, match='pool_size=2 exceeds E2B max_sandboxes=1'): + await asyncio.wait_for(manager.initialize_pool(explicit_size, config={'template': 'explicit'}), 1) + sdk.create.assert_not_awaited() + finally: + await manager.stop() + + +async def test_pool_defaults_to_e2b(sdk: SimpleNamespace) -> None: + manager = E2BSandboxManager({'api_key': 'placeholder', 'max_sandboxes': 1}) + await manager.start() + try: + pool = await manager.initialize_pool(1, config={'template': 'explicit'}) + assert len(pool) == 1 + assert (await manager.get_sandbox(pool[0])).sandbox_type == SandboxType.E2B + finally: + await manager.stop() + sdk.remote.kill.assert_awaited_once() + + +@pytest.mark.parametrize('tool', ['shell_executor', 'python_executor']) +@pytest.mark.parametrize('status', [ExecutionStatus.SUCCESS, ExecutionStatus.ERROR, ExecutionStatus.TIMEOUT]) +async def test_tools(tool: str, status: ExecutionStatus) -> None: + box = running_sandbox() + box.config.tools_config = {tool: {}} + await box.initialize_tools() + from ms_enclave.sandbox.model import CommandResult + box.execute_command = AsyncMock(return_value=CommandResult( + command='cmd', status=status, exit_code=0 if status == ExecutionStatus.SUCCESS else -1, stdout='out', stderr='err')) + params = {'command': 'echo hi'} if tool == 'shell_executor' else {'code': 'print(1)'} + result = await box.execute_tool(tool, params) + assert result.status == status + + +def test_optional_dependency() -> None: + script = """ +import sys +sys.modules['e2b'] = None +import ms_enclave.sandbox +from ms_enclave.sandbox.manager import E2BSandboxManager +try: + E2BSandboxManager() +except ImportError as exc: + assert "pip install 'ms-enclave[e2b]'" in str(exc) +else: + raise AssertionError('missing installation error') +""" + subprocess.run([sys.executable, '-c', script], check=True, capture_output=True) diff --git a/tests/test_tool.py b/tests/test_tool.py index 8016921..f1d6855 100644 --- a/tests/test_tool.py +++ b/tests/test_tool.py @@ -25,7 +25,7 @@ def test_create_python_executor_tool(self): self.assertEqual(tool.name, 'python_executor') self.assertIsNotNone(tool.description) self.assertIsNotNone(tool.schema) - self.assertEqual(tool.required_sandbox_types, [SandboxType.DOCKER, SandboxType.VOLCENGINE]) + self.assertEqual(tool.required_sandbox_types, [SandboxType.DOCKER, SandboxType.VOLCENGINE, SandboxType.E2B]) def test_create_unknown_tool_raises_error(self): """Test that creating unknown tool raises appropriate error."""