feat(images): read image documents with no tarball, and pull them by name

src/agent_env/artifact/artifacts/docker_image.py

@@ -23,7 +23,7 @@
 
 from agent_env.artifact.artifact import Artifact, _write_twin
 from agent_env.store.ids import fs_safe, image_repository, is_local_id
-from agent_env.store.image_store.oci_registry_credentials import is_loopback_host, registry_host_from_ref
+from agent_env.store.image_store.oci_registry_credentials import is_loopback_host, names_registry, registry_host_from_ref
 from agent_env.utils.deprecation import OMITTED, renamed_keyword
 
 logger = logging.getLogger(__name__)
@@ -63,16 +63,19 @@ def _git_clone_commands(owner: str, repo: str, ref: str | None, token: str | Non
     ]
 
 class DockerImageArtifact(Artifact):
-    """A Docker image artifact stored as tar.gz in the object store."""
+    """A Docker image artifact: a tar.gz of the image in the object store, or, without one, a registry reference a
+    sandbox pulls."""
 
     model_config = ConfigDict(populate_by_name=True)
 
     DOCKER_SAVE_TIMEOUT_SECONDS: ClassVar[int] = 900
 
     type: Literal["docker_image"] = "docker_image"
     description: str = Field(description="Description of the Docker image")
-    image_name: str = Field(description="Docker image name/tag")
-    tar_gz_object_url: str = Field(alias="tar_gz_s3_url", description="Object-store locator of the tar.gz file")
+    image_name: str = Field(description="Docker image name/tag; with no tar.gz, the registry reference pulled")
+    tar_gz_object_url: str | None = Field(
+        default=None, alias="tar_gz_s3_url", description="Object-store locator of the tar.gz file, if it has one"
+    )
     build_context_object_url: str | None = Field(default=None, alias="build_context_s3_url", description="Object-store locator of the build context tar.gz")
 
     # No return annotation: pydantic builds the serialization schema from one, and a dict drops the fields.
@@ -188,6 +191,8 @@ def put_tar(
         )
         if tar_gz_object_url is None:
             raise TypeError(f"{owner}() missing required keyword argument: 'tar_gz_object_url'")
+        if not tar_gz_object_url:
+            raise ValueError(f"{owner}(): tar_gz_object_url is empty")
         for url in (tar_gz_object_url, build_context_object_url):
             if url:
                 get_config().check_object_url(id, url)
