agentenv-framework 0.9.1254__py3-none-any.whl

agent_env/providers/sandbox_providers/e2b/sandbox.py

@@ -0,0 +1,342 @@
+"""E2B's async sandbox presented through the :class:`VmSandbox` contract.
+
+The E2B SDK intentionally returns a completed ``CommandResult`` from
+``commands.run``.  ``VmSandbox``, on the other hand, is built on a small
+process interface whose stdout and stderr are async byte streams.  This module
+bridges those two shapes without initializing E2B SDK clients at module import
+time, keeping provider registration and offline unit tests isolated.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import shlex
+import time
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, Any
+from urllib.parse import urlparse
+
+from agent_env.providers.sandbox_providers.sandbox import NetworkMode, NetworkPolicy, VmSandbox
+from agent_env.providers.sandbox_providers.sandbox_provider import SANDBOX_MODE_VM
+
+if TYPE_CHECKING:
+    from e2b import AsyncSandbox
+
+logger = logging.getLogger(__name__)
+
+E2B_ALL_TRAFFIC = "0.0.0.0/0"
+
+
+class _BytesReader:
+    """The small ``StreamReader`` subset consumed by ``Sandbox.exec_with_output``."""
+
+    def __init__(self, value: str | bytes | None):
+        self._value = value.encode() if isinstance(value, str) else (value or b"")
+
+    async def read(self) -> bytes:
+        return self._value
+
+
+class _E2BCompletedProcess:
+    """Adapt E2B's completed command result to the common process protocol."""
+
+    def __init__(self, result: Any):
+        self.stdout = _BytesReader(getattr(result, "stdout", ""))
+        self.stderr = _BytesReader(getattr(result, "stderr", ""))
+        self._exit_code = int(getattr(result, "exit_code", -1))
+
+    async def wait(self) -> int:
+        return self._exit_code
+
+
+class _E2BTunnelURLs(dict[int, str]):
+    """A port map that asks E2B for an external host on first use.
+
+    E2B doesn't allocate persistent tunnel objects.  Its ``get_host(port)``
+    derives the public endpoint for any port, so laziness both avoids a stale
+    static port list after reconnect and preserves the normal ``tunnel_urls``
+    indexing used throughout agent-env.
+    """
+
+    def __init__(self, sandbox: E2BSandbox):
+        super().__init__()
+        self._sandbox = sandbox
+
+    def __missing__(self, port: int) -> str:
+        return self._sandbox.get_host(port)
+
+    def get(self, port: int, default: str | None = None) -> str | None:
+        # ``dict.get`` is used as a capability query by deployment callers.
+        # Unlike indexing, it must not imply that every arbitrary port has a
+        # reachable E2B endpoint (nor should it mutate this cache).
+        return super().get(port, default)
+
+
+class E2BSandbox(VmSandbox):
+    """An E2B Linux sandbox with Docker, exposed as an agent-env VM sandbox.
+
+    ``sandbox`` is an already-created ``e2b.AsyncSandbox``.  Keeping creation
+    in the provider lets this adapter stay focused on VM process, lifecycle,
+    reconnect, and public-port behavior.
+    """
+
+    type = "e2b"
+    # ``docker info`` should normally return in well under one second.  Keep
+    # each readiness probe short so a wedged daemon cannot consume the entire
+    # VM-ready budget in one await.
+    _DOCKER_READINESS_COMMAND_TIMEOUT = 10
+    _MIN_COMMAND_TIMEOUT_SECONDS = 0.001
+    _UNBOUNDED_COMMAND_TIMEOUT_SECONDS = 0
+    _DOCKER_LOG_TAIL_LINES = 40
+    _UNKNOWN_EXIT_CODE = -1
+
+    def __init__(
+        self,
+        sandbox: AsyncSandbox | Any,
+        *,
+        exposed_ports: Iterable[int] = (),
+        network_policy: NetworkPolicy | None = None,
+    ):
+        self._sandbox = sandbox
+        self.sandbox_id = sandbox.sandbox_id
+        self.vnc_url = None
+        self.mode = SANDBOX_MODE_VM
+        self.network_policy = network_policy
+        self.tunnel_urls = _E2BTunnelURLs(self)
+        for port in exposed_ports:
+            self.get_host(port)
+
+    async def terminate(self) -> None:
+        """Irreversibly stop the E2B sandbox (``kill`` is idempotent)."""
+        killed = await self._sandbox.kill()
+        if not killed:
+            logger.info("E2B sandbox %s was already absent during terminate", self.sandbox_id)
+
+    async def exec(self, *command: str) -> _E2BCompletedProcess:
+        """Run argv safely through E2B's shell-command API.
+
+        E2B accepts one command string and itself invokes Bash.  ``shlex.join``
+        keeps arguments distinct, notably the scripts supplied to inherited
+        ``exec_script``.  Do not strip ``sudo``: E2B commands run as the normal
+        sandbox user and Docker administration requires the inherited sudo
+        prefixes (``sudo bash`` / ``sudo docker``).
+        """
+        return await self._run_command(
+            *command,
+            timeout=self._UNBOUNDED_COMMAND_TIMEOUT_SECONDS,
+        )
+
+    async def _run_command(
+        self,
+        *command: str,
+        timeout: float,
+    ) -> _E2BCompletedProcess:
+        """Run an E2B command with an explicit SDK command timeout.
+
+        Normal VM operations remain unbounded (``timeout=0``) because image
+        loads and agent commands can legitimately run for a long time.
+        Readiness uses this primitive with a finite timeout so its wall-clock
+        budget remains meaningful even when Docker is wedged.
+        """
+        try:
+            result = await self._sandbox.commands.run(shlex.join(command), timeout=timeout)
+        except Exception as exc:
+            # E2B raises CommandExitException for a non-zero foreground command;
+            # it is also a CommandResult, so expose that failure through the
+            # common process result rather than treating it as an SDK transport
+            # error.  Other exceptions (network/auth/etc.) must still surface.
+            if not all(hasattr(exc, attribute) for attribute in ("exit_code", "stdout", "stderr")):
+                raise
+            result = exc
+        return _E2BCompletedProcess(result)
+
+    async def _exec_with_output_bounded(
+        self,
+        *command: str,
+        timeout: float,
+    ) -> tuple[int, str, str]:
+        """Execute one short readiness/diagnostic command with an E2B timeout."""
+        process = await self._run_command(*command, timeout=timeout)
+        stdout, stderr = await asyncio.gather(process.stdout.read(), process.stderr.read())
+        return await process.wait(), stdout.decode(), stderr.decode()
+
+    async def setup_vm_for_gateway(self, exposed_ports: list[int] | None = None) -> None:
+        """Start Docker and verify the VM has the Compose v2 gateway runtime.
+
+        E2B publishes every sandbox port through its dynamic ``get_host``
+        endpoint, so unlike the KubeVirt implementation there is no guest
+        firewall to configure. ``EnvironmentGatewayProvider`` starts all co-located
+        services with ``docker compose up``, so accepting a Docker-only base
+        template here would defer a deterministic configuration error until
+        deploy-env is already staging images and files.
+        """
+        await self.wait_for_vm()
+        exit_code, version, stderr = await self._exec_with_output_bounded(
+            "sudo",
+            "docker",
+            "compose",
+            "version",
+            timeout=self._DOCKER_READINESS_COMMAND_TIMEOUT,
+        )
+        if exit_code != 0:
+            detail = (stderr or version).strip()
+            raise RuntimeError(
+                "E2B base template must include the Docker Compose v2 plugin "
+                "required by deploy-env"
+                + (f": {detail}" if detail else "")
+            )
+        logger.info(
+            "Docker Compose ready in E2B sandbox %s: %s",
+            self.sandbox_id,
+            version.strip(),
+        )
+
+    async def wait_for_vm(self) -> None:
+        """Ensure E2B's Docker daemon is available to inherited VM helpers."""
+        await self._ensure_docker_running()
+
+    async def _ensure_docker_running(self) -> None:
+        # The base template normally has systemd, but retain service and direct
+        # dockerd fallbacks for custom templates.  `nohup ... &` detaches the
+        # last fallback from E2B's completed-command stream.
+        deadline = time.monotonic() + self._VM_READY_TIMEOUT
+        start_command = (
+            "sudo systemctl start docker || sudo service docker start || "
+            "(sudo nohup dockerd > /var/log/dockerd.log 2>&1 &)"
+        )
+        try:
+            start_timeout = max(
+                self._MIN_COMMAND_TIMEOUT_SECONDS,
+                min(self._DOCKER_READINESS_COMMAND_TIMEOUT, self._VM_READY_TIMEOUT),
+            )
+            start_result = await self._run_command("bash", "-c", start_command, timeout=start_timeout)
+            if await start_result.wait() != 0:
+                logger.debug("E2B Docker startup command exited nonzero in %s", self.sandbox_id)
+        except Exception as exc:
+            logger.debug("E2B Docker startup command did not complete: %s", exc)
+
+        attempts = 0
+        while time.monotonic() < deadline:
+            attempts += 1
+            remaining = deadline - time.monotonic()
+            command_timeout = min(self._DOCKER_READINESS_COMMAND_TIMEOUT, remaining)
+            try:
+                exit_code, _, _ = await self._exec_with_output_bounded(
+                    "sudo",
+                    "docker",
+                    "info",
+                    timeout=command_timeout,
+                )
+            except Exception as exc:
+                # A timed-out or temporarily unavailable command is simply a
+                # failed readiness probe.  The enclosing deadline, rather than
+                # a single SDK call, owns the readiness budget.
+                logger.debug(
+                    "Docker readiness probe failed in E2B sandbox %s: %s",
+                    self.sandbox_id,
+                    exc,
+                )
+                exit_code = self._UNKNOWN_EXIT_CODE
+            if exit_code == 0:
+                logger.info("Docker ready in E2B sandbox %s after %s polls", self.sandbox_id, attempts)
+                return
+            sleep_seconds = min(self._VM_READY_POLL_INTERVAL, max(0, deadline - time.monotonic()))
+            if sleep_seconds:
+                await asyncio.sleep(sleep_seconds)
+
+        try:
+            _, log_tail, _ = await self._exec_with_output_bounded(
+                "sudo",
+                "tail",
+                "-n",
+                str(self._DOCKER_LOG_TAIL_LINES),
+                "/var/log/dockerd.log",
+                timeout=self._DOCKER_READINESS_COMMAND_TIMEOUT,
+            )
+        except Exception as exc:
+            log_tail = f"unavailable ({type(exc).__name__}: {exc})"
+        raise RuntimeError(
+            f"Docker not ready in E2B sandbox {self.sandbox_id} after "
+            f"{self._VM_READY_TIMEOUT}s ({attempts} polls); "
+            f"/var/log/dockerd.log tail:\n{log_tail}"
+        )
+
+    def get_host(self, port: int) -> str:
+        """Return and cache E2B's current HTTPS URL for ``port``.
+
+        ``AsyncSandbox.get_host`` returns a hostname, not a URL.  Normalizing it
+        here keeps ``tunnel_urls`` compatible with the other providers.
+        """
+        normalized_port = int(port)
+        host = self._sandbox.get_host(normalized_port)
+        url = host if host.startswith(("http://", "https://")) else f"https://{host}"
+        self.tunnel_urls[normalized_port] = url
+        return url
+
+    async def apply_network_policy(self, policy: NetworkPolicy) -> None:
+        """Atomically replace E2B egress policy with ``policy``.
+
+        E2B's ``update_network`` clears omitted fields.  ALLOW_ALL explicitly
+        enables internet access.  ALLOWLIST combines E2B's canonical
+        deny-all sentinel with explicit hostname/CIDR exceptions.
+        """
+        # ``allow_public_traffic`` is creation-only and is silently discarded
+        # by update_network.  Do not confuse ingress configuration with the
+        # workload's outbound policy here.
+        network: dict[str, Any] = {"allow_internet_access": True}
+        if policy.mode is NetworkMode.ALLOWLIST:
+            network["allow_out"] = [*policy.allow_hosts, *policy.allow_cidrs]
+            network["deny_out"] = [E2B_ALL_TRAFFIC]
+        await self._sandbox.update_network(network)
+        self.network_policy = policy
+
+    async def load_docker_images(self, artifacts: list) -> None:
+        """Load images after adding their signed-download hosts to the policy.
+
+        With a restrictive policy, the object store that supplies image tarballs
+        is infrastructure rather than a workload destination.  E2B replaces an
+        egress policy atomically, so union the exact signed-URL hostnames into the
+        sandbox's applied policy. CIDRs stay untouched through
+        :meth:`NetworkPolicy.with_hosts`.
+        """
+        policy = self.network_policy
+        if not artifacts:
+            return
+        if policy is None:
+            raise RuntimeError(
+                f"Cannot load Docker images in reconnected E2B sandbox {self.sandbox_id}: "
+                "its applied network policy is unknown, so signed download hosts cannot "
+                "be added safely"
+            )
+        if not policy.restricts_egress:
+            await super().load_docker_images(artifacts)
+            return
+
+        signed_urls = await self._signed_image_urls(artifacts)
+        signed_hosts = {
+            parsed.hostname
+            for signed_url in signed_urls
+            if signed_url and (parsed := urlparse(signed_url)).hostname
+        }
+        if signed_hosts:
+            await self.apply_network_policy(policy.with_hosts(sorted(signed_hosts)))
+        await self._load_docker_images(artifacts, signed_urls)
+
+    @classmethod
+    async def reconnect(
+        cls,
+        sandbox_id: str,
+        *,
+        api_key: str | None = None,
+        sandbox_cls: type[AsyncSandbox] | Any | None = None,
+        exposed_ports: Iterable[int] = (),
+        network_policy: NetworkPolicy | None = None,
+    ) -> E2BSandbox:
+        """Reconnect to an existing sandbox without relying on ambient credentials."""
+        if sandbox_cls is None:
+            # Optional dependency: importing this adapter must not require E2B.
+            from e2b import AsyncSandbox as sandbox_cls
+
+        sandbox = await sandbox_cls.connect(sandbox_id, api_key=api_key)
+        return cls(sandbox, exposed_ports=exposed_ports, network_policy=network_policy)

