diff --git a/README.md b/README.md index 67aa102e..df754043 100644 --- a/README.md +++ b/README.md @@ -44,6 +44,24 @@ uv sync --extra dev source .venv/bin/activate ``` +## Run agents in local microVMs + +Install `agentenv-framework[smol]` and select the `smol_vm` sandbox provider in +`.agentenv/config.toml`: + +```toml +[sandbox] +default = "smol_vm" +agent_default = "smol_vm" +``` + +Each agent and environment gets its own local SmolVM machine with a Docker daemon, +so workloads are isolated from the host Docker daemon. This needs a supported +Apple Silicon macOS or Linux host; on Linux, `/dev/kvm` must be available. AgentEnv can still +use its default local object and image stores: guest containers can fetch local +object grants, and the VM forwards pulls from the host's loopback registry. +Cloud image and object stores continue to work through the same provider. + ## Plugins A package of your own can add env types, providers, task steps and stores to agent-env, and `agent-env plugin` lists, checks, adds and removes installed plugins. The [plugin docs](https://www.agentenvframework.com/docs/plugins) show how. diff --git a/pyproject.toml b/pyproject.toml index 4cc53812..958bc9bb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,7 +73,7 @@ sail = [ "sail~=0.12.8", # Sailbox exec/fs/listener/egress APIs validated against the 0.12 line ] dev = [ - "agentenv-framework[aws,explorer,gcp,sail]", # so the AWS, explorer, GCP backend and Sail provider tests run in CI + "agentenv-framework[aws,explorer,gcp,sail,smol]", # exercise the optional providers and stores in CI "fastmcp>=3.4,<4", # the standalone-fastmcp mount path of the protocol; <4 keeps mcp 1.x "griffelib==2.3.0", # the plugin API check; exact, since what it reports as a break changes between releases "moto>=5.0.0", @@ -85,6 +85,9 @@ dev = [ "pytest-timeout>=2.3.0", "pytest-xdist>=3.0.0", ] +smol = [ + "smolmachines>=1.27.0", +] [project.scripts] agent-env = "agent_env.cli:cli" diff --git a/src/agent_env/providers/sandbox_providers/chained_sandbox_provider.py b/src/agent_env/providers/sandbox_providers/chained_sandbox_provider.py index 5fadf549..53a7a0f3 100644 --- a/src/agent_env/providers/sandbox_providers/chained_sandbox_provider.py +++ b/src/agent_env/providers/sandbox_providers/chained_sandbox_provider.py @@ -66,6 +66,17 @@ async def create_sandbox( attribution: Attribution | None = None, **kwargs, ) -> Sandbox: + return await self._create("create_sandbox", attribution=attribution, **kwargs) + + async def create_container( + self, + *, + attribution: Attribution | None = None, + **kwargs, + ) -> Sandbox: + return await self._create("create_container", attribution=attribution, **kwargs) + + async def _create(self, method: str, *, attribution: Attribution | None = None, **kwargs) -> Sandbox: sandbox_providers = self.filter_sandbox_providers(self._providers, kwargs.get("network_policy")) image = self._images.get(kwargs.get("image_name")) errors: list[tuple[str, Exception]] = [] @@ -76,9 +87,9 @@ async def create_sandbox( # outside the deadline, which bounds creating the sandbox, not a build await p.prepare_image(image, attribution=attribution) async with asyncio.timeout(_PROVISION_DEADLINE_SECONDS): - sandbox = await p.create_sandbox(attribution=attribution, **kwargs) + sandbox = await getattr(p, method)(attribution=attribution, **kwargs) except Exception as e: - logger.warning(f"{name}.create_sandbox failed: {e}; trying next provider") + logger.warning(f"{name}.{method} failed: {e}; trying next provider") errors.append((name, e)) continue # A VM refuses an image it can't get or run, and a container runs only what its provider's containers do; diff --git a/src/agent_env/providers/sandbox_providers/local_sandbox.py b/src/agent_env/providers/sandbox_providers/local_sandbox.py index 3ece7364..f0fb5f56 100644 --- a/src/agent_env/providers/sandbox_providers/local_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/local_sandbox.py @@ -45,6 +45,7 @@ Accepts, SandboxProvider, refuse_unenforceable_policy, + _sandbox_provider_class, ) from agent_env.store.object_store.local.tls import local_ca from agent_env.store.object_store.local.store import LocalFilesystemObjectStore @@ -259,6 +260,9 @@ def _rewrite_app_arg(self, arg: str) -> str: def _rewrite_app_script(self, script: str) -> str: return _APP_PATH_PATTERN.sub(lambda _: str(self._work_dir), script) + async def copy_trust_dir_into_container(self, trust_dir: Path, container: str, destination: str) -> None: + await asyncio.to_thread(_copy_into_container, trust_dir, container, destination) + async def terminate(self) -> None: """Tear down whatever this sandbox is running. @@ -507,7 +511,7 @@ async def start_trusting(sandbox: VmSandbox, container: str, trust_dir: Path) -> """Copy ``trust_dir`` into the created ``container`` where ``LOCAL_TRUST_ENV`` points, then start it. A container that cannot be given the files or started is removed, so its name is free for the next attempt.""" try: - await asyncio.to_thread(_copy_into_container, trust_dir, container, _TRUST_DIR) + await sandbox.copy_trust_dir_into_container(trust_dir, container, _TRUST_DIR) await sandbox.exec_script(f"docker start {shlex.quote(container)} > /dev/null") except Exception: await sandbox.exec_script(f"docker rm -f {shlex.quote(container)} >/dev/null 2>&1 || true") @@ -684,6 +688,8 @@ def get_external_url(cls, url: str) -> str: def host_url_for(url: str, sandbox_type: Optional[str]) -> str: """``url``, which this machine reaches, as a container on ``sandbox_type`` reaches it: only a local sandbox's containers run on this machine, and they reach its loopback by another name.""" + if sandbox_type == "smol_vm": + return _sandbox_provider_class(sandbox_type).get_external_url(url) return LocalSandboxProvider.get_external_url(url) if sandbox_type == LocalSandbox.type else url diff --git a/src/agent_env/providers/sandbox_providers/sandbox.py b/src/agent_env/providers/sandbox_providers/sandbox.py index 2cb83a35..ddf830e0 100644 --- a/src/agent_env/providers/sandbox_providers/sandbox.py +++ b/src/agent_env/providers/sandbox_providers/sandbox.py @@ -19,6 +19,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass, replace from enum import Enum +from pathlib import Path from typing import IO, TYPE_CHECKING, Any, AsyncIterator, Callable, ClassVar, Iterable, Optional from agent_env.config import get_config @@ -473,6 +474,17 @@ async def docker_cp(self, source: str, destination: str, *, remove_source: bool f"docker cp {source} {destination} failed (exit {exit_code}):\nstdout: {stdout[-1500:]}\nstderr: {stderr[-1500:]}" ) + async def copy_trust_dir_into_container(self, trust_dir: Path, container: str, destination: str) -> None: + """Install the local transfer CA in a container on this VM's Docker host.""" + vm_dir = self._staging_path("trust", container) + await self.exec_script(f"mkdir -p {shlex.quote(vm_dir)}") + try: + for path in trust_dir.iterdir(): + await self.write_host_file(await asyncio.to_thread(path.read_bytes), f"{vm_dir}/{path.name}") + await self.docker_cp(f"{vm_dir}/.", f"{container}:{destination}") + finally: + await self.exec_script(f"rm -rf {shlex.quote(vm_dir)}") + async def _copy_into_container(self, vm_path: str, destination_path: str) -> None: parent = os.path.dirname(destination_path) if parent: diff --git a/src/agent_env/providers/sandbox_providers/sandbox_provider.py b/src/agent_env/providers/sandbox_providers/sandbox_provider.py index 39854700..97986410 100644 --- a/src/agent_env/providers/sandbox_providers/sandbox_provider.py +++ b/src/agent_env/providers/sandbox_providers/sandbox_provider.py @@ -108,12 +108,19 @@ async def create_container( Caller still owns readiness polling for the container's own endpoints. ``private_network`` puts it on a network this provider's other such containers share, at its ``private_host``. - Default implementation: create_vm, authenticate Docker to ECR if applicable, docker pull, - docker run. Backends that bake the image in at sandbox-creation time (Modal) override. + Default implementation: a container-only backend uses create_sandbox; a VM backend creates + a VM, pulls the image, and starts it. Backends that bake the image in at creation override. """ if private_network: raise ValueError(f"{type(self).__name__} runs each container in a VM of its own, so it can't put " "containers on a private network they share") + if (type(self).create_vm is SandboxProvider.create_vm + and type(self).create_sandbox is not SandboxProvider.create_sandbox): + return await self.create_sandbox( + image_name=image_name, port=port, env=env, cpu=cpu, memory=memory, + disk_size_gb=disk_size_gb, timeout=timeout, attribution=attribution, + network_policy=network_policy, + ) sandbox = await self.create_vm( cpu=cpu, memory=memory, disk_size_gb=disk_size_gb, exposed_ports=[port], timeout=timeout, @@ -287,6 +294,7 @@ def effective_network_policy(cls, policy: Optional[NetworkPolicy]) -> NetworkPol "modal_vm": "agent_env.providers.sandbox_providers.modal_vm_sandbox:ModalVmSandboxProvider", "e2b": "agent_env.providers.sandbox_providers.e2b:E2BSandboxProvider", "sail_vm": "agent_env.providers.sandbox_providers.sail_vm.provider:SailVmSandboxProvider", + "smol_vm": "agent_env.providers.sandbox_providers.smol_vm.provider:SmolVmSandboxProvider", "local": "agent_env.providers.sandbox_providers.local_sandbox:LocalSandboxProvider", } diff --git a/src/agent_env/providers/sandbox_providers/smol_vm/__init__.py b/src/agent_env/providers/sandbox_providers/smol_vm/__init__.py new file mode 100644 index 00000000..d4b6e275 --- /dev/null +++ b/src/agent_env/providers/sandbox_providers/smol_vm/__init__.py @@ -0,0 +1 @@ +"""Local SmolVM sandbox provider (imported only when the optional extra is selected).""" diff --git a/src/agent_env/providers/sandbox_providers/smol_vm/provider.py b/src/agent_env/providers/sandbox_providers/smol_vm/provider.py new file mode 100644 index 00000000..4c3060ff --- /dev/null +++ b/src/agent_env/providers/sandbox_providers/smol_vm/provider.py @@ -0,0 +1,176 @@ +"""Run AgentEnv workloads in local SmolVM machines.""" + +from __future__ import annotations + +import asyncio +import ipaddress +import math +import shlex +import socket +import uuid +from urllib.parse import urlsplit +from typing import ClassVar + +from smol import ConnectOptions, Machine, MachineConfig, PortSpec, ResourceSpec + +from agent_env.attribution import Attribution +from agent_env.config import get_config +from agent_env.store.image_store.local_registry_image_store import LocalRegistryImageStore +from agent_env.store.routing import LocalRunImageStore +from agent_env.providers.sandbox_providers.sandbox import NetworkPolicy +from agent_env.providers.sandbox_providers.local_sandbox import LOCAL_TRUST_ENV, local_grant_trust, start_trusting +from agent_env.providers.sandbox_providers.sandbox_provider import ( + Accepts, SandboxProvider, refuse_unenforceable_policy, +) +from agent_env.providers.sandbox_providers.smol_vm.sandbox import SmolVmSandbox, _host_url + +_DUMMY_PORT = 60832 + + +def _host_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +class SmolVmSandboxProvider(SandboxProvider): + _DOCKER_READY_ATTEMPTS: ClassVar[int] = 40 + _LOCAL_REGISTRY_DEFAULT_PORT: ClassVar[int] = 5000 + ON_THIS_MACHINE: ClassVar[bool] = True + CREATES_VMS: ClassVar[bool] = True + SANDBOX_ACCEPTS: ClassVar[Accepts] = Accepts.LOADABLE + CONTAINER_ACCEPTS: ClassVar[Accepts] = Accepts.NAME + + @classmethod + def shares_network_with(cls, sandbox_type: str | None) -> bool: + return False + + @classmethod + def get_external_url(cls, url: str) -> str: + return _host_url(url) + + def url_from_sandbox(self, url: str) -> str: + return _host_url(url) + + async def create_sandbox(self, *, image_name: str, port: int, env: dict[str, str], + cpu: float = 1.0, memory: int = 8192, disk_size_gb: float = 10, + timeout: int = 7200, attribution: Attribution | None = None, + network_policy: NetworkPolicy | None = None) -> SmolVmSandbox: + return await self.create_vm( + exposed_ports=[port], cpu=cpu, memory=memory, disk_size_gb=disk_size_gb, + timeout=timeout, attribution=attribution, network_policy=network_policy, + ) + + async def create_vm(self, *, image: str | None = None, boot_mode: str | None = None, + cpu: float = 1.0, memory: int = 8192, disk_size_gb: float = 10, + timeout: int = 7200, exposed_ports: list[int] | None = None, + setup_for_gateway: bool = True, attribution: Attribution | None = None, + network_policy: NetworkPolicy | None = None) -> SmolVmSandbox: + refuse_unenforceable_policy(self, network_policy) + if image is not None or boot_mode is not None: + raise ValueError("smol_vm provisions its own Docker-capable guest; image and boot_mode overrides are unsupported") + if cpu < 1 or memory <= 0 or disk_size_gb <= 0: + raise ValueError("CPU, memory, and storage must be positive") + name = f"agentenv-{uuid.uuid4().hex[:12]}" + ports = {port: _host_port() for port in dict.fromkeys(exposed_ports or [])} + # Even a portless VM needs a published mapping to select routed virtio-net for nested Docker. + published = ports or {_DUMMY_PORT: _host_port()} + config = MachineConfig( + name=name, persistent=True, wait_for_ports=False, + resources=ResourceSpec(cpus=math.ceil(cpu), memory_mb=memory, storage_gb=math.ceil(disk_size_gb), network=True), + ports=[PortSpec(host=host, guest=guest) for guest, host in published.items()], + ) + creation = asyncio.create_task(asyncio.to_thread(Machine.create, config, ConnectOptions(target="local"))) + try: + machine = await asyncio.shield(creation) + except asyncio.CancelledError: + # The native create call keeps running after its awaiting task is cancelled. + # Wait for it and delete the result so cancellation cannot orphan a VM. + try: + created = await creation + except Exception: + pass + else: + await asyncio.to_thread(created.delete) + raise + sandbox = SmolVmSandbox(machine, ports) + try: + await self._setup_docker(sandbox) + await self._configure_host_access(sandbox) + await sandbox.persist_ports() + sandbox.network_policy = self.effective_network_policy(network_policy) + return sandbox + except BaseException: + await sandbox.terminate() + raise + + async def _setup_docker(self, sandbox: SmolVmSandbox) -> None: + code, stdout, stderr = await sandbox.exec_with_output( + "sh", "-c", "command -v bash >/dev/null && command -v docker >/dev/null && " + "command -v curl >/dev/null && command -v python3 >/dev/null && command -v unzip >/dev/null && command -v socat >/dev/null " + "|| apk add --no-cache docker docker-cli-compose bash curl python3 unzip socat", + ) + if code: + raise RuntimeError(f"Could not install guest Docker dependencies ({code}): {stdout[-500:]} {stderr[-500:]}") + await sandbox.exec_script(f""" +set -e +mkdir -p /storage/docker /storage/containerd /var/lib/docker /var/lib/containerd +mountpoint -q /var/lib/docker || mount --bind /storage/docker /var/lib/docker +mountpoint -q /var/lib/containerd || mount --bind /storage/containerd /var/lib/containerd +if ! docker info >/dev/null 2>&1; then + rm -f /var/run/docker.pid /run/containerd/containerd.pid + dockerd --storage-driver=overlay2 >/tmp/agentenv-dockerd.log 2>&1 & + for i in $(seq 1 {self._DOCKER_READY_ATTEMPTS}); do docker info >/dev/null 2>&1 && exit 0; sleep 1; done + tail -80 /tmp/agentenv-dockerd.log + exit 1 +fi +""") + await self._ensure_local_registry_proxy(sandbox) + + async def _ensure_local_registry_proxy(self, sandbox: SmolVmSandbox) -> None: + store = get_config().get_image_store() + stores = (store.configured, store.local) if isinstance(store, LocalRunImageStore) else (store,) + ports = set() + for candidate in stores: + if isinstance(candidate, LocalRegistryImageStore): + registry = urlsplit(f"//{candidate.registry_host}") + if registry.hostname in ("localhost", "127.0.0.1"): + ports.add(registry.port or self._LOCAL_REGISTRY_DEFAULT_PORT) + for port in sorted(ports): + pid_file = f"/run/agentenv-registry-proxy-{port}.pid" + await sandbox.exec_script(f""" +if ! test -f {pid_file} || ! kill -0 "$(cat {pid_file})" 2>/dev/null; then + socat TCP-LISTEN:{port},bind=127.0.0.1,reuseaddr,fork TCP:host.smolvm.internal:{port} \ + >/tmp/agentenv-registry-proxy-{port}.log 2>&1 & + echo $! > {pid_file} +fi +sleep 0.1 +kill -0 "$(cat {pid_file})" +""") + + async def get_sandbox(self, sandbox_id: str) -> SmolVmSandbox: + sandbox = await SmolVmSandbox.reconnect(sandbox_id) + await self._setup_docker(sandbox) + await self._configure_host_access(sandbox) + sandbox.network_policy = self.effective_network_policy(None) + return sandbox + + async def _configure_host_access(self, sandbox: SmolVmSandbox) -> None: + code, stdout, stderr = await sandbox.exec_with_output("getent", "ahostsv4", "host.smolvm.internal") + if code or not stdout.strip(): + raise RuntimeError(f"Cannot resolve SmolVM host gateway: {stderr[-500:]}") + gateway = str(ipaddress.IPv4Address(stdout.split()[0])) + sandbox.extra_hosts = (f"host.docker.internal:{gateway}",) + + def _container_args(self, sandbox: SmolVmSandbox, *, image_name: str, port: int, env: dict[str, str]) -> str: + args = super()._container_args(sandbox, image_name=image_name, port=port, env=env) + return " ".join(f"--add-host {shlex.quote(entry)}" for entry in sandbox.extra_hosts) + " " + args + + async def _start_container(self, sandbox: SmolVmSandbox, *, image_name: str, port: int, env: dict[str, str]) -> None: + trust_dir = await asyncio.to_thread(local_grant_trust) + if trust_dir is None: + await super()._start_container(sandbox, image_name=image_name, port=port, env=env) + return + args = self._container_args(sandbox, image_name=image_name, port=port, env={**LOCAL_TRUST_ENV, **env}) + await sandbox.exec_script(f"docker create {args} > /dev/null") + await start_trusting(sandbox, sandbox.container_name, trust_dir) diff --git a/src/agent_env/providers/sandbox_providers/smol_vm/sandbox.py b/src/agent_env/providers/sandbox_providers/smol_vm/sandbox.py new file mode 100644 index 00000000..557169ff --- /dev/null +++ b/src/agent_env/providers/sandbox_providers/smol_vm/sandbox.py @@ -0,0 +1,99 @@ +"""A local SmolVM machine with a Docker daemon, exposed through the VM sandbox contract.""" + +from __future__ import annotations + +import asyncio +import json +import shlex +from typing import ClassVar +from urllib.parse import urlsplit, urlunsplit + +from smol import ConnectOptions, Machine + +from agent_env.config import get_config +from agent_env.providers.sandbox_providers.sandbox import CURL_RETRY_FLAGS, VmSandbox +from agent_env.store.object_store.local.store import LocalFilesystemObjectStore +from agent_env.store.object_store.local.tls import local_ca +from agent_env.providers.sandbox_providers.sandbox_provider import SANDBOX_MODE_VM + +_PORTS_FILE = "/storage/agentenv-ports.json" +_HOST_ALIAS = "host.smolvm.internal" + + +def _host_url(url: str) -> str: + parts = urlsplit(url) + if parts.hostname not in ("localhost", "127.0.0.1", "::1", "0.0.0.0", "::", "host.docker.internal"): + return url + userinfo, at, _ = parts.netloc.rpartition("@") + port = f":{parts.port}" if parts.port is not None else "" + return urlunsplit(parts._replace(netloc=f"{userinfo}{at}{_HOST_ALIAS}{port}")) + + +class SmolVmSandbox(VmSandbox): + _READ_GRANT_LIFETIME_SECONDS: ClassVar[int] = 3600 + type = "smol_vm" + ON_THIS_MACHINE = True + + def __init__(self, machine: Machine, port_map: dict[int, int]): + self._machine = machine + self._port_map = port_map + self.sandbox_id = machine.id + self.tunnel_urls = {port: machine.endpoint(port).http_url for port in port_map} + self.vnc_url = None + self.mode = SANDBOX_MODE_VM + self.network_policy = None + self.extra_hosts = () + + def url_from_sandbox(self, url: str) -> str: + return _host_url(url) + + async def terminate(self) -> None: + await asyncio.to_thread(self._machine.delete) + + async def exec_with_output(self, *args: str) -> tuple[int, str, str]: + command = list(args[1:] if args and args[0] == "sudo" else args) + result = await asyncio.to_thread(self._machine.exec, command) + return result.exit_code, result.stdout, result.stderr + + async def write_host_file(self, data: bytes, vm_path: str) -> None: + await asyncio.to_thread(self._machine.write_file, vm_path, data) + + async def _download_object_to_vm(self, object_url: str, vm_path: str) -> None: + store = get_config().get_object_store_at(object_url) + signed = await asyncio.to_thread(store.signed_get_url, object_url) + local_grant = isinstance(store, LocalFilesystemObjectStore) and store.supports_transfer_grants and store.grants_reach(self.type) + if signed is None and local_grant: + signed = (await asyncio.to_thread( + store.issue_read_grant, object_url, expires_in=self._READ_GRANT_LIFETIME_SECONDS, + )).url + if signed is None: + await self._write_unsigned_object(store, object_url, vm_path) + return + rewritten = _host_url(signed) + local_https = local_grant or (rewritten != signed and urlsplit(signed).scheme == "https") + if local_https: + await self.write_host_file(local_ca().bundle_path.read_bytes(), "/etc/agentenv-ca-bundle.pem") + ca = "SSL_CERT_FILE=/etc/agentenv-ca-bundle.pem " if local_https else "" + await self.exec_script(f"{ca}curl -fsSL {CURL_RETRY_FLAGS} {shlex.quote(rewritten)} -o {shlex.quote(vm_path)}") + + async def write_file_from_url(self, url: str, destination_path: str) -> None: + vm_path = self._staging_path("url", destination_path) + rewritten = _host_url(url) + local_https = rewritten != url and urlsplit(url).scheme == "https" + try: + if local_https: + await self.write_host_file(local_ca().bundle_path.read_bytes(), "/etc/agentenv-ca-bundle.pem") + ca = "SSL_CERT_FILE=/etc/agentenv-ca-bundle.pem " if local_https else "" + await self.exec_script(f"{ca}curl -fsSL {CURL_RETRY_FLAGS} {shlex.quote(rewritten)} -o {shlex.quote(vm_path)}") + await self._copy_into_container(vm_path, destination_path) + finally: + await self._remove_vm_temp_file(vm_path) + + async def persist_ports(self) -> None: + await self.write_host_file(json.dumps(self._port_map).encode(), _PORTS_FILE) + + @classmethod + async def reconnect(cls, name: str) -> SmolVmSandbox: + machine = await asyncio.to_thread(Machine.connect, name, ConnectOptions(target="local", wait_for_ports=False)) + data = await asyncio.to_thread(machine.read_file, _PORTS_FILE) + return cls(machine, {int(port): int(host) for port, host in json.loads(data).items()}) diff --git a/src/agent_env/store/object_store/local/grant_server.py b/src/agent_env/store/object_store/local/grant_server.py index 6647e7d5..6e9f4d27 100644 --- a/src/agent_env/store/object_store/local/grant_server.py +++ b/src/agent_env/store/object_store/local/grant_server.py @@ -22,7 +22,7 @@ import time from pathlib import Path from collections.abc import AsyncIterator -from typing import TYPE_CHECKING, BinaryIO +from typing import TYPE_CHECKING, BinaryIO, ClassVar import uvicorn from python_multipart.exceptions import FormParserError @@ -134,6 +134,8 @@ class GrantServer: """Serves the grants of the local stores registered with it, over HTTPS on one port. Started by the first grant it issues; it runs in a daemon thread for the rest of the process.""" + _BIND_ATTEMPTS: ClassVar[int] = 5 + def __init__(self, bind_host: str, advertise_host: str) -> None: self.bind_host = bind_host self.advertise_host = advertise_host @@ -192,16 +194,35 @@ def _start(self) -> int: if self._port is not None: return self._port address = ipaddress.ip_address(self.bind_host) - sock = socket.socket(socket.AF_INET6 if address.version == 6 else socket.AF_INET, socket.SOCK_STREAM) + # On Linux the Docker bridge is the primary bind address, but SmolVM's + # host.smolvm.internal forwards to host loopback. Serve both on one port + # without exposing the grant server on public interfaces. + bind_hosts = [self.bind_host] + if self.bind_host == default_bind_host() and self.bind_host != "127.0.0.1": + bind_hosts.append("127.0.0.1") + sockets: list[socket.socket] = [] try: - sock.bind((self.bind_host, 0)) - hosts = ["localhost", "host.docker.internal", "127.0.0.1", self.advertise_host] + for attempt in range(self._BIND_ATTEMPTS): + try: + for host in bind_hosts: + sock = socket.socket(socket.AF_INET6 if address.version == 6 and host == self.bind_host else socket.AF_INET, socket.SOCK_STREAM) + sockets.append(sock) + sock.bind((host, sockets[0].getsockname()[1] if len(sockets) > 1 else 0)) + break + except OSError: + for sock in sockets: + sock.close() + sockets.clear() + if attempt == self._BIND_ATTEMPTS - 1: + raise + hosts = ["localhost", "host.docker.internal", "host.smolvm.internal", "127.0.0.1", self.advertise_host] with contextlib.suppress(ValueError): # a name the CA cannot vouch for would void the certificate check_local_host(self.bind_host) hosts.append(self.bind_host) context = server_context(local_ca(), hosts) except (OSError, ValueError) as e: - sock.close() + for sock in sockets: + sock.close() raise GrantUnavailableError(f"the local grant server cannot listen on {self.bind_host}: {e}") from e config = uvicorn.Config( self._app, lifespan="off", log_config=None, http=_UnloggedH11Protocol, @@ -212,7 +233,7 @@ def _start(self) -> int: def serve() -> None: try: - asyncio.run(server.serve(sockets=[sock])) + asyncio.run(server.serve(sockets=sockets)) except BaseException as e: # uvicorn ends a failed startup with SystemExit failure.append(e) @@ -222,15 +243,16 @@ def serve() -> None: while not server.started: if failure or not thread.is_alive() or time.monotonic() > deadline: server.should_exit = True - sock.close() + for sock in sockets: + sock.close() raise GrantUnavailableError( f"the local grant server did not start on {self.bind_host}" + (f": {failure[0]!r}" if failure else "") ) time.sleep(0.01) self._uvicorn, self._thread = server, thread - self._port = sock.getsockname()[1] - logger.info("Local grant server listening on %s:%d", self.bind_host, self._port) + self._port = sockets[0].getsockname()[1] + logger.info("Local grant server listening on %s:%d", ", ".join(bind_hosts), self._port) return self._port async def _handle(self, request: Request) -> Response: diff --git a/src/agent_env/store/object_store/local/store.py b/src/agent_env/store/object_store/local/store.py index fe4860d6..fdec14c2 100644 --- a/src/agent_env/store/object_store/local/store.py +++ b/src/agent_env/store/object_store/local/store.py @@ -194,7 +194,11 @@ def get_object_key(self, object_url: str) -> str: return path.relative_to(root).as_posix() def grants_reach(self, sandbox_type: str | None) -> bool: - return sandbox_type == _LOCAL_SANDBOX_TYPE + return sandbox_type == _LOCAL_SANDBOX_TYPE or ( + sandbox_type == "smol_vm" + and self._grant_bind_host in (None, "127.0.0.1") + and self._grant_advertise_host in (None, "host.docker.internal", "host.smolvm.internal") + ) def issue_read_grant(self, object_url: str, *, expires_in: int | None = None) -> HttpGetGrant: expires_in = self.grant_lifetime_seconds if expires_in is None else expires_in diff --git a/src/agent_env/store/object_store/local/tls.py b/src/agent_env/store/object_store/local/tls.py index d3ae6d90..016159d5 100644 --- a/src/agent_env/store/object_store/local/tls.py +++ b/src/agent_env/store/object_store/local/tls.py @@ -32,7 +32,7 @@ from agent_env.store.local_state import ensure_state_dir # What the CA may vouch for: the names containers reach the host by, and loopback and private addresses. -LOCAL_NAMES = ("localhost", "host.docker.internal") +LOCAL_NAMES = ("localhost", "host.docker.internal", "host.smolvm.internal") LOCAL_NETWORKS = tuple( ipaddress.ip_network(n) for n in ("127.0.0.0/8", "10.0.0.0/8", "172.16.0.0/12", "192.168.0.0/16", "::1/128", "fc00::/7") diff --git a/src/agent_env/task_step/task_steps/deploy_sandbox.py b/src/agent_env/task_step/task_steps/deploy_sandbox.py index f0d8cd37..581e15b3 100644 --- a/src/agent_env/task_step/task_steps/deploy_sandbox.py +++ b/src/agent_env/task_step/task_steps/deploy_sandbox.py @@ -152,7 +152,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: if self.image is None or self.port is None: raise ValueError("container-mode sandbox requires `image` and `port`") provider = skip_refusing(provider, self._pull_refusal, what=f"sandbox '{self.sandbox_name}'") - sandbox = await provider.create_sandbox( + sandbox = await provider.create_container( image_name=self.image, port=self.port, env=self.env_vars or {}, diff --git a/tst/manual/smol_vm_agent_check.py b/tst/manual/smol_vm_agent_check.py new file mode 100644 index 00000000..9909d467 --- /dev/null +++ b/tst/manual/smol_vm_agent_check.py @@ -0,0 +1,93 @@ +"""Run explicitly with ``pytest tst/manual/smol_vm_agent_check.py`` on a KVM host with Docker and Smol installed.""" + +import uuid + +import httpx +import pytest + +from agent_env.a2a_agent.a2a_agent import DEFAULT_A2A_PORT +from agent_env.config import get_config, set_image_store +from agent_env.providers.sandbox_providers.smol_vm.provider import SmolVmSandboxProvider +from agent_env.store.routing import LocalRunImageStore +from agent_env.task.teardown import teardown_run +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep +from agent_env.task_step.task_steps.deploy_sandbox import DeploySandboxTaskStep +from tst.integration.store.local_parallel_deploy_test import local_stack # noqa: F401 +from tst.util.a2a_test_agent import build_test_agent_image, put_test_agent + + +@pytest.mark.asyncio +async def test_agent_answers_from_smol_vm_and_cleans_up(local_stack): + uid = uuid.uuid4().hex[:8] + agent = put_test_agent(f"smol-agent-{uid}") + context = TaskStepContext(instance_id=f"smol-agent-run-{uid}") + + try: + await DeploySandboxTaskStep( + id="box", version=None, sandbox_name="box", sandbox_mode="vm", sandbox_type="smol_vm", + exposed_ports=[DEFAULT_A2A_PORT], cpu=2, memory_mb=4096, disk_size_gb=16, + ).execute(context) + await DeployAgentTaskStep( + id="agent", version=None, env_ids=[], a2a_agent_id=agent.id, agent_name="solver", sandbox_name="box", + ).execute(context) + + assert len(context.deployed_agents) == 1 + url = context.deployed_agents[0].a2a_url + assert url == context.deployed_sandboxes[0].tunnel_urls[str(DEFAULT_A2A_PORT)] + async with httpx.AsyncClient() as client: + response = await client.get(f"{url}/.well-known/agent.json", timeout=30) + assert response.status_code == 200 + finally: + report = await teardown_run(context) + assert not report.still_up + + +@pytest.mark.asyncio +async def test_agent_deploys_into_its_own_smol_vm(local_stack): + agent = put_test_agent(f"smol-own-vm-{uuid.uuid4().hex[:8]}") + + try: + deployed = await agent.deploy(sandbox_type="smol_vm", cpu=2, memory=4096, disk_size_gb=16) + assert deployed.sandbox_type == "smol_vm" + async with httpx.AsyncClient() as client: + response = await client.get(f"{deployed.a2a_url}/.well-known/agent.json", timeout=30) + assert response.status_code == 200 + finally: + await agent.close() + + +@pytest.mark.asyncio +async def test_local_run_registry_image_pulls_inside_smol_vm(local_stack): + artifact = build_test_agent_image(f"smol-local-run-{uuid.uuid4().hex[:8]}") + registry = get_config().get_image_store() + set_image_store(LocalRunImageStore(registry, registry)) + sandbox = await SmolVmSandboxProvider().create_vm(cpu=2, memory=4096, disk_size_gb=16) + + try: + await sandbox.pull_images([artifact.image_name]) + code, _, stderr = await sandbox.exec_with_output("docker", "image", "inspect", artifact.image_name) + assert code == 0, stderr + finally: + await sandbox.terminate() + + +@pytest.mark.asyncio +async def test_container_mode_starts_image_with_environment_inside_smol_vm(local_stack): + agent = put_test_agent(f"smol-container-{uuid.uuid4().hex[:8]}") + context = TaskStepContext(instance_id=f"smol-container-run-{uuid.uuid4().hex[:8]}") + step = DeploySandboxTaskStep( + id="box", version=None, sandbox_name="box", sandbox_mode="container", sandbox_type="smol_vm", + image=agent.docker_image_artifact.image_name, port=DEFAULT_A2A_PORT, + env_vars={**agent.default_env_vars, "A2A_PORT": str(DEFAULT_A2A_PORT)}, + cpu=2, memory_mb=4096, disk_size_gb=16, + ) + + try: + await step.execute(context) + assert context.deployed_sandboxes[0].sandbox_mode == "container" + url = context.deployed_sandboxes[0].tunnel_urls[str(DEFAULT_A2A_PORT)] + assert (await agent._wait_for_agent_card(url))["name"] + finally: + report = await teardown_run(context) + assert not report.still_up diff --git a/tst/manual/smol_vm_gateway_check.py b/tst/manual/smol_vm_gateway_check.py new file mode 100644 index 00000000..7b8730d2 --- /dev/null +++ b/tst/manual/smol_vm_gateway_check.py @@ -0,0 +1,86 @@ +"""Run explicitly with ``pytest tst/manual/smol_vm_gateway_check.py`` on a KVM host with Docker and Smol installed.""" + +import json +import uuid + +import pytest +from agentenv_protocol import client as protocol_v1 + +from agent_env.a2a_agent.a2a_agent import DEFAULT_A2A_PORT +from agent_env.env import legacy_protocol +from agent_env.env.env import DeployedGatewayEnv, Env +from agent_env.env.envs.mcp_server import MCPServerEnv +from agent_env.env.envs.multi_env import MultiEnv +from agent_env.providers.sandbox_providers.smol_vm.provider import SmolVmSandboxProvider +from agent_env.task.teardown import teardown_run +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep +from agent_env.task_step.task_steps.deploy_sandbox import DeploySandboxTaskStep +from tst.integration.env.envs.multi_env_local_test import _items_universe +from tst.integration.env.envs.server_env_local_test import _put_items_image +from tst.integration.store.local_parallel_deploy_test import _seed_bootstrap_envs, local_stack # noqa: F401 +from tst.util.a2a_test_agent import put_test_agent + + +@pytest.mark.asyncio +async def test_gateway_restores_and_serves_data_on_smol_vm(local_stack): + _seed_bootstrap_envs() + uid = uuid.uuid4().hex[:8] + items = MCPServerEnv.put( + id=f"smol-items-{uid}", + environment_name="items", + docker_image_artifact=_put_items_image(f"smol-items-img-{uid}"), + ) + env = MultiEnv.put(id=f"smol-multi-{uid}", mcp_server_envs=[items], name="smol-suite") + + context = TaskStepContext(instance_id=f"smol-env-agent-run-{uid}") + + try: + deployed = await env.deploy(sandbox_type="smol_vm", ttl_seconds=900, disk_size_gb=16, cpu=2, memory_mb=4096) + assert isinstance(deployed, DeployedGatewayEnv) + assert deployed.sandbox_type == "smol_vm" + assert deployed.gateway_url.startswith("http://127.0.0.1:") + + restored = await Env.from_instance_id(deployed.instance_id) + await restored.load_environment_universe_artifact(_items_universe(uid)) + base_url = await legacy_protocol.v1_base_url(deployed, deployed.gateway_url, "items") + payload = (await protocol_v1.get_data(base_url)).model_dump()["parts"][0]["data"] + assert payload == {"items": ["snap-x", "snap-y"]} + + context.deployed_envs.append(deployed) + agent = put_test_agent(f"smol-agent-{uid}") + await DeploySandboxTaskStep( + id="box", version=None, sandbox_name="box", sandbox_mode="vm", sandbox_type="smol_vm", + exposed_ports=[DEFAULT_A2A_PORT], cpu=2, memory_mb=4096, disk_size_gb=16, + ).execute(context) + await DeployAgentTaskStep( + id="agent", version=None, env_ids=[env.id], a2a_agent_id=agent.id, + agent_name="solver", sandbox_name="box", + ).execute(context) + + agent_vm = await SmolVmSandboxProvider().get_sandbox(context.deployed_sandboxes[0].sandbox_id) + code, names, stderr = await agent_vm.exec_with_output("docker", "ps", "--format", "{{.Names}}") + assert code == 0, stderr + assert len(names.splitlines()) == 1 + reachable = SmolVmSandboxProvider.get_external_url(deployed.gateway_url) + code, status, stderr = await agent_vm.exec_with_output( + "docker", "exec", names.strip(), "python", "-c", + "import sys; from urllib.request import urlopen; print(urlopen(sys.argv[1], timeout=20).status)", + f"{reachable}/.well-known/agent-env.json", + ) + assert code == 0, stderr + assert status.strip() == "200" + + code, data, stderr = await agent_vm.exec_with_output( + "docker", "exec", names.strip(), "python", "-c", + "import asyncio, sys; from agentenv_protocol import client; " + "print(asyncio.run(client.get_data(sys.argv[1])).model_dump_json())", + SmolVmSandboxProvider.get_external_url(base_url), + ) + assert code == 0, stderr + assert json.loads(data)["parts"][0]["data"] == {"items": ["snap-x", "snap-y"]} + finally: + report = await teardown_run(context) + if env._sandbox is not None: + await env.close() + assert not report.still_up diff --git a/tst/unit/providers/sandbox_providers/network_policy_stamping_test.py b/tst/unit/providers/sandbox_providers/network_policy_stamping_test.py index 0b801a53..91155cfa 100644 --- a/tst/unit/providers/sandbox_providers/network_policy_stamping_test.py +++ b/tst/unit/providers/sandbox_providers/network_policy_stamping_test.py @@ -19,6 +19,8 @@ from agent_env.providers.sandbox_providers.modal_vm_sandbox import ModalVmSandboxProvider from agent_env.providers.sandbox_providers.sail_vm import _sdk as sail_sdk from agent_env.providers.sandbox_providers.sail_vm.provider import SailVmSandboxProvider +from agent_env.providers.sandbox_providers.smol_vm import provider as smol_module +from agent_env.providers.sandbox_providers.smol_vm.provider import SmolVmSandboxProvider from agent_env.providers.sandbox_providers.sandbox import Sandbox from agent_env.providers.sandbox_providers.sandbox_provider import _BUILTIN_SANDBOX_PROVIDERS @@ -86,11 +88,23 @@ async def _sail() -> Sandbox: return await provider.create_vm(exposed_ports=[], setup_for_gateway=False) +async def _smol_vm() -> Sandbox: + machine = MagicMock(id="agentenv-test") + with ( + patch.object(smol_module, "Machine", SimpleNamespace(create=lambda config, conn: machine)), + patch.object(smol_module, "_host_port", return_value=44112), + patch.object(SmolVmSandboxProvider, "_setup_docker", new_callable=AsyncMock), + patch.object(SmolVmSandboxProvider, "_configure_host_access", new_callable=AsyncMock), + ): + return await SmolVmSandboxProvider().create_vm(exposed_ports=[]) + + FACTORIES = { "modal": _modal, "modal_vm": _modal_vm, "e2b": _e2b, "sail_vm": _sail, + "smol_vm": _smol_vm, "local": _local, } diff --git a/tst/unit/providers/sandbox_providers/smol_vm_provider_test.py b/tst/unit/providers/sandbox_providers/smol_vm_provider_test.py new file mode 100644 index 00000000..1c3af69f --- /dev/null +++ b/tst/unit/providers/sandbox_providers/smol_vm_provider_test.py @@ -0,0 +1,188 @@ +"""Local SmolVM lifecycle rules that must hold even when cloud credentials are present.""" + +from __future__ import annotations + +import asyncio +import json +import threading +from unittest.mock import AsyncMock, MagicMock +from types import SimpleNamespace + +import pytest + +from agent_env.providers.sandbox_providers.smol_vm import provider as module +from agent_env.providers.sandbox_providers.smol_vm import sandbox as sandbox_module +from agent_env.providers.sandbox_providers.smol_vm.sandbox import SmolVmSandbox, _host_url +from agent_env.store.image_store.local_registry_image_store import LocalRegistryImageStore +from agent_env.store.object_store.local.store import LocalFilesystemObjectStore +from agent_env.store.routing import LocalRunImageStore + + +@pytest.mark.asyncio +async def test_agent_sandbox_returns_vm_to_load_saved_or_context_only_images(monkeypatch): + provider = module.SmolVmSandboxProvider() + create_vm = AsyncMock(return_value=SimpleNamespace(mode="vm")) + monkeypatch.setattr(provider, "create_vm", create_vm) + + sandbox = await provider.create_sandbox(image_name="local-image", port=8080, env={"A2A_PORT": "8080"}) + + assert sandbox.mode == "vm" + assert create_vm.await_args.kwargs["exposed_ports"] == [8080] + + +@pytest.mark.asyncio +async def test_local_run_proxies_its_underlying_loopback_registries(monkeypatch): + store = LocalRunImageStore( + LocalRegistryImageStore("localhost:5001"), LocalRegistryImageStore("localhost:5000"), + ) + monkeypatch.setattr(module, "get_config", lambda: SimpleNamespace(get_image_store=lambda: store)) + sandbox = SimpleNamespace(exec_script=AsyncMock()) + + await module.SmolVmSandboxProvider()._ensure_local_registry_proxy(sandbox) + + commands = [call.args[0] for call in sandbox.exec_script.await_args_list] + assert len(commands) == 2 + assert "TCP-LISTEN:5000" in commands[0] + assert "TCP-LISTEN:5001" in commands[1] + + +@pytest.mark.asyncio +async def test_vm_creation_selects_local_and_can_publish_an_unstarted_application(monkeypatch): + created = [] + + class FakeMachine: + id = "agentenv-test" + + def endpoint(self, port): + return SimpleNamespace(http_url=f"http://127.0.0.1:44000") + + def write_file(self, path, data): + created.append((path, data)) + + def delete(self): + created.append(("deleted", None)) + + def create(config, conn): + assert conn.target == "local" + assert config.wait_for_ports is False + assert config.resources.network is True + assert [(p.host, p.guest) for p in config.ports] == [(44000, 8080)] + return FakeMachine() + + async def setup(self, sandbox): + sandbox.extra_hosts = ("host.docker.internal:100.96.0.1",) + + monkeypatch.setattr(module, "Machine", SimpleNamespace(create=create)) + monkeypatch.setattr(module, "_host_port", lambda: 44000) + monkeypatch.setattr(module.SmolVmSandboxProvider, "_setup_docker", setup) + monkeypatch.setattr(module.SmolVmSandboxProvider, "_configure_host_access", setup) + sandbox = await module.SmolVmSandboxProvider().create_vm(exposed_ports=[8080], memory=1024) + assert sandbox.tunnel_urls[8080] == "http://127.0.0.1:44000" + assert ("/storage/agentenv-ports.json", b'{"8080": 44000}') in created + await sandbox.terminate() + assert ("deleted", None) in created + + +@pytest.mark.asyncio +async def test_cancellation_during_native_creation_reclaims_vm(monkeypatch): + started, release = threading.Event(), threading.Event() + deleted = [] + + class FakeMachine: + def delete(self): + deleted.append(True) + + def create(config, conn): + started.set() + assert release.wait(timeout=10) + return FakeMachine() + + monkeypatch.setattr(module, "Machine", SimpleNamespace(create=create)) + monkeypatch.setattr(module, "_host_port", lambda: 44000) + task = asyncio.create_task(module.SmolVmSandboxProvider().create_vm(memory=1024)) + assert await asyncio.to_thread(started.wait, 5) + task.cancel() + release.set() + with pytest.raises(asyncio.CancelledError): + await task + assert deleted == [True] + + +@pytest.mark.asyncio +async def test_reconnect_uses_explicit_agent_only_readiness(monkeypatch): + options = [] + + class FakeMachine: + id = "agentenv-test" + + def read_file(self, path): + assert path == "/storage/agentenv-ports.json" + return json.dumps({"8080": 44000}).encode() + + def endpoint(self, port): + assert port == 8080 + return SimpleNamespace(http_url="http://127.0.0.1:44000") + + def connect(name, conn): + assert name == "agentenv-test" + options.append(conn) + return FakeMachine() + + monkeypatch.setattr("agent_env.providers.sandbox_providers.smol_vm.sandbox.Machine", SimpleNamespace(connect=connect)) + sandbox = await SmolVmSandbox.reconnect("agentenv-test") + assert sandbox.tunnel_urls == {8080: "http://127.0.0.1:44000"} + assert options[0].target == "local" and options[0].wait_for_ports is False + + +def test_local_urls_are_rewritten_for_vm_host(): + assert _host_url("http://127.0.0.1:5000/v2/") == "http://host.smolvm.internal:5000/v2/" + assert _host_url("https://example.org/path") == "https://example.org/path" + + +@pytest.mark.asyncio +async def test_local_object_download_uses_a_scoped_https_grant(monkeypatch, tmp_path): + store = LocalFilesystemObjectStore(str(tmp_path)) + object_url = store.put("image.tar.gz", b"archive") + issued = MagicMock(return_value=SimpleNamespace(url="https://host.docker.internal:1234/v1/grants/token")) + monkeypatch.setattr(store, "issue_read_grant", issued) + monkeypatch.setattr(sandbox_module, "get_config", lambda: SimpleNamespace(get_object_store_at=lambda url: store)) + machine = MagicMock(id="agentenv-test") + sandbox = SmolVmSandbox(machine, {}) + sandbox.write_host_file = AsyncMock() + sandbox.exec_script = AsyncMock() + sandbox._write_unsigned_object = AsyncMock() + + await sandbox._download_object_to_vm(object_url, "/storage/image.tar.gz") + + issued.assert_called_once_with(object_url, expires_in=3600) + sandbox._write_unsigned_object.assert_not_called() + sandbox.write_host_file.assert_awaited_once() + command = sandbox.exec_script.await_args.args[0] + assert "SSL_CERT_FILE=/etc/agentenv-ca-bundle.pem" in command + assert "https://host.smolvm.internal:1234/v1/grants/token" in command + assert "-o /storage/image.tar.gz" in command + + +@pytest.mark.asyncio +async def test_unreachable_local_grant_falls_back_to_object_push(monkeypatch, tmp_path): + store = LocalFilesystemObjectStore(str(tmp_path), grant_bind_host="172.17.0.1") + object_url = store.put("image.tar.gz", b"archive") + monkeypatch.setattr(sandbox_module, "get_config", lambda: SimpleNamespace(get_object_store_at=lambda url: store)) + sandbox = SmolVmSandbox(MagicMock(id="agentenv-test"), {}) + sandbox._write_unsigned_object = AsyncMock() + + await sandbox._download_object_to_vm(object_url, "/storage/image.tar.gz") + sandbox._write_unsigned_object.assert_awaited_once_with(store, object_url, "/storage/image.tar.gz") + + +@pytest.mark.asyncio +async def test_loopback_advertised_grant_uses_object_push(monkeypatch, tmp_path): + store = LocalFilesystemObjectStore(str(tmp_path), grant_advertise_host="localhost") + object_url = store.put("image.tar.gz", b"archive") + monkeypatch.setattr(sandbox_module, "get_config", lambda: SimpleNamespace(get_object_store_at=lambda url: store)) + sandbox = SmolVmSandbox(MagicMock(id="agentenv-test"), {}) + sandbox._write_unsigned_object = AsyncMock() + + await sandbox._download_object_to_vm(object_url, "/storage/image.tar.gz") + + sandbox._write_unsigned_object.assert_awaited_once_with(store, object_url, "/storage/image.tar.gz") diff --git a/tst/unit/providers/sandbox_providers/test_url_rewrites.py b/tst/unit/providers/sandbox_providers/test_url_rewrites.py index 149d0c7b..632e6d71 100644 --- a/tst/unit/providers/sandbox_providers/test_url_rewrites.py +++ b/tst/unit/providers/sandbox_providers/test_url_rewrites.py @@ -115,9 +115,10 @@ def test_local_provider_leaves_a_url_whose_host_is_not_loopback(url): @pytest.mark.parametrize("sandbox_type, expected", [ - ("local", "http://host.docker.internal:4000"), ("modal", "http://localhost:4000"), (None, "http://localhost:4000"), + ("local", "http://host.docker.internal:4000"), + ("smol_vm", "http://host.smolvm.internal:4000"), ("modal", "http://localhost:4000"), (None, "http://localhost:4000"), ]) -def test_only_a_local_container_reaches_this_machine_by_another_name(sandbox_type, expected): +def test_a_local_or_smol_vm_container_reaches_this_machine_by_another_name(sandbox_type, expected): assert host_url_for("http://localhost:4000", sandbox_type) == expected diff --git a/tst/unit/store/local_grant_server_test.py b/tst/unit/store/local_grant_server_test.py index d88de1a1..1987b2b9 100644 --- a/tst/unit/store/local_grant_server_test.py +++ b/tst/unit/store/local_grant_server_test.py @@ -2,6 +2,7 @@ grant cases, and the bounds and refusals S3 would apply to the same requests.""" import os +import ssl import subprocess import sys import threading @@ -295,6 +296,17 @@ def test_the_root_defaults_to_the_state_root(): assert LocalFilesystemObjectStore().root == state_root() / "object_store" +def test_default_bridge_listener_also_serves_loopback_on_the_same_port(store, monkeypatch, tmp_path): + # A second loopback IP stands in for the Linux Docker bridge address. + monkeypatch.setattr(server_module, "default_bind_host", lambda: "127.0.0.2") + bridged = LocalFilesystemObjectStore(str(tmp_path / "bridged"), grant_bind_host="127.0.0.2", grant_advertise_host="localhost") + grant = bridged.issue_read_grant(bridged.put("k", b"both")) + with httpx.Client(verify=ssl.create_default_context(cafile=str(local_ca().bundle_path)), trust_env=False) as client: + assert client.get(grant.url).content == b"both" + assert client.get(grant.url.replace("localhost", "127.0.0.2")).content == b"both" + grant_server("127.0.0.2", "localhost").close() + + class TestDefaultBindHost: @pytest.fixture(autouse=True) def _fresh(self): diff --git a/tst/unit/store/local_tls_test.py b/tst/unit/store/local_tls_test.py index 9c557472..509e551f 100644 --- a/tst/unit/store/local_tls_test.py +++ b/tst/unit/store/local_tls_test.py @@ -38,10 +38,10 @@ def _handshake(server: ssl.SSLContext, *, cafile: Path, hostname: str) -> None: raise AssertionError("the handshake did not finish") -@pytest.mark.parametrize("hostname", ["host.docker.internal", "localhost", "127.0.0.1"]) +@pytest.mark.parametrize("hostname", ["host.docker.internal", "host.smolvm.internal", "localhost", "127.0.0.1"]) def test_a_server_certificate_verifies_for_the_names_it_was_made_for(hostname): ca = local_ca() - context = server_context(ca, ["localhost", "host.docker.internal", "127.0.0.1"]) + context = server_context(ca, ["localhost", "host.docker.internal", "host.smolvm.internal", "127.0.0.1"]) _handshake(context, cafile=ca.cert_path, hostname=hostname) _handshake(context, cafile=ca.bundle_path, hostname=hostname) @@ -132,7 +132,7 @@ def test_an_unreadable_ca_is_replaced(): @pytest.mark.parametrize("host", [ - "localhost", "host.docker.internal", "a.localhost", "127.0.0.1", "172.17.0.1", "192.168.5.2", "10.0.0.8", + "localhost", "host.docker.internal", "host.smolvm.internal", "a.localhost", "127.0.0.1", "172.17.0.1", "192.168.5.2", "10.0.0.8", "::1", "fd00::1", ]) def test_local_hosts_are_accepted(host): diff --git a/tst/unit/store/test_object_transfer_grants.py b/tst/unit/store/test_object_transfer_grants.py index 144775a4..71a98ce1 100644 --- a/tst/unit/store/test_object_transfer_grants.py +++ b/tst/unit/store/test_object_transfer_grants.py @@ -9,15 +9,36 @@ from tst.util.granting_object_store import GrantingObjectStore -def test_local_store_grants_reach_only_local_sandboxes(tmp_path) -> None: +def test_local_store_grants_reach_local_and_smol_vm_sandboxes(tmp_path) -> None: store = LocalFilesystemObjectStore(str(tmp_path)) assert store.supports_transfer_grants assert store.grants_reach("local") + assert store.grants_reach("smol_vm") assert not store.grants_reach("modal") assert not store.grants_reach(None) +def test_custom_bridge_only_grants_are_not_claimed_reachable_from_smol_vm(tmp_path) -> None: + store = LocalFilesystemObjectStore(str(tmp_path), grant_bind_host="172.17.0.1") + assert store.grants_reach("local") + assert not store.grants_reach("smol_vm") + + +def test_custom_private_advertised_host_is_not_claimed_reachable_from_smol_vm(tmp_path) -> None: + store = LocalFilesystemObjectStore(str(tmp_path), grant_advertise_host="172.17.0.1") + assert store.grants_reach("local") + assert not store.grants_reach("smol_vm") + + +@pytest.mark.parametrize("host", ["localhost", "127.0.0.1"]) +def test_loopback_grant_url_is_not_claimed_reachable_from_agent_container(tmp_path, host) -> None: + store = LocalFilesystemObjectStore(str(tmp_path), grant_advertise_host=host) + + assert store.grants_reach("local") + assert not store.grants_reach("smol_vm") + + def test_local_store_grants_can_be_turned_off(tmp_path) -> None: assert not LocalFilesystemObjectStore(str(tmp_path), grants="off").supports_transfer_grants with pytest.raises(ValueError, match="grants must be one of"): diff --git a/tst/unit/task_step/deploy_sandbox_image_test.py b/tst/unit/task_step/deploy_sandbox_image_test.py index bbc038ea..52ab0441 100644 --- a/tst/unit/task_step/deploy_sandbox_image_test.py +++ b/tst/unit/task_step/deploy_sandbox_image_test.py @@ -1,12 +1,14 @@ """A container-mode deploy_sandbox runs its image by name, so a provider whose sandboxes are elsewhere can't pull one from this machine's registry; one on this machine can, and a chain skips the rest for it.""" +from types import SimpleNamespace from unittest.mock import AsyncMock, patch import pytest from agent_env.providers.sandbox_providers.chained_sandbox_provider import ChainedSandboxProvider from agent_env.providers.sandbox_providers.sandbox import NetworkPolicy +from agent_env.providers.sandbox_providers.sandbox_provider import SandboxProvider from agent_env.task_step.context import TaskStepContext from agent_env.task_step.task_steps.deploy_sandbox import DeploySandboxTaskStep @@ -17,10 +19,46 @@ def _provider(on_this_machine: bool): sandbox.tunnel_urls, sandbox.vnc_url, sandbox.network_policy = {}, None, NetworkPolicy() provider = AsyncMock(ON_THIS_MACHINE=on_this_machine) provider.links = (provider,) - provider.create_sandbox.return_value = sandbox + provider.create_container.return_value = sandbox return provider +class _CustomContainerProvider(SandboxProvider): + ON_THIS_MACHINE = True + + def __init__(self): + self.calls = [] + + async def create_sandbox(self, **kwargs): + self.calls.append(kwargs) + return SimpleNamespace( + sandbox_id="custom-1", mode="container", type="custom", tunnel_urls={}, + vnc_url=None, network_policy=NetworkPolicy(), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("chained", [False, True], ids=["direct", "fallback-chain"]) +async def test_custom_container_only_provider_still_deploys(chained): + step = DeploySandboxTaskStep( + id="s", version=1, sandbox_name="box", sandbox_mode="container", + image="example.com/box:v1", port=80, env_vars={"MODE": "test"}, + ) + custom = _CustomContainerProvider() + if chained: + first = _provider(on_this_machine=True) + first.create_container.side_effect = RuntimeError("unavailable") + provider = ChainedSandboxProvider([first, custom]) + else: + provider = custom + + with patch("agent_env.providers.sandbox_providers.sandbox_provider.get_sandbox_provider", return_value=provider): + context = await step.execute(TaskStepContext()) + + assert custom.calls[0]["env"] == {"MODE": "test"} + assert context.deployed_sandboxes[0].sandbox_mode == "container" + + @pytest.mark.asyncio @pytest.mark.parametrize("on_this_machine", [False, True], ids=["elsewhere", "on-this-machine"]) async def test_a_container_sandbox_pulls_an_image_from_this_machines_registry_only_on_this_machine(on_this_machine): @@ -31,13 +69,13 @@ async def test_a_container_sandbox_pulls_an_image_from_this_machines_registry_on with patch("agent_env.providers.sandbox_providers.sandbox_provider.get_sandbox_provider", return_value=provider): if on_this_machine: await step.execute(TaskStepContext()) - provider.create_sandbox.assert_awaited_once() + provider.create_container.assert_awaited_once() else: with pytest.raises(ValueError, match="Can't deploy sandbox 'box' with AsyncMock, whose sandboxes pull it by " "name: localhost:5000/box:v1 is in a registry on this machine, which a " "sandbox elsewhere can't pull from"): await step.execute(TaskStepContext()) - provider.create_sandbox.assert_not_awaited() + provider.create_container.assert_not_awaited() @pytest.mark.asyncio @@ -50,6 +88,25 @@ async def test_a_chain_pulls_an_image_from_this_machines_registry_on_its_provide return_value=ChainedSandboxProvider([elsewhere, here])): await step.execute(TaskStepContext()) - elsewhere.create_sandbox.assert_not_awaited() - here.create_sandbox.assert_awaited_once() + elsewhere.create_container.assert_not_awaited() + here.create_container.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_a_container_deploy_falls_back_when_the_first_provider_fails(): + step = DeploySandboxTaskStep( + id="s", version=1, sandbox_name="box", sandbox_mode="container", + image="localhost:5000/box:v1", port=80, env_vars={"MODE": "test"}, + ) + first, second = _provider(on_this_machine=True), _provider(on_this_machine=True) + first.create_container.side_effect = RuntimeError("unavailable") + context = TaskStepContext() + + with patch("agent_env.providers.sandbox_providers.sandbox_provider.get_sandbox_provider", + return_value=ChainedSandboxProvider([first, second])): + await step.execute(context) + first.create_container.assert_awaited_once() + second.create_container.assert_awaited_once() + assert second.create_container.await_args.kwargs["env"] == {"MODE": "test"} + assert context.deployed_sandboxes[0].sandbox_mode == "container" diff --git a/tst/unit/task_step/deploy_sandbox_network_policy_test.py b/tst/unit/task_step/deploy_sandbox_network_policy_test.py index da4b7b61..82c051d4 100644 --- a/tst/unit/task_step/deploy_sandbox_network_policy_test.py +++ b/tst/unit/task_step/deploy_sandbox_network_policy_test.py @@ -46,7 +46,7 @@ def test_an_unknown_mode_is_rejected_when_the_task_is_built(): @pytest.mark.asyncio @pytest.mark.parametrize("mode, create_call, extra", [ ("vm", "create_vm", {}), - ("container", "create_sandbox", {"port": 8000}), # container mode requires image + port + ("container", "create_container", {"port": 8000}), # container mode requires image + port ]) async def test_the_policy_reaches_both_create_paths(mode, create_call, extra): step = _step(sandbox_mode=mode, image="img", network_policy=ALLOWLIST, **extra) diff --git a/tst/unit/task_step/test_deploy_ttl_from_budget.py b/tst/unit/task_step/test_deploy_ttl_from_budget.py index 74f062c5..bd1e60eb 100644 --- a/tst/unit/task_step/test_deploy_ttl_from_budget.py +++ b/tst/unit/task_step/test_deploy_ttl_from_budget.py @@ -55,7 +55,7 @@ async def test_deploy_sandbox_uses_budget_override(): async def test_deploy_sandbox_container_mode_uses_budget_override(): sb = MagicMock(sandbox_id="s1", mode="container", type="remote", tunnel_urls={}, vnc_url=None) provider = MagicMock() - provider.create_sandbox = AsyncMock(return_value=sb) + provider.create_container = AsyncMock(return_value=sb) step = DeploySandboxTaskStep( id="d", version=None, sandbox_name="sbx", sandbox_mode="container", image="img", port=8080, ttl_seconds=1800, @@ -63,7 +63,7 @@ async def test_deploy_sandbox_container_mode_uses_budget_override(): ctx = TaskStepContext(metadata={"user_overrides": {"ttl_seconds": 30600}}) with patch("agent_env.providers.sandbox_providers.sandbox_provider.get_sandbox_provider", return_value=provider): await step.execute(ctx) - assert provider.create_sandbox.await_args.kwargs["timeout"] == 30600 + assert provider.create_container.await_args.kwargs["timeout"] == 30600 @pytest.mark.asyncio diff --git a/uv.lock b/uv.lock index 66ca8166..a6cd81c0 100644 --- a/uv.lock +++ b/uv.lock @@ -5,8 +5,8 @@ resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", "python_full_version >= '3.14' and sys_platform != 'win32'", "python_full_version == '3.13.*' and sys_platform == 'win32'", - "python_full_version < '3.13' and sys_platform == 'win32'", "python_full_version == '3.13.*' and sys_platform != 'win32'", + "python_full_version < '3.13' and sys_platform == 'win32'", "python_full_version < '3.13' and sys_platform != 'win32'", ] @@ -88,6 +88,7 @@ dev = [ { name = "pytest-xdist" }, { name = "requests" }, { name = "sail" }, + { name = "smolmachines" }, { name = "uvicorn" }, ] explorer = [ @@ -104,11 +105,14 @@ gcp = [ sail = [ { name = "sail" }, ] +smol = [ + { name = "smolmachines" }, +] [package.metadata] requires-dist = [ { name = "a2a-sdk", specifier = "==0.3.26" }, - { name = "agentenv-framework", extras = ["aws", "explorer", "gcp", "sail"], marker = "extra == 'dev'" }, + { name = "agentenv-framework", extras = ["aws", "explorer", "gcp", "sail", "smol"], marker = "extra == 'dev'" }, { name = "agentenv-framework-protocol", editable = "packages/agentenv-protocol" }, { name = "boto3", marker = "extra == 'aws'", specifier = ">=1.35.0" }, { name = "botocore", marker = "extra == 'aws'", specifier = ">=1.35.0" }, @@ -142,11 +146,12 @@ requires-dist = [ { name = "pyyaml", specifier = ">=6.0" }, { name = "requests", marker = "extra == 'gcp'", specifier = ">=2.31" }, { name = "sail", marker = "extra == 'sail'", specifier = "~=0.12.8" }, + { name = "smolmachines", marker = "extra == 'smol'", specifier = ">=1.27.0" }, { name = "starlette", specifier = ">=0.40" }, { name = "uvicorn", specifier = ">=0.47" }, { name = "uvicorn", marker = "extra == 'explorer'", specifier = ">=0.29" }, ] -provides-extras = ["aws", "explorer", "gcp", "sail", "dev"] +provides-extras = ["aws", "explorer", "gcp", "sail", "dev", "smol"] [[package]] name = "agentenv-framework-protocol" @@ -3372,6 +3377,16 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, ] +[[package]] +name = "smolmachines" +version = "1.27.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/58/fe/bfb482f18eae870518335333b6ec30cf8f9f9e60051dfdec3ec37d9e763f/smolmachines-1.27.0-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:c0134c621e6a48161b378d8331016b8f62747fd45f6c7b8aeb09ef0899e92bec", size = 45937433, upload-time = "2026-10-12T03:31:01.34Z" }, + { url = "https://files.pythonhosted.org/packages/6a/a7/15aaa80bab6a1db909c310b444f7826f89ee2576c3494b19460f76aa413a/smolmachines-1.27.0-cp39-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:8f68e6b88f4b5465a4fea9a9387a3bf40b13cf90f10ebdec572b1c2e7f4cfc32", size = 44827771, upload-time = "2026-10-12T03:31:05.607Z" }, + { url = "https://files.pythonhosted.org/packages/c2/49/fb0e9ab762b75329ca4c6872a753404e942720d38b8525b5055792b6c611/smolmachines-1.27.0-cp39-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:fbcf8fd2956a2df237ae783b9e06fd300dc5a5bb017c4a1be26eedb0f709a5b8", size = 46527401, upload-time = "2026-10-12T03:31:09.721Z" }, +] + [[package]] name = "sniffio" version = "1.3.1"