@@ -205,8 +210,18 @@ def put_tar(
 
     def load(self) -> bytes:
         from agent_env.artifact.store import get_artifact_store
+        if not self.tar_gz_object_url:
+            raise ValueError(f"{self.id!r} v{self.version} has no tar.gz; its image is pulled from {self.image_name}")
         return get_artifact_store().get_object(self.tar_gz_object_url)
 
+    def load_problem(self) -> str | None:
+        """Why no sandbox can get this image, or None when one can: a tar.gz is loaded, and with none, ``image_name``
+        is pulled, so it must name its registry."""
+        if self.tar_gz_object_url or names_registry(self.image_name):
+            return None
+        return (f"{self.id!r} v{self.version} has no tar.gz, and its image name {self.image_name!r} doesn't name a "
+                "registry to pull it from")
+
     @classmethod
     async def put_from_github(
         cls,

src/agent_env/bundle/preflight.py

@@ -95,6 +95,7 @@ class _Image:
 
     what: str
     local_only: str | None  # why only this machine has it, or None
+    unloadable: str | None = None  # why no sandbox can get it, or None
 
 
 @dataclass(frozen=True)
@@ -179,6 +180,9 @@ def _env(self, where: str, step: DeployEnvTaskStep) -> None:
         else:  # a plugin's provider, or one deploy_env's own preflight refuses for this env
             return
         env_id, images = step.env_id, deployment.images
+        if GATEWAY in kinds and (loaders := [link for link in _links(provider)
+                                             if not isinstance(link, ModalSandboxProvider)]):
+            self._loadable(where, loaders, images)  # a gateway VM loads its images; Modal's runs each by name
         if _local_link(provider):
             self.infra |= kinds
             self.docker_users.append(where)
@@ -256,8 +260,7 @@ def _planned_image(self, env_id: str, ref: Reference) -> _Image | None:
         if ref.id in self.written:
             return None  # another of the bundle's writes, which materialize refuses or writes first
         try:
-            return _Image(what, _local_only(DockerImageArtifact.get(ref.id, self._planned_version(ref.kind, ref.id,
-                                                                                                    ref.version))))
+            return _image(what, DockerImageArtifact.get(ref.id, self._planned_version(ref.kind, ref.id, ref.version)))
         except NotFoundError:
             return None  # the plan reports a store image that isn't there
 
@@ -288,10 +291,14 @@ def _agent_deploy(self, where: str, provider: SandboxProvider, agent_id: str | N
         if agent_id is None:
             self.default_agent_users.append((where, unnamed))
             agent_id = get_config().get_default_a2a_agent_id()
-        if remote := _remote_links(provider):
-            image = self._agent_image(agent_id, version)
-            if image is not None:
-                self._reachable(where, remote, [image])
+        loaders = [link for link in _links(provider) if not isinstance(link, (LocalSandboxProvider, ModalSandboxProvider))]
+        remote = _remote_links(provider)
+        if not (loaders or remote) or (image := self._agent_image(agent_id, version)) is None:
+            return
+        if loaders:
+            self._loadable(where, loaders, [image])  # those two run an agent in a container, by image name
+        if remote:
+            self._reachable(where, remote, [image])
 
     def _sandbox(self, where: str, step: DeploySandboxTaskStep) -> None:
         provider = _provider(self.sandbox or step.sandbox_type, get_sandbox_provider)
@@ -308,8 +315,8 @@ def _sandbox(self, where: str, step: DeploySandboxTaskStep) -> None:
     # What a deploy runs
 
     def _agent_image(self, agent_id: str, version: int | None) -> _Image | None:
-        """The image the agent ``agent_id`` runs, or None when the store doesn't hold the agent, which the plan or
-        ``_default_agent`` reports."""
+        """The image the agent ``agent_id`` runs, or None when the store doesn't hold the agent or can't read it, which
+        the plan or ``_default_agent`` reports."""
         what = f"agent {agent_id!r}'s image"
         if (agent := self.agents.get(agent_id)) is not None:
             image_id, image_version = parse_toml_ref(EntityKind.ARTIFACT, agent.config.get("image"))
@@ -319,12 +326,20 @@ def _agent_image(self, agent_id: str, version: int | None) -> _Image | None:
             if image_id in self.written:
                 return None  # another of the bundle's writes, which materialize refuses or writes first
             planned = self._planned_version(EntityKind.ARTIFACT, image_id, image_version)
-            return _Image(what, _local_only(DockerImageArtifact.get(image_id, planned)))
+            return _image(what, DockerImageArtifact.get(image_id, planned))
         try:
-            return _Image(what, _local_only(A2AAgent.get(agent_id, version).docker_image_artifact))
-        except NotFoundError:
+            return _image(what, A2AAgent.get(agent_id, version).docker_image_artifact)
+        except (NotFoundError, ValueError, KeyError, TypeError):
             return None
 
+    def _loadable(self, where: str, loaders: list[SandboxProvider], images: list[_Image]) -> None:
+        """Refuse each of ``images`` a VM on ``loaders`` can't get: it loads an image's tar.gz, or pulls an image with
+        none by name."""
+        for image in images:
+            if image.unloadable:
+                self._problem(where, f"deploys {image.what} on the {_shown(loaders[0])} sandbox provider, which can't "
+                                     f"load it: {image.unloadable}")
+
     def _reachable(self, where: str, remote: list[SandboxProvider], images: list[_Image]) -> None:
         for image in images:
             if image.local_only:
@@ -341,10 +356,12 @@ def _remote_infra(self, where: str, provider: SandboxProvider, kinds: set[str])
                                      f"{env_id!r}, and the store doesn't hold it; agent-env builds it only for the local "
                                      f"sandbox provider, so put it in a store that provider can reach (`{put_command(kind)}`)")
                 continue
-            images = ([env.db_docker_image_artifact, env.db_web_docker_image_artifact, env.db_mcp_docker_image_artifact]
-                      if isinstance(env, ServiceDBEnv) else [env.docker_image_artifact])
-            self._reachable(where, [provider], [_Image(f"the {kind} env {env_id!r}'s image {image.id!r}", _local_only(image))
-                                                for image in images])
+            artifacts = ([env.db_docker_image_artifact, env.db_web_docker_image_artifact,
+                          env.db_mcp_docker_image_artifact] if isinstance(env, ServiceDBEnv) else [env.docker_image_artifact])
+            images = [_image(f"the {kind} env {env_id!r}'s image {image.id!r}", image) for image in artifacts]
+            if not isinstance(provider, ModalSandboxProvider):  # whose containers run by name, or are swapped out
+                self._loadable(where, [provider], images)
+            self._reachable(where, [provider], images)
 
     def _default_agent(self) -> None:
         if not self.default_agent_users:
@@ -373,8 +390,7 @@ def _stored_images(env: Env) -> tuple[list[_Image], bool]:
     """The images a store env runs through a gateway, and whether it has websites."""
     topology = _gateway_topology(env)
     images = [*topology.mcp_server_images, *(topology.website_images or [])]
-    return [_Image(f"env {env.id!r}'s image {image.id!r}", _local_only(image)) for image in images], bool(
-        topology.website_configs)
+    return [_image(f"env {env.id!r}'s image {image.id!r}", image) for image in images], bool(topology.website_configs)
 
 
 def _state_type(step: DeployEnvTaskStep) -> str | None:
@@ -422,12 +438,16 @@ def _name(provider: SandboxProvider) -> str:
     return cls.__name__
 
 
+def _image(what: str, image: DockerImageArtifact) -> _Image:
+    return _Image(what, _local_only(image), image.load_problem())
+
+
 def _local_only(image: DockerImageArtifact | str) -> str | None:
     """Why only this machine has ``image``: its reference names a registry on this machine, or it's saved in this
     machine's object store. None when neither."""
     ref = image if isinstance(image, str) else image.image_name
     if is_loopback_host(registry_host_from_ref(ref)):
         return f"{ref} is in a registry on this machine"
-    if not isinstance(image, str) and urlparse(image.tar_gz_object_url).scheme == "file":
+    if not isinstance(image, str) and image.tar_gz_object_url and urlparse(image.tar_gz_object_url).scheme == "file":
         return f"{image.id!r} is saved in this machine's object store"
     return None

src/agent_env/providers/sandbox_providers/e2b/sandbox.py

@@ -291,8 +291,8 @@ async def apply_network_policy(self, policy: NetworkPolicy) -> None:
         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.
+    async def _load_tarballs(self, artifacts: list) -> None:
+        """Load image tarballs 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
@@ -301,16 +301,14 @@ async def load_docker_images(self, artifacts: list) -> None:
         :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)
+            await super()._load_tarballs(artifacts)
             return
 
         signed_urls = await self._signed_image_urls(artifacts)

src/agent_env/providers/sandbox_providers/sail_vm/sandbox.py

@@ -464,10 +464,8 @@ async def _allow_download_hosts(self, urls: list[str | None], purpose: str) -> N
                 )
             await self._apply_policy(policy.with_hosts(missing))
 
-    async def load_docker_images(self, artifacts: list) -> None:
-        """Load images, first adding their signed-download hosts to a restrictive policy."""
-        if not artifacts:
-            return
+    async def _load_tarballs(self, artifacts: list) -> None:
+        """Load image tarballs, first adding their signed-download hosts to a restrictive policy."""
         self._known_policy("load Docker images")
         signed_urls = await self._signed_image_urls(artifacts)
         await self._allow_download_hosts(signed_urls, "load Docker images")

src/agent_env/providers/sandbox_providers/sandbox.py

@@ -22,6 +22,7 @@
 from typing import IO, TYPE_CHECKING, Any, AsyncIterator, Callable, Iterable, Optional
 
 from agent_env.config import get_config
+from agent_env.store.image_store.oci_registry_credentials import registry_host_from_ref
 from agent_env.utils.deprecation import warn_deprecated
 from agent_env.utils.paths import validate_relative_filename
 
@@ -284,11 +285,46 @@ async def wait_for_vm(self) -> None:
         )
 
     async def load_docker_images(self, artifacts: list) -> None:
-        """Load Docker images from DockerImageArtifacts into the sandbox in parallel."""
+        """Put each DockerImageArtifact's image on the VM: those with a tar.gz are loaded from it, in parallel, and the
+        rest are pulled by image name. An image no sandbox can get is refused before anything is loaded."""
         if not artifacts:
             return
+        if problems := [problem for artifact in artifacts if (problem := artifact.load_problem())]:
+            raise RuntimeError(f"Can't load images: {'; '.join(problems)}")
+        if tarballs := [artifact for artifact in artifacts if artifact.tar_gz_object_url]:
+            await self._load_tarballs(tarballs)
+        if pulled := [artifact.image_name for artifact in artifacts if not artifact.tar_gz_object_url]:
+            await self.pull_images(pulled)
+
+    async def _load_tarballs(self, artifacts: list) -> None:
+        """Load the images of ``artifacts``, each with a tar.gz, in parallel."""
         await self._load_docker_images(artifacts, await self._signed_image_urls(artifacts))
 
+    async def pull_images(self, image_names: list[str]) -> None:
+        """``docker pull`` each of ``image_names``, concurrently, after logging in to each registry the image store
+        holds credentials for: once per registry, since minting a login can be a network round trip. The first pull
+        to fail cancels the rest, since one that stalls would otherwise hold the failure back; the command a cancelled
+        exec started keeps running on the VM until it ends or the VM does. A network policy restricting egress isn't
+        widened for them: it must allow the registries itself."""
+        image_names = list(dict.fromkeys(image_names))
+        store = get_config().get_image_store()
+        one_per_registry = {registry_host_from_ref(image_name): image_name for image_name in reversed(image_names)}
+        for image_name in one_per_registry.values():
+            auth = await asyncio.to_thread(store.auth, image_name)
+            if auth is not None:
+                await self.exec_script(
+                    f"echo {shlex.quote(auth.password)} | docker login "
+                    f"--username {shlex.quote(auth.username)} --password-stdin {shlex.quote(auth.registry)}"
+                )
+        pulls = [asyncio.ensure_future(_pull(self, image_name)) for image_name in image_names]
+        try:
+            await asyncio.gather(*pulls)
+        except BaseException:
+            for pull in pulls:
+                pull.cancel()
+            await asyncio.gather(*pulls, return_exceptions=True)
+            raise
+
     @staticmethod
     async def _signed_image_urls(artifacts: list) -> list[str | None]:
         """Each image tarball's signed URL, or None where the store cannot sign one. Signed
@@ -455,6 +491,19 @@ def port_bindings(host_ips: Iterable[str], host_port: int, container_port: int)
     return [f"{ip}:{host_port}:{container_port}" for ip in host_ips] or [f"{host_port}:{container_port}"]
 
 
+async def _pull(sandbox: VmSandbox, image_name: str) -> None:
+    """``docker pull image_name``. An image built for linux/amd64 only has nothing for an arm64 host (an
+    Apple Silicon Mac running the local provider), so that pull falls back to the amd64 image, which the
+    host's Docker runs emulated."""
+    try:
+        await sandbox.exec_script(f"docker pull {shlex.quote(image_name)}")
+    except RuntimeError as e:
+        if "no matching manifest" not in str(e):
+            raise
+        logger.warning("%s has no image for this host's platform; pulling linux/amd64, which runs emulated", image_name)
+        await sandbox.exec_script(f"docker pull --platform linux/amd64 {shlex.quote(image_name)}")
+
+
 async def stage_files_into_container(sandbox: Sandbox, file_artifacts: dict[str, Any], destination: str) -> dict[str, str]:
     """Write each object-store file into the sandbox's container under ``destination``; returns ``{name: path}``.
     A VM-backed sandbox stages through its host into ``container_name``; a container sandbox is written directly."""

src/agent_env/providers/sandbox_providers/sandbox_provider.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import asyncio
 import logging
 import os
 import shlex
@@ -36,19 +35,6 @@ class SandboxProviderTypeError(ConfigError):
     """A config-registered provider produced a Sandbox whose ``.type`` != its ``[sandbox.providers.<name>]`` key."""
 
 
-async def _pull(sandbox: VmSandbox, image_name: str) -> None:
-    """``docker pull image_name``. An image built for linux/amd64 only has nothing for an arm64 host (an
-    Apple Silicon Mac running the local provider), so that pull falls back to the amd64 image, which the
-    host's Docker runs emulated."""
-    try:
-        await sandbox.exec_script(f"docker pull {shlex.quote(image_name)}")
-    except RuntimeError as e:
-        if "no matching manifest" not in str(e):
-            raise
-        logger.warning("%s has no image for this host's platform; pulling linux/amd64, which runs emulated", image_name)
-        await sandbox.exec_script(f"docker pull --platform linux/amd64 {shlex.quote(image_name)}")
-
-
 class SandboxProvider(ABC):
     """Compute backend that provisions sandboxes. Selected via [sandbox] default / build_sandbox_provider(); swap via set_sandbox_provider()."""
 
@@ -111,15 +97,7 @@ async def create_container(
             network_policy=network_policy,
         )
         try:
-            from agent_env.config import get_config
-
-            auth = await asyncio.to_thread(get_config().get_image_store().auth, image_name)
-            if auth is not None:
-                await sandbox.exec_script(
-                    f"echo {shlex.quote(auth.password)} | docker login "
-                    f"--username {shlex.quote(auth.username)} --password-stdin {shlex.quote(auth.registry)}"
-                )
-            await _pull(sandbox, image_name)
+            await sandbox.pull_images([image_name])
             await self._start_container(sandbox, image_name=image_name, port=port, env=env)
             sandbox.mode = SANDBOX_MODE_CONTAINER
             return sandbox

src/agent_env/store/image_store/__init__.py

@@ -16,6 +16,7 @@
     OciRegistryCredentials,
     RegistryAuth,
     SecretStoreCredentials,
+    names_registry,
     normalize_registry_host,
     registry_host_from_ref,
 )
@@ -25,6 +26,7 @@
     "OciRegistryImageStore",
     "RegistryAuth",
     "OciRegistryCredentials",
+    "names_registry",
     "normalize_registry_host",
     "registry_host_from_ref",
     "EcrCredentials",