agent_env/providers/sandbox_providers/e2b/template.py

@@ -0,0 +1,158 @@
+"""Deterministic E2B template aliases for requested sandbox sizing.
+
+E2B fixes a sandbox's CPU and memory at template-build time.  This resolver
+derives a template from a configured base template and gives that derivative a
+stable alias, so equivalent requests reuse the same template.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import math
+import threading
+from numbers import Real
+from typing import Any, ClassVar
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_BUILD_TIMEOUT_SECONDS = 15 * 60
+_LOCK_POLL_INTERVAL_SECONDS = 0.05
+
+
+def _normalize_positive_integer(value: Real, *, field: str) -> int:
+    """Return an integral positive resource value, rejecting ambiguous sizes."""
+    if (
+        isinstance(value, bool)
+        or not isinstance(value, Real)
+        or not math.isfinite(value)
+        or value <= 0
+        or int(value) != value
+    ):
+        raise ValueError(f"{field} must be a positive integer; got {value!r}")
+    return int(value)
+
+
+def normalize_e2b_cpu(cpu: float) -> int:
+    """Validate an exact CPU request accepted by E2B's integer API."""
+    return _normalize_positive_integer(cpu, field="cpu")
+
+
+def normalize_e2b_memory_mb(memory_mb: int) -> int:
+    """Validate an exact memory request accepted by E2B's integer API."""
+    return _normalize_positive_integer(memory_mb, field="memory_mb")
+
+
+class E2BTemplateResolver:
+    """Resolve a size-specific E2B template, building it at most once locally.
+
+    Locks are shared by all resolver instances because providers may be
+    constructed more than once in the same process.  E2B itself remains the
+    authority across processes: a failed build is accepted if a follow-up
+    existence check shows another worker created the same alias.
+    """
+
+    _locks: ClassVar[dict[str, threading.Lock]] = {}
+    _locks_guard: ClassVar[threading.Lock] = threading.Lock()
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        async_template_cls: type[Any] | None = None,
+        template_cls: type[Any] | None = None,
+        build_timeout_seconds: float = _DEFAULT_BUILD_TIMEOUT_SECONDS,
+    ) -> None:
+        self._api_key = api_key
+        self._async_template_cls = async_template_cls
+        self._template_cls = template_cls
+        if build_timeout_seconds <= 0:
+            raise ValueError("build_timeout_seconds must be positive")
+        self._build_timeout_seconds = build_timeout_seconds
+
+    @classmethod
+    def _lock_for(cls, name: str) -> threading.Lock:
+        # A threading lock is safe when separate worker threads each run their
+        # own event loop. ``resolve`` acquires it with non-blocking polls so it
+        # neither blocks the event loop nor leaks an acquired lock if cancelled.
+        with cls._locks_guard:
+            return cls._locks.setdefault(name, threading.Lock())
+
+    @staticmethod
+    async def _acquire_lock(lock: threading.Lock) -> None:
+        while not lock.acquire(blocking=False):
+            await asyncio.sleep(_LOCK_POLL_INTERVAL_SECONDS)
+
+    @staticmethod
+    def _log_build(name: str, entry: Any) -> None:
+        level = getattr(entry, "level", "info")
+        message = getattr(entry, "message", str(entry))
+        log = logger.warning if level in {"warn", "error"} else logger.info
+        log("E2B template build %s: %s", name, message)
+
+    def _sdk(self) -> tuple[type[Any], type[Any]]:
+        if self._async_template_cls is None or self._template_cls is None:
+            # Defer SDK initialization until an E2B backend is actually used.
+            from e2b import AsyncTemplate, Template
+
+            return AsyncTemplate, Template
+        return self._async_template_cls, self._template_cls
+
+    async def resolve(self, base_template: str, *, cpu: float, memory_mb: int) -> str:
+        """Return the E2B alias for ``base_template`` at the requested size."""
+        exact_cpu = normalize_e2b_cpu(cpu)
+        exact_memory_mb = normalize_e2b_memory_mb(memory_mb)
+        name = f"{base_template}-{exact_cpu}c-{exact_memory_mb}m"
+        api_params = {"api_key": self._api_key} if self._api_key is not None else {}
+
+        async_template_cls, template_cls = self._sdk()
+        lock = self._lock_for(name)
+        await self._acquire_lock(lock)
+        try:
+            if await async_template_cls.exists(name, **api_params):
+                return name
+
+            template = template_cls().from_template(base_template)
+            logger.info(
+                "Building E2B template %s from %s (cpu=%s, memory=%sMB, timeout=%ss)",
+                name,
+                base_template,
+                exact_cpu,
+                exact_memory_mb,
+                self._build_timeout_seconds,
+            )
+            try:
+                async with asyncio.timeout(self._build_timeout_seconds):
+                    await async_template_cls.build(
+                        template,
+                        name,
+                        cpu_count=exact_cpu,
+                        memory_mb=exact_memory_mb,
+                        on_build_logs=lambda entry: self._log_build(name, entry),
+                        **api_params,
+                    )
+            except Exception:
+                # A separate process may have created the deterministic alias
+                # after our initial exists check.  In that case its template is
+                # exactly the resource variant we needed, so reuse it.
+                if await async_template_cls.exists(name, **api_params):
+                    return name
+                raise
+            return name
+        finally:
+            lock.release()
+
+
+async def resolve_e2b_template(
+    base_template: str,
+    *,
+    api_key: str | None = None,
+    cpu: float,
+    memory_mb: int,
+) -> str:
+    """Resolve a size-specific E2B template using the default SDK classes."""
+    return await E2BTemplateResolver(api_key=api_key).resolve(
+        base_template,
+        cpu=cpu,
+        memory_mb=memory_mb,
+    )