diff --git a/libs/python/cua-sandbox/cua_sandbox/builder/build.py b/libs/python/cua-sandbox/cua_sandbox/builder/build.py index 932eab43ad..d1b15bc2eb 100644 --- a/libs/python/cua-sandbox/cua_sandbox/builder/build.py +++ b/libs/python/cua-sandbox/cua_sandbox/builder/build.py @@ -317,6 +317,31 @@ async def build_user_image( return user_path +async def resolve_backing_disk(image: Image) -> Path: + """Resolve the disk a local session overlays, preferring the registry containerDisk. + + Built-in images (and explicit ``Image.from_registry(...)`` refs) map to a + KubeVirt containerDisk in the registry — the very image Fleet cloud boots. Pulling + it keeps local QEMU runs on the same disk as the cloud instead of a separately + built base. Images with no registry counterpart fall back to the local base build. + """ + from cua_sandbox.image import cloud_registry_image + + ref = cloud_registry_image(image) + if ref is not None: + from cua_sandbox.registry.container_disk import pull_container_disk + + logger.info(f"Resolving containerDisk {ref} for local session...") + try: + # Network + multi-GB extraction: keep it off the event loop. + return await asyncio.to_thread(pull_container_disk, ref) + except FileNotFoundError as exc: + # Not a containerDisk (e.g. a lume/tart/qemu-format VM image) — fall back. + logger.info(f"{ref} is not a containerDisk ({exc}); falling back to base image") + + return await ensure_base_image(image.os_type, image.version) + + async def create_session_disk( image: Image, name: str, @@ -326,8 +351,9 @@ async def create_session_disk( """Create a session overlay for a sandbox run. If the image has layers and a cached user image exists, overlay on that. - Otherwise overlay on the base image. If no base exists, returns the - image's _disk_path directly (no overlay). + Otherwise overlay on the registry containerDisk (the same disk Fleet cloud + boots) or on a locally built base image. If the image carries a direct disk + path and no layers, that disk is returned as-is (no overlay). Returns the disk path to boot. """ @@ -341,8 +367,7 @@ async def create_session_disk( elif image._disk_path: backing = Path(image._disk_path) else: - # Auto-build base if it doesn't exist - backing = await ensure_base_image(image.os_type, image.version) + backing = await resolve_backing_disk(image) # If there are user layers, check for cached user image if image._layers: diff --git a/libs/python/cua-sandbox/cua_sandbox/image.py b/libs/python/cua-sandbox/cua_sandbox/image.py index 22d8cc76dc..e66520bca7 100644 --- a/libs/python/cua-sandbox/cua_sandbox/image.py +++ b/libs/python/cua-sandbox/cua_sandbox/image.py @@ -28,6 +28,8 @@ logger = logging.getLogger(__name__) +DEFAULT_LINUX_REGISTRY_IMAGE = "public.ecr.aws/k5j5w0x5/cua-ubuntu-24.04:main-38352d34" + _IMAGE_CACHE = Path.home() / ".cua" / "cua-sandbox" / "image-cache" @@ -477,3 +479,17 @@ def __repr__(self) -> str: f"Image({self.os_type}/{self.distro}:{self.version}, " f"kind={self.kind}, {len(self._layers)} layers{reg})" ) + + +def cloud_registry_image(image: Image) -> Optional[str]: + """Return the explicit or built-in registry image used by Fleet cloud.""" + if image._registry is not None: + return image._registry + if ( + image.os_type == "linux" + and image.distro == "ubuntu" + and image.version == "24.04" + and image.kind == "vm" + ): + return DEFAULT_LINUX_REGISTRY_IMAGE + return None diff --git a/libs/python/cua-sandbox/cua_sandbox/pool.py b/libs/python/cua-sandbox/cua_sandbox/pool.py index 757f8c5b22..f80835762a 100644 --- a/libs/python/cua-sandbox/cua_sandbox/pool.py +++ b/libs/python/cua-sandbox/cua_sandbox/pool.py @@ -7,7 +7,7 @@ import logging from typing import Any, Callable, Coroutine, Generic, TypeVar, cast -from cua_sandbox.image import Image +from cua_sandbox.image import Image, cloud_registry_image from cua_sandbox.sandbox import Sandbox from cua_sandbox.transport.fleet import FleetTransport from cua_sandbox.transport.fleet_cloud import FleetCloudTransport, _FleetClient @@ -244,7 +244,7 @@ async def apply( if name is None: identity = json.dumps( { - "image": image._registry, + "image": cloud_registry_image(image), "replicas": replicas, "cpu": cpu, "memory_mb": memory_mb, diff --git a/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py b/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py new file mode 100644 index 0000000000..8757deed61 --- /dev/null +++ b/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py @@ -0,0 +1,232 @@ +"""Pull KubeVirt containerDisk images through the OCI registry API.""" + +from __future__ import annotations + +import hashlib +import logging +import os +import platform as _platform +import shutil +import tarfile +import tempfile +import time +from collections.abc import Callable +from pathlib import Path +from typing import Any, Optional + +import oras.provider +import requests +from cua_sandbox.registry.cache import CACHE_ROOT +from cua_sandbox.registry.media_types import VM_MEDIA_TYPES + +logger = logging.getLogger(__name__) + +_CONTAINER_DISK_PATHS = {"disk/disk.img", "./disk/disk.img"} +_LOCK_POLL_INTERVAL_SECONDS = 0.1 + +# Registries disagree about how to authenticate, and oras cannot negotiate on its own: +# public.ecr.aws -> WWW-Authenticate: Bearer ... ; "token" works, "basic" raises +# AttributeError: 'BasicAuth' object has no attribute '_basic_auth' +# private ECR -> WWW-Authenticate: Basic ... ; "basic" works, "token" raises +# ValueError: Cannot respond to request for authentication +# So read the challenge off the registry's /v2/ endpoint and pick the matching backend. +_TOKEN_AUTH_BACKEND = "token" +_BASIC_AUTH_BACKEND = "basic" +_PING_TIMEOUT_SECONDS = 10 + +_INDEX_MEDIA_TYPES = frozenset( + { + "application/vnd.oci.image.index.v1+json", + "application/vnd.docker.distribution.manifest.list.v2+json", + } +) + +_ARCHITECTURE_ALIASES = { + "x86_64": "amd64", + "x86-64": "amd64", + "amd64": "amd64", + "aarch64": "arm64", + "arm64": "arm64", +} + + +def pull_container_disk( + ref: str, + *, + cache_root: Path | None = None, + architecture: Optional[str] = None, + auth_backend: Optional[str] = None, + registry_factory: Callable[..., Any] = oras.provider.Registry, +) -> Path: + """Pull a KubeVirt containerDisk and cache its qcow2 disk locally.""" + root = cache_root or (CACHE_ROOT / "container-disks") + destination = root / hashlib.sha256(ref.encode()).hexdigest() / "disk.qcow2" + if destination.exists(): + return destination + + destination.parent.mkdir(parents=True, exist_ok=True) + lock_path = destination.with_suffix(".lock") + lock_fd = _acquire_cache_lock(lock_path, destination) + if lock_fd is None: + return destination + try: + if destination.exists(): + return destination + + registry = registry_factory(auth_backend=auth_backend or _detect_auth_backend(ref)) + container = registry.get_container(ref) + registry.auth.load_configs(container) + manifest = _resolve_platform_manifest(registry, ref, architecture or _host_architecture()) + + temporary_fd, temporary_name = tempfile.mkstemp( + dir=destination.parent, prefix="disk.", suffix=".tmp" + ) + os.close(temporary_fd) + temporary = Path(temporary_name) + try: + for layer in reversed(manifest.get("layers", [])): + if _is_vm_layer(layer): + continue + logger.info( + "Pulling containerDisk layer %s (%s bytes) from %s", + layer["digest"], + layer.get("size", "?"), + ref, + ) + response = registry.get_blob(container, layer["digest"], stream=True) + try: + with tarfile.open(fileobj=response.raw, mode="r|*") as archive: + for member in archive: + if member.name not in _CONTAINER_DISK_PATHS or not member.isfile(): + continue + source = archive.extractfile(member) + if source is None: + continue + with temporary.open("wb") as output: + shutil.copyfileobj(source, output) + temporary.replace(destination) + logger.info("Cached containerDisk %s at %s", ref, destination) + return destination + finally: + close = getattr(response, "close", None) + if close is not None: + close() + finally: + temporary.unlink(missing_ok=True) + finally: + os.close(lock_fd) + lock_path.unlink(missing_ok=True) + + raise FileNotFoundError(f"OCI image {ref!r} does not contain /disk/disk.img") + + +def _registry_host(ref: str) -> Optional[str]: + """Return the registry host of a reference, or None when it is implicit.""" + head = _repository(ref).split("/", 1)[0] + if head == "localhost" or "." in head or ":" in head: + return head + return None + + +def _detect_auth_backend(ref: str) -> str: + """Pick the oras auth backend from the registry's /v2/ WWW-Authenticate challenge. + + Bearer-token registries (public.ecr.aws, ghcr.io, Docker Hub) need the "token" + backend; registries that answer with Basic (private ECR) need "basic". Defaults to + "token", the OCI-standard flow, when the challenge is missing or unreachable. + """ + host = _registry_host(ref) + if host is None: + return _TOKEN_AUTH_BACKEND + + try: + response = requests.get(f"https://{host}/v2/", timeout=_PING_TIMEOUT_SECONDS) + challenge = response.headers.get("WWW-Authenticate", "") + except requests.RequestException as exc: + logger.debug("Could not probe %s for an auth challenge (%s); assuming bearer", host, exc) + return _TOKEN_AUTH_BACKEND + + backend = ( + _BASIC_AUTH_BACKEND + if challenge.strip().lower().startswith("basic") + else _TOKEN_AUTH_BACKEND + ) + logger.debug("%s answered %r; using the %s auth backend", host, challenge, backend) + return backend + + +def _host_architecture() -> str: + machine = _platform.machine().lower() + return _ARCHITECTURE_ALIASES.get(machine, machine) + + +def _repository(ref: str) -> str: + """Strip the tag or digest from an image reference.""" + repository = ref.split("@", 1)[0] + head, separator, tail = repository.rpartition(":") + if separator and "/" not in tail: + return head + return repository + + +def _is_vm_layer(layer: dict) -> bool: + """True for chunked VM-disk layers (lume/tart/qemu), which are never containerDisks.""" + media_type = layer.get("mediaType", "") + return media_type in VM_MEDIA_TYPES or "part.number=" in media_type + + +def _resolve_platform_manifest(registry: Any, ref: str, architecture: str) -> dict: + """Fetch the manifest for ``ref``, following image indexes to the platform child. + + Multi-arch images publish an OCI image index whose entries are per-platform child + manifests, so it carries ``manifests`` and no ``layers`` of its own. + """ + manifest = registry.get_manifest(ref) + repository = _repository(ref) + seen: set[str] = set() + + while manifest.get("mediaType") in _INDEX_MEDIA_TYPES or manifest.get("manifests"): + entries = manifest.get("manifests", []) + entry = _select_platform_entry(entries, architecture) + if entry is None: + available = ", ".join( + sorted( + f"{(candidate.get('platform') or {}).get('os')}/" + f"{(candidate.get('platform') or {}).get('architecture')}" + for candidate in entries + ) + or [""] + ) + raise FileNotFoundError( + f"OCI image {ref!r} has no linux/{architecture} manifest (available: {available})" + ) + digest = entry["digest"] + if digest in seen: + raise FileNotFoundError(f"OCI image {ref!r} has a cyclic manifest index at {digest}") + seen.add(digest) + manifest = registry.get_manifest(f"{repository}@{digest}") + + return manifest + + +def _select_platform_entry(entries: list, architecture: str) -> Optional[dict]: + for entry in entries: + platform = entry.get("platform") or {} + if platform.get("os") != "linux" or platform.get("architecture") != architecture: + continue + # Buildx publishes provenance/SBOM attestations as extra index entries. + reference_type = (entry.get("annotations") or {}).get("vnd.docker.reference.type", "") + if "attestation" in reference_type: + continue + return entry + return None + + +def _acquire_cache_lock(lock_path: Path, destination: Path) -> int | None: + while True: + try: + return os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_RDWR) + except FileExistsError: + if destination.exists(): + return None + time.sleep(_LOCK_POLL_INTERVAL_SECONDS) diff --git a/libs/python/cua-sandbox/cua_sandbox/sandbox.py b/libs/python/cua-sandbox/cua_sandbox/sandbox.py index 5f7569c2a1..ed21c6be74 100644 --- a/libs/python/cua-sandbox/cua_sandbox/sandbox.py +++ b/libs/python/cua-sandbox/cua_sandbox/sandbox.py @@ -592,9 +592,11 @@ async def create( await _save_fleet_claim_or_close(sandbox, claim_name, pool_name) return sandbox + from cua_sandbox.image import cloud_registry_image + fleet_image = ( image is not None - and image._registry is not None + and cloud_registry_image(image) is not None and cls._uses_fleet(api_key) and not local and runtime is None @@ -739,11 +741,12 @@ async def ephemeral( server_port: int = 8000, telemetry_enabled: bool = True, ) -> AsyncIterator["Sandbox"]: + from cua_sandbox.image import cloud_registry_image from cua_sandbox.pool import Pool fleet_image = ( image is not None - and image._registry is not None + and cloud_registry_image(image) is not None and cls._uses_fleet(api_key) and not local and runtime is None diff --git a/libs/python/cua-sandbox/cua_sandbox/transport/fleet_cloud.py b/libs/python/cua-sandbox/cua_sandbox/transport/fleet_cloud.py index fb73034a95..6e27878a84 100644 --- a/libs/python/cua-sandbox/cua_sandbox/transport/fleet_cloud.py +++ b/libs/python/cua-sandbox/cua_sandbox/transport/fleet_cloud.py @@ -19,7 +19,7 @@ get_fleet_token, get_token_url, ) -from cua_sandbox.image import Image +from cua_sandbox.image import Image, cloud_registry_image from cua_sandbox.transport.cyclops_http_client import CyclopsHttpClient from cua_sandbox.transport.fleet import FleetTransport from fleet_sdk import ( @@ -565,7 +565,7 @@ def _template_request(self) -> CreateTemplateRequest: ] vm_template_builder = ( VmTemplateBuilder() - .container_disk_image(self._image._registry) + .container_disk_image(cloud_registry_image(self._image)) .image_pull_secret("ecr-credentials") .probes( PreservedJson.from_json( @@ -606,8 +606,10 @@ def _service_names(template: Any) -> list[str]: @staticmethod def _validate_image(image: Image) -> None: - if not image._registry: - raise NotImplementedError("Fleet cloud sandboxes require Image.from_registry(...)") + if not cloud_registry_image(image): + raise NotImplementedError( + "Fleet cloud sandboxes require a supported built-in image or Image.from_registry(...)" + ) if ( image._layers or image._env diff --git a/libs/python/cua-sandbox/tests/live/test_fleet_ephemeral.py b/libs/python/cua-sandbox/tests/live/test_fleet_ephemeral.py index 228b01a3f7..aeac86809f 100644 --- a/libs/python/cua-sandbox/tests/live/test_fleet_ephemeral.py +++ b/libs/python/cua-sandbox/tests/live/test_fleet_ephemeral.py @@ -20,8 +20,8 @@ ) IMAGE = ( - "296062593712.dkr.ecr.us-west-2.amazonaws.com/desktop-workspace-duo" - "@sha256:5b9cb82f482834f7541901b87be956e7544d0db13fabc0b372cbc5eca5a74180" + "public.ecr.aws/k5j5w0x5/cua-ubuntu-24.04" + "@sha256:82702ebdd32d1f8fc05f2ea409a7c67d0ba9f8f8e4e9f1a89ce40989d5f4475d" ) diff --git a/libs/python/cua-sandbox/tests/test_cloud.py b/libs/python/cua-sandbox/tests/test_cloud.py index 56b23aacbe..9b458eb153 100644 --- a/libs/python/cua-sandbox/tests/test_cloud.py +++ b/libs/python/cua-sandbox/tests/test_cloud.py @@ -73,16 +73,19 @@ async def disconnect(self): ) fleet_sandbox = await Sandbox.create(Image.from_registry("example:latest"), name="fleet-demo") + default_linux_sandbox = await Sandbox.create(Image.linux(), name="linux-demo") legacy_sandbox = await Sandbox.create( Image.from_registry("example:latest"), name="legacy-demo", api_key="sk-explicit" ) await fleet_sandbox.disconnect() + await default_linux_sandbox.disconnect() await legacy_sandbox.disconnect() assert Sandbox._uses_fleet(None) assert not Sandbox._uses_fleet("sk-explicit") assert [(route, values["name"]) for route, values in routes] == [ ("fleet", "fleet-demo"), + ("fleet", "linux-demo"), ("legacy", "legacy-demo"), ] diff --git a/libs/python/cua-sandbox/tests/test_container_disk.py b/libs/python/cua-sandbox/tests/test_container_disk.py new file mode 100644 index 0000000000..fd2f2f9082 --- /dev/null +++ b/libs/python/cua-sandbox/tests/test_container_disk.py @@ -0,0 +1,522 @@ +import gzip +import hashlib +import io +import tarfile +import threading +from types import SimpleNamespace + +import pytest +import requests +from cua_sandbox.image import DEFAULT_LINUX_REGISTRY_IMAGE, Image +from cua_sandbox.registry import container_disk +from cua_sandbox.registry.container_disk import ( + _LOCK_POLL_INTERVAL_SECONDS, + _detect_auth_backend, + pull_container_disk, +) +from requests.structures import CaseInsensitiveDict + +# The pinned image really is served as a two-child OCI index: the linux/amd64 disk and a +# buildx provenance manifest that reports platform unknown/unknown. Taking "the first" +# or "any" child would grab the attestation and fail confusingly. +AMD64_CHILD = "sha256:85b3f5022bf9ecc864472f5fade5e2f1f54ab88ce429af7873a1415d668dc2ea" +ATTESTATION_CHILD = "sha256:a961974dc086404082c2299cfbdb6bac77dbf930da9846f6a05437f60592e426" + +INDEX_MANIFEST = { + "schemaVersion": 2, + "mediaType": "application/vnd.oci.image.index.v1+json", + "manifests": [ + { + "mediaType": "application/vnd.oci.image.manifest.v1+json", + "digest": AMD64_CHILD, + "size": 483, + "platform": {"architecture": "amd64", "os": "linux"}, + }, + { + "mediaType": "application/vnd.oci.image.manifest.v1+json", + "digest": ATTESTATION_CHILD, + "size": 563, + "annotations": { + "vnd.docker.reference.digest": AMD64_CHILD, + "vnd.docker.reference.type": "attestation-manifest", + }, + "platform": {"architecture": "unknown", "os": "unknown"}, + }, + ], +} + + +def _layer_with_disk(contents: bytes, *, path: str = "disk/disk.img") -> bytes: + raw = io.BytesIO() + with tarfile.open(fileobj=raw, mode="w") as archive: + info = tarfile.TarInfo(path) + info.size = len(contents) + archive.addfile(info, io.BytesIO(contents)) + return gzip.compress(raw.getvalue()) + + +class TestAuthBackendSelection: + """Registries disagree about the challenge scheme; oras cannot negotiate one.""" + + def _challenge(self, monkeypatch, header): + probed = [] + + def get(url, **kwargs): + probed.append(url) + return SimpleNamespace(headers=CaseInsensitiveDict({"WWW-Authenticate": header})) + + monkeypatch.setattr(container_disk.requests, "get", get) + return probed + + def test_a_basic_challenge_selects_basic_auth(self, monkeypatch): + """Private ECR answers with Basic; oras' token backend cannot respond to it.""" + probed = self._challenge( + monkeypatch, 'Basic realm="https://12345.dkr.ecr.us-west-2.amazonaws.com/"' + ) + + backend = _detect_auth_backend("12345.dkr.ecr.us-west-2.amazonaws.com/workspace:tag") + + assert backend == "basic" + assert probed == ["https://12345.dkr.ecr.us-west-2.amazonaws.com/v2/"] + + @pytest.mark.parametrize( + "header", + [ + 'Bearer realm="https://public.ecr.aws/token/",service="public.ecr.aws"', + 'Bearer realm="https://ghcr.io/token",service="ghcr.io"', + "", + ], + ) + def test_a_bearer_or_absent_challenge_selects_token_auth(self, monkeypatch, header): + self._challenge(monkeypatch, header) + + assert _detect_auth_backend("public.ecr.aws/example/workspace:tag") == "token" + + def test_the_pinned_linux_ref_selects_token_auth(self, monkeypatch): + """public.ecr.aws is anonymous-but-bearer, so basic auth has nothing to send.""" + self._challenge(monkeypatch, 'Bearer realm="https://public.ecr.aws/token/"') + + assert _detect_auth_backend(DEFAULT_LINUX_REGISTRY_IMAGE) == "token" + + def test_an_implicit_registry_host_skips_the_probe(self, monkeypatch): + """A short ref names no registry, so there is no host to read a challenge from.""" + + def get(url, **kwargs): + raise AssertionError("a ref without a registry host must not be probed") + + monkeypatch.setattr(container_disk.requests, "get", get) + + assert _detect_auth_backend("ubuntu:24.04") == "token" + + def test_an_unreachable_registry_falls_back_to_token_auth(self, monkeypatch): + def get(url, **kwargs): + raise requests.ConnectionError("no route to host") + + monkeypatch.setattr(container_disk.requests, "get", get) + + assert _detect_auth_backend("registry.example/workspace:tag") == "token" + + +def test_explicit_auth_backend_skips_the_challenge_probe(tmp_path, monkeypatch): + calls = [] + + def boom(url, **kwargs): + raise AssertionError("an explicit auth_backend must not trigger a probe") + + monkeypatch.setattr(container_disk.requests, "get", boom) + + class Registry: + def __init__(self, *, auth_backend): + calls.append(("init", auth_backend)) + self.auth = SimpleNamespace(load_configs=lambda container: None) + + def get_container(self, ref): + return "container" + + def get_manifest(self, ref): + return {"layers": [{"digest": "sha256:layer"}]} + + def get_blob(self, container, digest, *, stream): + return SimpleNamespace(raw=io.BytesIO(_layer_with_disk(b"qcow2"))) + + disk = pull_container_disk( + "registry.example/workspace:latest", + cache_root=tmp_path, + auth_backend="basic", + registry_factory=Registry, + ) + + assert disk.read_bytes() == b"qcow2" + assert calls == [("init", "basic")] + + +def test_pull_container_disk_uses_oras_credentials_and_caches_qcow2(tmp_path): + calls = [] + manifest = { + "layers": [ + { + "digest": "sha256:layer", + "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip", + } + ] + } + + class Registry: + def __init__(self, *, auth_backend): + calls.append(("init", auth_backend)) + self.auth = SimpleNamespace( + load_configs=lambda container: calls.append(("auth", container)) + ) + + def get_container(self, ref): + calls.append(("container", ref)) + return "container" + + def get_manifest(self, ref): + calls.append(("manifest", ref)) + return manifest + + def get_blob(self, container, digest, *, stream): + calls.append(("blob", container, digest, stream)) + return SimpleNamespace(raw=io.BytesIO(_layer_with_disk(b"qcow2"))) + + disk = pull_container_disk( + DEFAULT_LINUX_REGISTRY_IMAGE, + cache_root=tmp_path, + auth_backend="token", + registry_factory=Registry, + ) + + assert disk.name == "disk.qcow2" + assert disk.read_bytes() == b"qcow2" + assert calls == [ + ("init", "token"), + ("container", DEFAULT_LINUX_REGISTRY_IMAGE), + ("auth", "container"), + ("manifest", DEFAULT_LINUX_REGISTRY_IMAGE), + ("blob", "container", "sha256:layer", True), + ] + + calls.clear() + assert ( + pull_container_disk( + DEFAULT_LINUX_REGISTRY_IMAGE, + cache_root=tmp_path, + auth_backend="token", + registry_factory=Registry, + ) + == disk + ) + assert calls == [] + + +def test_pull_probes_for_an_auth_backend_when_none_is_given(tmp_path, monkeypatch): + """Production omits auth_backend, so the challenge probe decides.""" + calls = [] + monkeypatch.setattr(container_disk, "_detect_auth_backend", lambda ref: "basic") + + class Registry: + def __init__(self, *, auth_backend): + calls.append(("init", auth_backend)) + self.auth = SimpleNamespace(load_configs=lambda container: None) + + def get_container(self, ref): + return "container" + + def get_manifest(self, ref): + return {"layers": [{"digest": "sha256:layer"}]} + + def get_blob(self, container, digest, *, stream): + return SimpleNamespace(raw=io.BytesIO(_layer_with_disk(b"qcow2"))) + + pull_container_disk( + "registry.example/workspace:latest", cache_root=tmp_path, registry_factory=Registry + ) + + assert calls == [("init", "basic")] + + +def test_pull_container_disk_resolves_platform_manifest_from_index(tmp_path): + """A multi-arch index has no ``layers`` — the platform child must be followed.""" + requested = [] + child = { + "mediaType": "application/vnd.oci.image.manifest.v1+json", + "layers": [ + { + "digest": "sha256:amd64layer", + "mediaType": "application/vnd.oci.image.layer.v1.tar+gzip", + } + ], + } + + class Registry: + def __init__(self, *, auth_backend): + self.auth = SimpleNamespace(load_configs=lambda container: None) + + def get_container(self, ref): + return "container" + + def get_manifest(self, ref): + requested.append(ref) + return INDEX_MANIFEST if ref == DEFAULT_LINUX_REGISTRY_IMAGE else child + + def get_blob(self, container, digest, *, stream): + assert digest == "sha256:amd64layer" + return SimpleNamespace(raw=io.BytesIO(_layer_with_disk(b"qcow2"))) + + disk = pull_container_disk( + DEFAULT_LINUX_REGISTRY_IMAGE, + cache_root=tmp_path, + architecture="amd64", + auth_backend="token", + registry_factory=Registry, + ) + + assert disk.read_bytes() == b"qcow2" + repository = DEFAULT_LINUX_REGISTRY_IMAGE.rsplit(":", 1)[0] + # The linux/amd64 child, never the unknown/unknown attestation sibling. + assert requested == [DEFAULT_LINUX_REGISTRY_IMAGE, f"{repository}@{AMD64_CHILD}"] + + +def test_index_descent_never_selects_the_attestation_child(tmp_path): + """Even listed first, the buildx provenance manifest must not be mistaken for a disk.""" + requested = [] + attestation_first = { + "mediaType": "application/vnd.oci.image.index.v1+json", + "manifests": [ + INDEX_MANIFEST["manifests"][1], # unknown/unknown attestation + INDEX_MANIFEST["manifests"][0], # linux/amd64 disk + ], + } + + class Registry: + def __init__(self, *, auth_backend): + self.auth = SimpleNamespace(load_configs=lambda container: None) + + def get_container(self, ref): + return "container" + + def get_manifest(self, ref): + requested.append(ref) + if not ref.endswith(AMD64_CHILD): + return attestation_first + return {"layers": [{"digest": "sha256:amd64layer"}]} + + def get_blob(self, container, digest, *, stream): + return SimpleNamespace(raw=io.BytesIO(_layer_with_disk(b"qcow2"))) + + disk = pull_container_disk( + "registry.example/workspace:latest", + cache_root=tmp_path, + architecture="amd64", + auth_backend="token", + registry_factory=Registry, + ) + + assert disk.read_bytes() == b"qcow2" + assert requested == [ + "registry.example/workspace:latest", + f"registry.example/workspace@{AMD64_CHILD}", + ] + assert ATTESTATION_CHILD not in " ".join(requested) + + +def test_pull_container_disk_rejects_index_without_matching_platform(tmp_path): + class Registry: + def __init__(self, *, auth_backend): + self.auth = SimpleNamespace(load_configs=lambda container: None) + + def get_container(self, ref): + return "container" + + def get_manifest(self, ref): + return { + "mediaType": "application/vnd.oci.image.index.v1+json", + "manifests": [ + { + "digest": "sha256:amd64child", + "platform": {"architecture": "amd64", "os": "linux"}, + } + ], + } + + def get_blob(self, container, digest, *, stream): + raise AssertionError("no blob should be fetched without a platform match") + + with pytest.raises(FileNotFoundError, match="no linux/arm64 manifest"): + pull_container_disk( + DEFAULT_LINUX_REGISTRY_IMAGE, + cache_root=tmp_path, + architecture="arm64", + auth_backend="token", + registry_factory=Registry, + ) + + +def test_pull_container_disk_skips_vm_disk_layers(tmp_path): + """Chunked VM images (lume/tart/qemu) are not containerDisks — don't stream them.""" + + class Registry: + def __init__(self, *, auth_backend): + self.auth = SimpleNamespace(load_configs=lambda container: None) + + def get_container(self, ref): + return "container" + + def get_manifest(self, ref): + return { + "layers": [ + { + "digest": "sha256:chunk", + "mediaType": "application/vnd.trycua.lume.disk.chunk.lz4", + } + ] + } + + def get_blob(self, container, digest, *, stream): + raise AssertionError("VM disk chunks must not be downloaded") + + with pytest.raises(FileNotFoundError, match="does not contain /disk/disk.img"): + pull_container_disk( + "registry.example/lume-vm:latest", + cache_root=tmp_path, + auth_backend="token", + registry_factory=Registry, + ) + + +def test_pull_container_disk_searches_all_layers(tmp_path): + empty_layer = _layer_with_disk(b"ignored", path="etc/example") + disk_layer = _layer_with_disk(b"qcow2") + + class Registry: + def __init__(self, *, auth_backend): + self.auth = SimpleNamespace(load_configs=lambda container: None) + + def get_container(self, ref): + return "container" + + def get_manifest(self, ref): + return { + "layers": [ + {"digest": "sha256:empty"}, + {"digest": "sha256:disk"}, + ] + } + + def get_blob(self, container, digest, *, stream): + payload = empty_layer if digest == "sha256:empty" else disk_layer + return SimpleNamespace(raw=io.BytesIO(payload)) + + disk = pull_container_disk( + "registry.example/workspace:latest", + cache_root=tmp_path, + auth_backend="token", + registry_factory=Registry, + ) + + assert disk.read_bytes() == b"qcow2" + + +def test_pull_container_disk_waits_for_existing_lock(tmp_path, monkeypatch): + destination = ( + tmp_path / hashlib.sha256(DEFAULT_LINUX_REGISTRY_IMAGE.encode()).hexdigest() / "disk.qcow2" + ) + lock_path = destination.with_suffix(".lock") + destination.parent.mkdir(parents=True, exist_ok=True) + lock_path.write_text("locked") + destination_bytes = b"cached" + sleeps = [] + + class Registry: + def __init__(self, *, auth_backend): + raise AssertionError("registry should not be used when another process fills the cache") + + def fake_sleep(interval): + sleeps.append(interval) + destination.write_bytes(destination_bytes) + lock_path.unlink() + + monkeypatch.setattr("cua_sandbox.registry.container_disk.time.sleep", fake_sleep) + + disk = pull_container_disk( + DEFAULT_LINUX_REGISTRY_IMAGE, + cache_root=tmp_path, + auth_backend="token", + registry_factory=Registry, + ) + + assert disk == destination + assert disk.read_bytes() == destination_bytes + assert sleeps == [_LOCK_POLL_INTERVAL_SECONDS] + + +async def test_default_linux_session_overlays_pulled_container_disk(tmp_path, monkeypatch): + """Local Image.linux() must boot the same containerDisk Fleet cloud boots.""" + from cua_sandbox.builder import build + from cua_sandbox.registry import container_disk + + pulled = tmp_path / "container-disk.qcow2" + pulled.write_bytes(b"containerdisk") + session_disk = tmp_path / "session.qcow2" + calls = [] + + def fake_pull(ref): + calls.append(("pull", ref, threading.get_ident())) + return pulled + + async def ensure_base_image(os_type, version): + raise AssertionError("built-in Linux must not fall back to a locally built base") + + monkeypatch.setattr(container_disk, "pull_container_disk", fake_pull) + monkeypatch.setattr(build, "ensure_base_image", ensure_base_image) + monkeypatch.setattr(build, "session_overlay_path", lambda name: session_disk) + monkeypatch.setattr( + build, + "create_overlay", + lambda backing, destination: calls.append(("overlay", backing, destination)), + ) + + result = await build.create_session_disk(Image.linux(), "demo") + + assert result == session_disk + assert [call[0] for call in calls] == ["pull", "overlay"] + assert calls[0][1] == DEFAULT_LINUX_REGISTRY_IMAGE + # The pull does network + multi-GB I/O, so it must not run on the event loop thread. + assert calls[0][2] != threading.get_ident() + assert calls[1][1:] == (pulled, session_disk) + + +async def test_non_container_disk_registry_image_falls_back_to_base(tmp_path, monkeypatch): + from cua_sandbox.builder import build + from cua_sandbox.registry import container_disk + + base_disk = tmp_path / "base.qcow2" + base_disk.write_bytes(b"base") + session_disk = tmp_path / "session.qcow2" + calls = [] + + def fake_pull(ref): + raise FileNotFoundError(f"OCI image {ref!r} does not contain /disk/disk.img") + + async def ensure_base_image(os_type, version): + calls.append(("base", os_type, version)) + return base_disk + + monkeypatch.setattr(container_disk, "pull_container_disk", fake_pull) + monkeypatch.setattr(build, "ensure_base_image", ensure_base_image) + monkeypatch.setattr(build, "session_overlay_path", lambda name: session_disk) + monkeypatch.setattr( + build, + "create_overlay", + lambda backing, destination: calls.append(("overlay", backing, destination)), + ) + + image = Image.from_registry("registry.example/lume-vm:latest") + result = await build.create_session_disk(image, "demo") + + assert result == session_disk + assert calls == [ + ("base", "linux", "latest"), + ("overlay", base_disk, session_disk), + ] diff --git a/libs/python/cua-sandbox/tests/test_fleet_cloud_transport.py b/libs/python/cua-sandbox/tests/test_fleet_cloud_transport.py index 915e985960..0b090f00e3 100644 --- a/libs/python/cua-sandbox/tests/test_fleet_cloud_transport.py +++ b/libs/python/cua-sandbox/tests/test_fleet_cloud_transport.py @@ -72,6 +72,14 @@ def test_registry_image_becomes_typed_template_request(): ] +def test_default_linux_image_becomes_typed_template_request(): + request = FleetCloudTransport(image=Image.linux(), name="demo")._template_request() + + assert request.spec.vm_template.container_disk_image == ( + "public.ecr.aws/k5j5w0x5/cua-ubuntu-24.04:main-38352d34" + ) + + def test_pool_request_uses_the_single_sandbox_name_and_requested_replicas(): request = FleetCloudTransport( image=Image.from_registry("registry.example/workspace@sha256:abc"), @@ -133,7 +141,8 @@ async def list_claims(self, namespace): @pytest.mark.parametrize( - "image", [Image.linux(), Image.from_registry("example:latest").apt_install("curl")] + "image", + [Image.linux("debian", "12"), Image.from_registry("example:latest").apt_install("curl")], ) def test_rejects_unsupported_images(image): with pytest.raises(NotImplementedError): diff --git a/libs/python/cua-sandbox/tests/test_image.py b/libs/python/cua-sandbox/tests/test_image.py index 99f1db9954..88b6d678ac 100644 --- a/libs/python/cua-sandbox/tests/test_image.py +++ b/libs/python/cua-sandbox/tests/test_image.py @@ -1,7 +1,11 @@ """Unit tests for the Image builder (no runtime needed).""" import pytest -from cua_sandbox.image import Image +from cua_sandbox.image import ( + DEFAULT_LINUX_REGISTRY_IMAGE, + Image, + cloud_registry_image, +) class TestImageBuilder: @@ -10,6 +14,24 @@ def test_linux_defaults(self): assert img.os_type == "linux" assert img.distro == "ubuntu" assert img.version == "24.04" + assert cloud_registry_image(img) == DEFAULT_LINUX_REGISTRY_IMAGE + + @pytest.mark.parametrize( + "image", + [ + Image.linux("debian", "12"), + Image.linux("ubuntu", "22.04"), + Image.linux("ubuntu", "24.04", kind="container"), + Image.windows(), + ], + ) + def test_non_default_images_have_no_cloud_registry_default(self, image): + assert cloud_registry_image(image) is None + + def test_explicit_registry_is_cloud_override(self): + image = Image.linux()._with(_registry="registry.example/custom:latest") + + assert cloud_registry_image(image) == "registry.example/custom:latest" def test_macos(self): img = Image.macos("15")