From 5d68d13fadc4aa28c193fed9126aadf47deec069 Mon Sep 17 00:00:00 2001 From: r33drichards <57335981+r33drichards@users.noreply.github.com> Date: Wed, 12 Aug 2026 06:00:27 +0000 Subject: [PATCH 1/8] feat(sandbox): default Linux to desktop containerDisk --- .../cua-sandbox/cua_sandbox/builder/build.py | 11 +- libs/python/cua-sandbox/cua_sandbox/image.py | 24 ++++ libs/python/cua-sandbox/cua_sandbox/pool.py | 4 +- .../cua_sandbox/registry/container_disk.py | 60 ++++++++ .../python/cua-sandbox/cua_sandbox/sandbox.py | 7 +- .../cua_sandbox/transport/fleet_cloud.py | 10 +- libs/python/cua-sandbox/tests/test_cloud.py | 3 + .../cua-sandbox/tests/test_container_disk.py | 129 ++++++++++++++++++ .../tests/test_fleet_cloud_transport.py | 12 +- libs/python/cua-sandbox/tests/test_image.py | 26 +++- 10 files changed, 274 insertions(+), 12 deletions(-) create mode 100644 libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py create mode 100644 libs/python/cua-sandbox/tests/test_container_disk.py diff --git a/libs/python/cua-sandbox/cua_sandbox/builder/build.py b/libs/python/cua-sandbox/cua_sandbox/builder/build.py index 932eab43ad..6211292840 100644 --- a/libs/python/cua-sandbox/cua_sandbox/builder/build.py +++ b/libs/python/cua-sandbox/cua_sandbox/builder/build.py @@ -341,8 +341,15 @@ 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) + from cua_sandbox.image import default_registry_image + + registry_image = default_registry_image(image) + if registry_image: + from cua_sandbox.registry.container_disk import pull_container_disk + + backing = await asyncio.to_thread(pull_container_disk, registry_image) + else: + backing = await ensure_base_image(image.os_type, image.version) # 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..306f34047e 100644 --- a/libs/python/cua-sandbox/cua_sandbox/image.py +++ b/libs/python/cua-sandbox/cua_sandbox/image.py @@ -28,6 +28,11 @@ logger = logging.getLogger(__name__) +DEFAULT_LINUX_REGISTRY_IMAGE = ( + "296062593712.dkr.ecr.us-west-2.amazonaws.com/" + "desktop-workspace-duo:main-38352d34" +) + _IMAGE_CACHE = Path.home() / ".cua" / "cua-sandbox" / "image-cache" @@ -477,3 +482,22 @@ def __repr__(self) -> str: f"Image({self.os_type}/{self.distro}:{self.version}, " f"kind={self.kind}, {len(self._layers)} layers{reg})" ) + + +def default_registry_image(image: Image) -> Optional[str]: + """Return the registry image backing a built-in image descriptor.""" + if image._registry is not None: + return None + 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 + + +def cloud_registry_image(image: Image) -> Optional[str]: + """Return the explicit or built-in registry image used by Fleet cloud.""" + return image._registry or default_registry_image(image) 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..00412e92f5 --- /dev/null +++ b/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py @@ -0,0 +1,60 @@ +"""Pull KubeVirt containerDisk images through the OCI registry API.""" + +from __future__ import annotations + +import hashlib +import shutil +import tarfile +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import oras.provider + +from cua_sandbox.registry.cache import CACHE_ROOT + +_CONTAINER_DISK_PATHS = {"disk/disk.img", "./disk/disk.img"} + + +def pull_container_disk( + ref: str, + *, + cache_root: Path | None = 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) + temporary = destination.with_suffix(".tmp") + registry = registry_factory(auth_backend="token") + container = registry.get_container(ref) + registry.auth.load_configs(container) + manifest = registry.get_manifest(ref) + + try: + for layer in reversed(manifest.get("layers", [])): + 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) + return destination + finally: + close = getattr(response, "close", None) + if close is not None: + close() + finally: + temporary.unlink(missing_ok=True) + + raise FileNotFoundError(f"OCI image {ref!r} does not contain /disk/disk.img") 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/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..c5cb6fc040 --- /dev/null +++ b/libs/python/cua-sandbox/tests/test_container_disk.py @@ -0,0 +1,129 @@ +import gzip +import io +import tarfile +from types import SimpleNamespace + +from cua_sandbox.image import DEFAULT_LINUX_REGISTRY_IMAGE +from cua_sandbox.registry.container_disk import pull_container_disk + + +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()) + + +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, + 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, + registry_factory=Registry, + ) + == disk + ) + assert calls == [] + + +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, registry_factory=Registry) + + assert disk.read_bytes() == b"qcow2" + + +async def test_default_linux_session_uses_container_disk_as_qemu_backing(tmp_path, monkeypatch): + from cua_sandbox.builder import build + from cua_sandbox.image import Image + + base_disk = tmp_path / "base.qcow2" + base_disk.write_bytes(b"base") + session_disk = tmp_path / "session.qcow2" + calls = [] + + monkeypatch.setattr( + "cua_sandbox.registry.container_disk.pull_container_disk", + lambda ref: calls.append(("pull", ref)) or base_disk, + ) + 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 calls == [ + ("pull", DEFAULT_LINUX_REGISTRY_IMAGE), + ("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..7207a88f2e 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,15 @@ 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 == ( + "296062593712.dkr.ecr.us-west-2.amazonaws.com/" + "desktop-workspace-duo: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 +142,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..eefc58c4a8 100644 --- a/libs/python/cua-sandbox/tests/test_image.py +++ b/libs/python/cua-sandbox/tests/test_image.py @@ -1,7 +1,12 @@ """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, + default_registry_image, +) class TestImageBuilder: @@ -10,6 +15,25 @@ def test_linux_defaults(self): assert img.os_type == "linux" assert img.distro == "ubuntu" assert img.version == "24.04" + assert default_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_registry_default(self, image): + assert default_registry_image(image) is None + + def test_explicit_registry_is_cloud_override_not_local_default(self): + image = Image.linux()._with(_registry="registry.example/custom:latest") + + assert default_registry_image(image) is None + assert cloud_registry_image(image) == "registry.example/custom:latest" def test_macos(self): img = Image.macos("15") From f1056ddc184f4d3a5ee4ded49e832b0ee5c4c82d Mon Sep 17 00:00:00 2001 From: r33drichards <57335981+r33drichards@users.noreply.github.com> Date: Wed, 12 Aug 2026 06:11:27 +0000 Subject: [PATCH 2/8] style(sandbox): apply Python formatters --- libs/python/cua-sandbox/cua_sandbox/image.py | 3 +-- .../cua-sandbox/cua_sandbox/registry/container_disk.py | 1 - libs/python/cua-sandbox/tests/test_container_disk.py | 8 ++++++-- .../cua-sandbox/tests/test_fleet_cloud_transport.py | 5 ++--- 4 files changed, 9 insertions(+), 8 deletions(-) diff --git a/libs/python/cua-sandbox/cua_sandbox/image.py b/libs/python/cua-sandbox/cua_sandbox/image.py index 306f34047e..79d23ec14b 100644 --- a/libs/python/cua-sandbox/cua_sandbox/image.py +++ b/libs/python/cua-sandbox/cua_sandbox/image.py @@ -29,8 +29,7 @@ logger = logging.getLogger(__name__) DEFAULT_LINUX_REGISTRY_IMAGE = ( - "296062593712.dkr.ecr.us-west-2.amazonaws.com/" - "desktop-workspace-duo:main-38352d34" + "296062593712.dkr.ecr.us-west-2.amazonaws.com/" "desktop-workspace-duo:main-38352d34" ) _IMAGE_CACHE = Path.home() / ".cua" / "cua-sandbox" / "image-cache" diff --git a/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py b/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py index 00412e92f5..bdb7f69771 100644 --- a/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py +++ b/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py @@ -10,7 +10,6 @@ from typing import Any import oras.provider - from cua_sandbox.registry.cache import CACHE_ROOT _CONTAINER_DISK_PATHS = {"disk/disk.img", "./disk/disk.img"} diff --git a/libs/python/cua-sandbox/tests/test_container_disk.py b/libs/python/cua-sandbox/tests/test_container_disk.py index c5cb6fc040..414072f75e 100644 --- a/libs/python/cua-sandbox/tests/test_container_disk.py +++ b/libs/python/cua-sandbox/tests/test_container_disk.py @@ -30,7 +30,9 @@ def test_pull_container_disk_uses_oras_credentials_and_caches_qcow2(tmp_path): class Registry: def __init__(self, *, auth_backend): calls.append(("init", auth_backend)) - self.auth = SimpleNamespace(load_configs=lambda container: calls.append(("auth", container))) + self.auth = SimpleNamespace( + load_configs=lambda container: calls.append(("auth", container)) + ) def get_container(self, ref): calls.append(("container", ref)) @@ -95,7 +97,9 @@ 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, registry_factory=Registry) + disk = pull_container_disk( + "registry.example/workspace:latest", cache_root=tmp_path, registry_factory=Registry + ) assert disk.read_bytes() == b"qcow2" 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 7207a88f2e..9c537429c1 100644 --- a/libs/python/cua-sandbox/tests/test_fleet_cloud_transport.py +++ b/libs/python/cua-sandbox/tests/test_fleet_cloud_transport.py @@ -76,8 +76,7 @@ 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 == ( - "296062593712.dkr.ecr.us-west-2.amazonaws.com/" - "desktop-workspace-duo:main-38352d34" + "296062593712.dkr.ecr.us-west-2.amazonaws.com/" "desktop-workspace-duo:main-38352d34" ) @@ -143,7 +142,7 @@ async def list_claims(self, namespace): @pytest.mark.parametrize( "image", - [Image.linux("debian", "12"), Image.from_registry("example:latest").apt_install("curl")] + [Image.linux("debian", "12"), Image.from_registry("example:latest").apt_install("curl")], ) def test_rejects_unsupported_images(image): with pytest.raises(NotImplementedError): From e3ef13924ff4f3ee098f32f3ee0b7cf93ac6dd5f Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 12 Aug 2026 17:12:55 +0000 Subject: [PATCH 3/8] fix(sandbox): address review feedback on linux containerDisk defaults Co-authored-by: r33drichards <57335981+r33drichards@users.noreply.github.com> --- .../cua-sandbox/cua_sandbox/builder/build.py | 10 +-- libs/python/cua-sandbox/cua_sandbox/image.py | 11 +-- .../cua_sandbox/registry/container_disk.py | 79 +++++++++++++------ .../cua-sandbox/tests/test_container_disk.py | 53 +++++++++++-- libs/python/cua-sandbox/tests/test_image.py | 10 +-- 5 files changed, 108 insertions(+), 55 deletions(-) diff --git a/libs/python/cua-sandbox/cua_sandbox/builder/build.py b/libs/python/cua-sandbox/cua_sandbox/builder/build.py index 6211292840..bf3adc629f 100644 --- a/libs/python/cua-sandbox/cua_sandbox/builder/build.py +++ b/libs/python/cua-sandbox/cua_sandbox/builder/build.py @@ -341,15 +341,7 @@ async def create_session_disk( elif image._disk_path: backing = Path(image._disk_path) else: - from cua_sandbox.image import default_registry_image - - registry_image = default_registry_image(image) - if registry_image: - from cua_sandbox.registry.container_disk import pull_container_disk - - backing = await asyncio.to_thread(pull_container_disk, registry_image) - else: - backing = await ensure_base_image(image.os_type, image.version) + backing = await ensure_base_image(image.os_type, image.version) # 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 79d23ec14b..a5e788f325 100644 --- a/libs/python/cua-sandbox/cua_sandbox/image.py +++ b/libs/python/cua-sandbox/cua_sandbox/image.py @@ -483,10 +483,10 @@ def __repr__(self) -> str: ) -def default_registry_image(image: Image) -> Optional[str]: - """Return the registry image backing a built-in image descriptor.""" +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 None + return image._registry if ( image.os_type == "linux" and image.distro == "ubuntu" @@ -495,8 +495,3 @@ def default_registry_image(image: Image) -> Optional[str]: ): return DEFAULT_LINUX_REGISTRY_IMAGE return None - - -def cloud_registry_image(image: Image) -> Optional[str]: - """Return the explicit or built-in registry image used by Fleet cloud.""" - return image._registry or default_registry_image(image) diff --git a/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py b/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py index bdb7f69771..6ed271faa8 100644 --- a/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py +++ b/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py @@ -3,8 +3,11 @@ from __future__ import annotations import hashlib +import os import shutil import tarfile +import tempfile +import time from collections.abc import Callable from pathlib import Path from typing import Any @@ -13,6 +16,7 @@ from cua_sandbox.registry.cache import CACHE_ROOT _CONTAINER_DISK_PATHS = {"disk/disk.img", "./disk/disk.img"} +_LOCK_POLL_INTERVAL_SECONDS = 0.1 def pull_container_disk( @@ -28,32 +32,57 @@ def pull_container_disk( return destination destination.parent.mkdir(parents=True, exist_ok=True) - temporary = destination.with_suffix(".tmp") - registry = registry_factory(auth_backend="token") - container = registry.get_container(ref) - registry.auth.load_configs(container) - manifest = registry.get_manifest(ref) - + lock_path = destination.with_suffix(".lock") + lock_fd = _acquire_cache_lock(lock_path, destination) + if lock_fd is None: + return destination try: - for layer in reversed(manifest.get("layers", [])): - 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) - return destination - finally: - close = getattr(response, "close", None) - if close is not None: - close() + if destination.exists(): + return destination + + registry = registry_factory(auth_backend="token") + container = registry.get_container(ref) + registry.auth.load_configs(container) + manifest = registry.get_manifest(ref) + + 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", [])): + 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) + return destination + finally: + close = getattr(response, "close", None) + if close is not None: + close() + finally: + temporary.unlink(missing_ok=True) finally: - temporary.unlink(missing_ok=True) + os.close(lock_fd) + lock_path.unlink(missing_ok=True) raise FileNotFoundError(f"OCI image {ref!r} does not contain /disk/disk.img") + + +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/tests/test_container_disk.py b/libs/python/cua-sandbox/tests/test_container_disk.py index 414072f75e..b3907c4080 100644 --- a/libs/python/cua-sandbox/tests/test_container_disk.py +++ b/libs/python/cua-sandbox/tests/test_container_disk.py @@ -1,10 +1,11 @@ import gzip +import hashlib import io import tarfile from types import SimpleNamespace -from cua_sandbox.image import DEFAULT_LINUX_REGISTRY_IMAGE -from cua_sandbox.registry.container_disk import pull_container_disk +from cua_sandbox.image import DEFAULT_LINUX_REGISTRY_IMAGE, Image +from cua_sandbox.registry.container_disk import _LOCK_POLL_INTERVAL_SECONDS, pull_container_disk def _layer_with_disk(contents: bytes, *, path: str = "disk/disk.img") -> bytes: @@ -104,18 +105,56 @@ def get_blob(self, container, digest, *, stream): assert disk.read_bytes() == b"qcow2" -async def test_default_linux_session_uses_container_disk_as_qemu_backing(tmp_path, monkeypatch): +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, + registry_factory=Registry, + ) + + assert disk == destination + assert disk.read_bytes() == destination_bytes + assert sleeps == [_LOCK_POLL_INTERVAL_SECONDS] + + +async def test_default_linux_session_uses_standard_base_image_backing(tmp_path, monkeypatch): from cua_sandbox.builder import build - from cua_sandbox.image import Image base_disk = tmp_path / "base.qcow2" base_disk.write_bytes(b"base") session_disk = tmp_path / "session.qcow2" calls = [] + async def ensure_base_image(os_type, version): + calls.append(("base", os_type, version)) + return base_disk + monkeypatch.setattr( - "cua_sandbox.registry.container_disk.pull_container_disk", - lambda ref: calls.append(("pull", ref)) or base_disk, + build, + "ensure_base_image", + ensure_base_image, ) monkeypatch.setattr(build, "session_overlay_path", lambda name: session_disk) monkeypatch.setattr( @@ -128,6 +167,6 @@ async def test_default_linux_session_uses_container_disk_as_qemu_backing(tmp_pat assert result == session_disk assert calls == [ - ("pull", DEFAULT_LINUX_REGISTRY_IMAGE), + ("base", "linux", "24.04"), ("overlay", base_disk, session_disk), ] diff --git a/libs/python/cua-sandbox/tests/test_image.py b/libs/python/cua-sandbox/tests/test_image.py index eefc58c4a8..88b6d678ac 100644 --- a/libs/python/cua-sandbox/tests/test_image.py +++ b/libs/python/cua-sandbox/tests/test_image.py @@ -5,7 +5,6 @@ DEFAULT_LINUX_REGISTRY_IMAGE, Image, cloud_registry_image, - default_registry_image, ) @@ -15,7 +14,7 @@ def test_linux_defaults(self): assert img.os_type == "linux" assert img.distro == "ubuntu" assert img.version == "24.04" - assert default_registry_image(img) == DEFAULT_LINUX_REGISTRY_IMAGE + assert cloud_registry_image(img) == DEFAULT_LINUX_REGISTRY_IMAGE @pytest.mark.parametrize( "image", @@ -26,13 +25,12 @@ def test_linux_defaults(self): Image.windows(), ], ) - def test_non_default_images_have_no_registry_default(self, image): - assert default_registry_image(image) is None + 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_not_local_default(self): + def test_explicit_registry_is_cloud_override(self): image = Image.linux()._with(_registry="registry.example/custom:latest") - assert default_registry_image(image) is None assert cloud_registry_image(image) == "registry.example/custom:latest" def test_macos(self): From c8ad94f3f8397f391e1c6cc0a60864c64f32f851 Mon Sep 17 00:00:00 2001 From: Robert Wendt Date: Wed, 12 Aug 2026 18:00:01 +0000 Subject: [PATCH 4/8] fix(sandbox): boot the cloud containerDisk for local Linux sessions pull_container_disk() was never reachable and did not work against the real registry, so Image.linux() still fell through to a locally built base image instead of the disk Fleet cloud boots. - Use oras' basic auth backend. ECR answers WWW-Authenticate: Basic, so the token backend failed with "This endpoint requires a token. Please use basic auth with a username or password." after ~2 minutes of retry backoff. - Follow OCI image indexes to the platform child manifest. The real image is a multi-arch index, which carries "manifests" and no "layers", so the extraction loop never ran and the pull always raised FileNotFoundError. Buildx attestation entries are skipped. - Skip chunked VM-disk layers (lume/tart/qemu) instead of streaming GBs of a non-containerDisk image before failing. - Wire the pull into create_session_disk() via resolve_backing_disk(), off the event loop with asyncio.to_thread. The session overlay is now backed by the pulled containerDisk; images with no registry counterpart still fall back to ensure_base_image(). Verified end to end on bare metal with KVM: Image.linux() pulls the pinned ECR containerDisk, overlays it, boots it under qemu-system-x86_64 -enable-kvm, and serves shell.run() and screenshot() from the guest computer-server. Co-Authored-By: Claude Opus 5 (1M context) --- .../cua-sandbox/cua_sandbox/builder/build.py | 32 ++- .../cua_sandbox/registry/container_disk.py | 109 +++++++++- .../cua-sandbox/tests/test_container_disk.py | 191 ++++++++++++++++-- 3 files changed, 314 insertions(+), 18 deletions(-) diff --git a/libs/python/cua-sandbox/cua_sandbox/builder/build.py b/libs/python/cua-sandbox/cua_sandbox/builder/build.py index bf3adc629f..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,7 +367,7 @@ async def create_session_disk( elif image._disk_path: backing = Path(image._disk_path) else: - 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/registry/container_disk.py b/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py index 6ed271faa8..775e7678e1 100644 --- a/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py +++ b/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py @@ -3,26 +3,53 @@ 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 +from typing import Any, Optional import oras.provider 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 +# ECR fronts the registry API with HTTP Basic auth, so oras' "token" backend fails with +# "This endpoint requires a token. Please use basic auth with a username or password." +# The basic backend reads the same ~/.docker/config.json entry and works for token-based +# registries too, because it only replays the credential the daemon config already holds. +_AUTH_BACKEND = "basic" + +_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, registry_factory: Callable[..., Any] = oras.provider.Registry, ) -> Path: """Pull a KubeVirt containerDisk and cache its qcow2 disk locally.""" @@ -40,10 +67,10 @@ def pull_container_disk( if destination.exists(): return destination - registry = registry_factory(auth_backend="token") + registry = registry_factory(auth_backend=_AUTH_BACKEND) container = registry.get_container(ref) registry.auth.load_configs(container) - manifest = registry.get_manifest(ref) + manifest = _resolve_platform_manifest(registry, ref, architecture or _host_architecture()) temporary_fd, temporary_name = tempfile.mkstemp( dir=destination.parent, prefix="disk.", suffix=".tmp" @@ -52,6 +79,14 @@ def pull_container_disk( 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: @@ -64,6 +99,7 @@ def pull_container_disk( 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) @@ -78,6 +114,73 @@ def pull_container_disk( raise FileNotFoundError(f"OCI image {ref!r} does not contain /disk/disk.img") +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: diff --git a/libs/python/cua-sandbox/tests/test_container_disk.py b/libs/python/cua-sandbox/tests/test_container_disk.py index b3907c4080..2698f91f32 100644 --- a/libs/python/cua-sandbox/tests/test_container_disk.py +++ b/libs/python/cua-sandbox/tests/test_container_disk.py @@ -2,11 +2,38 @@ import hashlib import io import tarfile +import threading from types import SimpleNamespace +import pytest from cua_sandbox.image import DEFAULT_LINUX_REGISTRY_IMAGE, Image from cua_sandbox.registry.container_disk import _LOCK_POLL_INTERVAL_SECONDS, pull_container_disk +# The real ECR image is a multi-arch index whose children are the per-platform manifests +# plus a buildx attestation manifest. +INDEX_MANIFEST = { + "schemaVersion": 2, + "mediaType": "application/vnd.oci.image.index.v1+json", + "manifests": [ + { + "mediaType": "application/vnd.oci.image.manifest.v1+json", + "digest": "sha256:arm64child", + "platform": {"architecture": "arm64", "os": "linux"}, + }, + { + "mediaType": "application/vnd.oci.image.manifest.v1+json", + "digest": "sha256:amd64child", + "platform": {"architecture": "amd64", "os": "linux"}, + }, + { + "mediaType": "application/vnd.oci.image.manifest.v1+json", + "digest": "sha256:attestation", + "annotations": {"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() @@ -55,8 +82,10 @@ def get_blob(self, container, digest, *, stream): assert disk.name == "disk.qcow2" assert disk.read_bytes() == b"qcow2" + # ECR rejects oras' token backend ("This endpoint requires a token. Please use basic + # auth with a username or password."), so the pull must use basic auth. assert calls == [ - ("init", "token"), + ("init", "basic"), ("container", DEFAULT_LINUX_REGISTRY_IMAGE), ("auth", "container"), ("manifest", DEFAULT_LINUX_REGISTRY_IMAGE), @@ -75,6 +104,108 @@ def get_blob(self, container, digest, *, stream): assert calls == [] +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", + registry_factory=Registry, + ) + + assert disk.read_bytes() == b"qcow2" + repository = DEFAULT_LINUX_REGISTRY_IMAGE.rsplit(":", 1)[0] + assert requested == [DEFAULT_LINUX_REGISTRY_IMAGE, f"{repository}@sha256:amd64child"] + + +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", + 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, + 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") @@ -107,9 +238,7 @@ def get_blob(self, container, digest, *, stream): 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" + 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) @@ -139,23 +268,60 @@ def fake_sleep(interval): assert sleeps == [_LOCK_POLL_INTERVAL_SECONDS] -async def test_default_linux_session_uses_standard_base_image_backing(tmp_path, monkeypatch): +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( - build, - "ensure_base_image", - ensure_base_image, - ) + 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, @@ -163,10 +329,11 @@ async def ensure_base_image(os_type, version): lambda backing, destination: calls.append(("overlay", backing, destination)), ) - result = await build.create_session_disk(Image.linux(), "demo") + 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", "24.04"), + ("base", "linux", "latest"), ("overlay", base_disk, session_disk), ] From f80691454ca767a3907d72ed1e9ed743d0aa2a9e Mon Sep 17 00:00:00 2001 From: r33drichards <57335981+r33drichards@users.noreply.github.com> Date: Wed, 12 Aug 2026 19:19:39 +0000 Subject: [PATCH 5/8] fix(sandbox): use public Ubuntu containerDisk --- libs/python/cua-sandbox/cua_sandbox/image.py | 4 +--- libs/python/cua-sandbox/tests/live/test_fleet_ephemeral.py | 4 ++-- libs/python/cua-sandbox/tests/test_fleet_cloud_transport.py | 2 +- 3 files changed, 4 insertions(+), 6 deletions(-) diff --git a/libs/python/cua-sandbox/cua_sandbox/image.py b/libs/python/cua-sandbox/cua_sandbox/image.py index a5e788f325..30c44ff9f6 100644 --- a/libs/python/cua-sandbox/cua_sandbox/image.py +++ b/libs/python/cua-sandbox/cua_sandbox/image.py @@ -28,9 +28,7 @@ logger = logging.getLogger(__name__) -DEFAULT_LINUX_REGISTRY_IMAGE = ( - "296062593712.dkr.ecr.us-west-2.amazonaws.com/" "desktop-workspace-duo:main-38352d34" -) +DEFAULT_LINUX_REGISTRY_IMAGE = "public.ecr.aws/k5j5w0x5/" "cua-ubuntu-24.04:main-e5d853a9" _IMAGE_CACHE = Path.home() / ".cua" / "cua-sandbox" / "image-cache" 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_fleet_cloud_transport.py b/libs/python/cua-sandbox/tests/test_fleet_cloud_transport.py index 9c537429c1..36943216d0 100644 --- a/libs/python/cua-sandbox/tests/test_fleet_cloud_transport.py +++ b/libs/python/cua-sandbox/tests/test_fleet_cloud_transport.py @@ -76,7 +76,7 @@ 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 == ( - "296062593712.dkr.ecr.us-west-2.amazonaws.com/" "desktop-workspace-duo:main-38352d34" + "public.ecr.aws/k5j5w0x5/" "cua-ubuntu-24.04:main-e5d853a9" ) From 24c3480a9884d505d350ac3935ef950886c4abd3 Mon Sep 17 00:00:00 2001 From: Robert Wendt Date: Wed, 12 Aug 2026 20:02:40 +0000 Subject: [PATCH 6/8] fix(sandbox): pick the oras auth backend from the registry challenge Repinning the Linux default to public.ecr.aws invalidated the hardcoded basic auth backend. oras fixes its backend at construction and cannot negotiate one, but the two registries want opposite schemes: 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 the backend is now read off the registry's /v2/ challenge instead of being hardcoded, which keeps both the public default and explicit private Image.from_registry(...) refs working. Callers can force one with the new auth_backend= keyword; the probe falls back to the OCI-standard bearer flow when the registry is unreachable or sends no challenge. Verified against both real registries: detection picks token for public.ecr.aws and basic for private ECR, and a cold-cache pull succeeds on each. The public image is also an OCI index, so the index-descent path is still exercised. Co-Authored-By: Claude Opus 5 (1M context) --- .../cua_sandbox/registry/container_disk.py | 53 ++++++++- .../cua-sandbox/tests/test_container_disk.py | 109 +++++++++++++++++- 2 files changed, 151 insertions(+), 11 deletions(-) diff --git a/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py b/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py index 775e7678e1..8757deed61 100644 --- a/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py +++ b/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py @@ -15,6 +15,7 @@ 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 @@ -23,11 +24,15 @@ _CONTAINER_DISK_PATHS = {"disk/disk.img", "./disk/disk.img"} _LOCK_POLL_INTERVAL_SECONDS = 0.1 -# ECR fronts the registry API with HTTP Basic auth, so oras' "token" backend fails with -# "This endpoint requires a token. Please use basic auth with a username or password." -# The basic backend reads the same ~/.docker/config.json entry and works for token-based -# registries too, because it only replays the credential the daemon config already holds. -_AUTH_BACKEND = "basic" +# 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( { @@ -50,6 +55,7 @@ def pull_container_disk( *, 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.""" @@ -67,7 +73,7 @@ def pull_container_disk( if destination.exists(): return destination - registry = registry_factory(auth_backend=_AUTH_BACKEND) + 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()) @@ -114,6 +120,41 @@ def pull_container_disk( 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) diff --git a/libs/python/cua-sandbox/tests/test_container_disk.py b/libs/python/cua-sandbox/tests/test_container_disk.py index 2698f91f32..baf520ff7f 100644 --- a/libs/python/cua-sandbox/tests/test_container_disk.py +++ b/libs/python/cua-sandbox/tests/test_container_disk.py @@ -6,8 +6,35 @@ from types import SimpleNamespace import pytest +import requests from cua_sandbox.image import DEFAULT_LINUX_REGISTRY_IMAGE, Image -from cua_sandbox.registry.container_disk import _LOCK_POLL_INTERVAL_SECONDS, pull_container_disk +from cua_sandbox.registry import container_disk +from cua_sandbox.registry.container_disk import ( + _LOCK_POLL_INTERVAL_SECONDS, + _detect_auth_backend, + pull_container_disk, +) + +BEARER_CHALLENGE = 'Bearer realm="https://public.ecr.aws/token/",service="public.ecr.aws"' +BASIC_CHALLENGE = 'Basic realm="https://296062593712.dkr.ecr.us-west-2.amazonaws.com/"' + + +def _fake_ping(monkeypatch, challenge, *, probes=None): + """Stub the /v2/ auth-challenge probe.""" + + def fake_get(url, timeout=None): + if probes is not None: + probes.append(url) + return SimpleNamespace(headers={"WWW-Authenticate": challenge} if challenge else {}) + + monkeypatch.setattr(container_disk.requests, "get", fake_get) + + +@pytest.fixture(autouse=True) +def _never_probe_a_real_registry(monkeypatch): + """Keep the auth-challenge probe off the network unless a test opts in.""" + _fake_ping(monkeypatch, BEARER_CHALLENGE) + # The real ECR image is a multi-arch index whose children are the per-platform manifests # plus a buildx attestation manifest. @@ -44,7 +71,79 @@ def _layer_with_disk(contents: bytes, *, path: str = "disk/disk.img") -> bytes: return gzip.compress(raw.getvalue()) -def test_pull_container_disk_uses_oras_credentials_and_caches_qcow2(tmp_path): +@pytest.mark.parametrize( + ("ref", "challenge", "expected"), + [ + # public.ecr.aws is anonymously readable but still uses the Bearer token flow. + ("public.ecr.aws/k5j5w0x5/cua-ubuntu-24.04:main-e5d853a9", BEARER_CHALLENGE, "token"), + # Private ECR answers with Basic and rejects oras' token backend. + ("296062593712.dkr.ecr.us-west-2.amazonaws.com/duo:main", BASIC_CHALLENGE, "basic"), + ("ghcr.io/trycua/workspace:latest", BEARER_CHALLENGE, "token"), + # No challenge at all (open registry) — assume the OCI-standard bearer flow. + ("registry.example/workspace:latest", "", "token"), + ], +) +def test_auth_backend_follows_the_registry_challenge(monkeypatch, ref, challenge, expected): + probes = [] + _fake_ping(monkeypatch, challenge, probes=probes) + + assert _detect_auth_backend(ref) == expected + assert probes == [f"https://{ref.split('/', 1)[0]}/v2/"] + + +def test_auth_backend_defaults_to_token_when_the_probe_fails(monkeypatch): + def boom(url, timeout=None): + raise requests.ConnectionError("no route to host") + + monkeypatch.setattr(container_disk.requests, "get", boom) + + assert _detect_auth_backend("registry.example/workspace:latest") == "token" + + +def test_auth_backend_skips_the_probe_for_a_hostless_ref(monkeypatch): + def boom(url, timeout=None): + raise AssertionError("a ref without a registry host must not be probed") + + monkeypatch.setattr(container_disk.requests, "get", boom) + + assert _detect_auth_backend("workspace:latest") == "token" + + +def test_explicit_auth_backend_overrides_detection(tmp_path, monkeypatch): + calls = [] + + def boom(url, timeout=None): + 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, monkeypatch): + _fake_ping(monkeypatch, BEARER_CHALLENGE) calls = [] manifest = { "layers": [ @@ -82,10 +181,10 @@ def get_blob(self, container, digest, *, stream): assert disk.name == "disk.qcow2" assert disk.read_bytes() == b"qcow2" - # ECR rejects oras' token backend ("This endpoint requires a token. Please use basic - # auth with a username or password."), so the pull must use basic auth. + # The backend comes from the registry's challenge — hardcoding either one breaks + # half the registries (see test_auth_backend_follows_the_registry_challenge). assert calls == [ - ("init", "basic"), + ("init", "token"), ("container", DEFAULT_LINUX_REGISTRY_IMAGE), ("auth", "container"), ("manifest", DEFAULT_LINUX_REGISTRY_IMAGE), From f84758c44313b0c8332c9eb1c6ed60b8cc425a50 Mon Sep 17 00:00:00 2001 From: Robert Wendt Date: Wed, 12 Aug 2026 20:13:02 +0000 Subject: [PATCH 7/8] fix(sandbox): negotiate the registry auth backend, repin to main-38352d34 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adopts the Windows branch's _auth_backend_for() verbatim so registry/container_disk.py stays byte-identical between the two branches. oras fixes its auth backend at construction and cannot negotiate, but the registries want opposite schemes, measured live: basic public.ecr.aws FAIL AttributeError: 'BasicAuth' object has no attribute '_basic_auth' token public.ecr.aws OK basic private ECR OK token private ECR FAIL ValueError: Cannot respond to request for authentication So the scheme is read off the registry's 401 challenge instead of guessed. pull_container_disk() takes an optional auth_backend; production omits it and probes, tests pass it explicitly so no unit test reaches the network. Repins the Linux default to public.ecr.aws/k5j5w0x5/cua-ubuntu-24.04:main-38352d34 rather than :main-e5d853a9. The public main-38352d34 is bit-identical to the private image the PR was written against — the cold pull fetches the same layer sha256:7bd992eca8bf... (1438806575 bytes) the private registry served — so only the registry host changes and existing boot evidence still describes these bytes. Index descent now tests against the real two-child shape: the linux/amd64 disk plus the buildx provenance manifest that reports platform unknown/unknown, with a case proving the attestation is skipped even when listed first. Co-Authored-By: Claude Opus 5 (1M context) --- libs/python/cua-sandbox/cua_sandbox/image.py | 2 +- .../cua_sandbox/registry/container_disk.py | 56 ++---- .../cua-sandbox/tests/test_container_disk.py | 190 ++++++++++++------ .../tests/test_fleet_cloud_transport.py | 2 +- 4 files changed, 144 insertions(+), 106 deletions(-) diff --git a/libs/python/cua-sandbox/cua_sandbox/image.py b/libs/python/cua-sandbox/cua_sandbox/image.py index 30c44ff9f6..1fbd0d90d6 100644 --- a/libs/python/cua-sandbox/cua_sandbox/image.py +++ b/libs/python/cua-sandbox/cua_sandbox/image.py @@ -28,7 +28,7 @@ logger = logging.getLogger(__name__) -DEFAULT_LINUX_REGISTRY_IMAGE = "public.ecr.aws/k5j5w0x5/" "cua-ubuntu-24.04:main-e5d853a9" +DEFAULT_LINUX_REGISTRY_IMAGE = "public.ecr.aws/k5j5w0x5/" "cua-ubuntu-24.04:main-38352d34" _IMAGE_CACHE = Path.home() / ".cua" / "cua-sandbox" / "image-cache" diff --git a/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py b/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py index 8757deed61..d8fd56db32 100644 --- a/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py +++ b/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py @@ -14,6 +14,7 @@ from pathlib import Path from typing import Any, Optional +import oras.container import oras.provider import requests from cua_sandbox.registry.cache import CACHE_ROOT @@ -24,15 +25,15 @@ _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" +# oras fixes its auth backend at construction time and cannot negotiate one, but +# registries disagree about the scheme. Private ECR challenges with Basic — the token +# backend fails there with "Cannot respond to request for authentication" — while +# public.ecr.aws, ghcr.io and Docker Hub challenge with Bearer, where basic auth has +# no credential to send. Asking the registry which it wants beats maintaining a +# host list. _BASIC_AUTH_BACKEND = "basic" -_PING_TIMEOUT_SECONDS = 10 +_TOKEN_AUTH_BACKEND = "token" +_CHALLENGE_TIMEOUT_SECONDS = 30 _INDEX_MEDIA_TYPES = frozenset( { @@ -73,7 +74,7 @@ def pull_container_disk( if destination.exists(): return destination - registry = registry_factory(auth_backend=auth_backend or _detect_auth_backend(ref)) + registry = registry_factory(auth_backend=auth_backend or _auth_backend_for(ref)) container = registry.get_container(ref) registry.auth.load_configs(container) manifest = _resolve_platform_manifest(registry, ref, architecture or _host_architecture()) @@ -120,38 +121,17 @@ def pull_container_disk( 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 - +def _auth_backend_for(ref: str) -> str: + """Name the oras auth backend that matches the registry's authentication challenge.""" + url = f"https://{oras.container.Container(ref).manifest_url()}" 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) + response = requests.get(url, timeout=_CHALLENGE_TIMEOUT_SECONDS) + except requests.RequestException: 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) + challenge = response.headers.get("Www-Authenticate", "").strip().lower() + backend = _BASIC_AUTH_BACKEND if challenge.startswith("basic") else _TOKEN_AUTH_BACKEND + logger.debug("Registry challenged %s with %r; using %s auth", ref, challenge, backend) return backend diff --git a/libs/python/cua-sandbox/tests/test_container_disk.py b/libs/python/cua-sandbox/tests/test_container_disk.py index baf520ff7f..86bb446e17 100644 --- a/libs/python/cua-sandbox/tests/test_container_disk.py +++ b/libs/python/cua-sandbox/tests/test_container_disk.py @@ -11,51 +11,34 @@ from cua_sandbox.registry import container_disk from cua_sandbox.registry.container_disk import ( _LOCK_POLL_INTERVAL_SECONDS, - _detect_auth_backend, + _auth_backend_for, pull_container_disk, ) -BEARER_CHALLENGE = 'Bearer realm="https://public.ecr.aws/token/",service="public.ecr.aws"' -BASIC_CHALLENGE = 'Basic realm="https://296062593712.dkr.ecr.us-west-2.amazonaws.com/"' +# 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" - -def _fake_ping(monkeypatch, challenge, *, probes=None): - """Stub the /v2/ auth-challenge probe.""" - - def fake_get(url, timeout=None): - if probes is not None: - probes.append(url) - return SimpleNamespace(headers={"WWW-Authenticate": challenge} if challenge else {}) - - monkeypatch.setattr(container_disk.requests, "get", fake_get) - - -@pytest.fixture(autouse=True) -def _never_probe_a_real_registry(monkeypatch): - """Keep the auth-challenge probe off the network unless a test opts in.""" - _fake_ping(monkeypatch, BEARER_CHALLENGE) - - -# The real ECR image is a multi-arch index whose children are the per-platform manifests -# plus a buildx attestation manifest. INDEX_MANIFEST = { "schemaVersion": 2, "mediaType": "application/vnd.oci.image.index.v1+json", "manifests": [ { "mediaType": "application/vnd.oci.image.manifest.v1+json", - "digest": "sha256:arm64child", - "platform": {"architecture": "arm64", "os": "linux"}, - }, - { - "mediaType": "application/vnd.oci.image.manifest.v1+json", - "digest": "sha256:amd64child", + "digest": AMD64_CHILD, + "size": 483, "platform": {"architecture": "amd64", "os": "linux"}, }, { "mediaType": "application/vnd.oci.image.manifest.v1+json", - "digest": "sha256:attestation", - "annotations": {"vnd.docker.reference.type": "attestation-manifest"}, + "digest": ATTESTATION_CHILD, + "size": 563, + "annotations": { + "vnd.docker.reference.digest": AMD64_CHILD, + "vnd.docker.reference.type": "attestation-manifest", + }, "platform": {"architecture": "unknown", "os": "unknown"}, }, ], @@ -71,48 +54,47 @@ def _layer_with_disk(contents: bytes, *, path: str = "disk/disk.img") -> bytes: return gzip.compress(raw.getvalue()) -@pytest.mark.parametrize( - ("ref", "challenge", "expected"), - [ - # public.ecr.aws is anonymously readable but still uses the Bearer token flow. - ("public.ecr.aws/k5j5w0x5/cua-ubuntu-24.04:main-e5d853a9", BEARER_CHALLENGE, "token"), - # Private ECR answers with Basic and rejects oras' token backend. - ("296062593712.dkr.ecr.us-west-2.amazonaws.com/duo:main", BASIC_CHALLENGE, "basic"), - ("ghcr.io/trycua/workspace:latest", BEARER_CHALLENGE, "token"), - # No challenge at all (open registry) — assume the OCI-standard bearer flow. - ("registry.example/workspace:latest", "", "token"), - ], -) -def test_auth_backend_follows_the_registry_challenge(monkeypatch, ref, challenge, expected): - probes = [] - _fake_ping(monkeypatch, challenge, probes=probes) +class TestAuthBackendSelection: + """Registries disagree about the challenge scheme; oras cannot negotiate one.""" - assert _detect_auth_backend(ref) == expected - assert probes == [f"https://{ref.split('/', 1)[0]}/v2/"] + def _challenge(self, monkeypatch, header): + def get(url, **kwargs): + return SimpleNamespace(headers={"Www-Authenticate": header} if header else {}) + monkeypatch.setattr(container_disk.requests, "get", get) -def test_auth_backend_defaults_to_token_when_the_probe_fails(monkeypatch): - def boom(url, timeout=None): - raise requests.ConnectionError("no route to host") + def test_a_basic_challenge_selects_basic_auth(self, monkeypatch): + """Private ECR answers with Basic; oras' token backend cannot respond to it.""" + self._challenge(monkeypatch, 'Basic realm="https://12345.dkr.ecr.us-west-2.amazonaws.com/"') - monkeypatch.setattr(container_disk.requests, "get", boom) + assert _auth_backend_for("12345.dkr.ecr.us-west-2.amazonaws.com/workspace:tag") == "basic" - assert _detect_auth_backend("registry.example/workspace:latest") == "token" + @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 _auth_backend_for("public.ecr.aws/example/workspace:tag") == "token" -def test_auth_backend_skips_the_probe_for_a_hostless_ref(monkeypatch): - def boom(url, timeout=None): - raise AssertionError("a ref without a registry host must not be probed") + 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", boom) + monkeypatch.setattr(container_disk.requests, "get", get) - assert _detect_auth_backend("workspace:latest") == "token" + assert _auth_backend_for("registry.example/workspace:tag") == "token" -def test_explicit_auth_backend_overrides_detection(tmp_path, monkeypatch): +def test_explicit_auth_backend_skips_the_challenge_probe(tmp_path, monkeypatch): calls = [] - def boom(url, timeout=None): + def boom(url, **kwargs): raise AssertionError("an explicit auth_backend must not trigger a probe") monkeypatch.setattr(container_disk.requests, "get", boom) @@ -142,8 +124,7 @@ def get_blob(self, container, digest, *, stream): assert calls == [("init", "basic")] -def test_pull_container_disk_uses_oras_credentials_and_caches_qcow2(tmp_path, monkeypatch): - _fake_ping(monkeypatch, BEARER_CHALLENGE) +def test_pull_container_disk_uses_oras_credentials_and_caches_qcow2(tmp_path): calls = [] manifest = { "layers": [ @@ -176,13 +157,12 @@ def get_blob(self, container, digest, *, stream): 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" - # The backend comes from the registry's challenge — hardcoding either one breaks - # half the registries (see test_auth_backend_follows_the_registry_challenge). assert calls == [ ("init", "token"), ("container", DEFAULT_LINUX_REGISTRY_IMAGE), @@ -196,6 +176,7 @@ def get_blob(self, container, digest, *, stream): pull_container_disk( DEFAULT_LINUX_REGISTRY_IMAGE, cache_root=tmp_path, + auth_backend="token", registry_factory=Registry, ) == disk @@ -203,6 +184,32 @@ def get_blob(self, container, digest, *, stream): 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, "_auth_backend_for", 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 = [] @@ -235,12 +242,57 @@ def get_blob(self, container, digest, *, stream): 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] - assert requested == [DEFAULT_LINUX_REGISTRY_IMAGE, f"{repository}@sha256:amd64child"] + # 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): @@ -270,6 +322,7 @@ def get_blob(self, container, digest, *, stream): DEFAULT_LINUX_REGISTRY_IMAGE, cache_root=tmp_path, architecture="arm64", + auth_backend="token", registry_factory=Registry, ) @@ -301,6 +354,7 @@ def get_blob(self, container, digest, *, stream): pull_container_disk( "registry.example/lume-vm:latest", cache_root=tmp_path, + auth_backend="token", registry_factory=Registry, ) @@ -329,7 +383,10 @@ def get_blob(self, container, digest, *, stream): return SimpleNamespace(raw=io.BytesIO(payload)) disk = pull_container_disk( - "registry.example/workspace:latest", cache_root=tmp_path, registry_factory=Registry + "registry.example/workspace:latest", + cache_root=tmp_path, + auth_backend="token", + registry_factory=Registry, ) assert disk.read_bytes() == b"qcow2" @@ -359,6 +416,7 @@ def fake_sleep(interval): disk = pull_container_disk( DEFAULT_LINUX_REGISTRY_IMAGE, cache_root=tmp_path, + auth_backend="token", registry_factory=Registry, ) 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 36943216d0..c3a7b75904 100644 --- a/libs/python/cua-sandbox/tests/test_fleet_cloud_transport.py +++ b/libs/python/cua-sandbox/tests/test_fleet_cloud_transport.py @@ -76,7 +76,7 @@ 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-e5d853a9" + "public.ecr.aws/k5j5w0x5/" "cua-ubuntu-24.04:main-38352d34" ) From ab69e608e4132988605886cadbb1dae8715bdf5d Mon Sep 17 00:00:00 2001 From: Robert Wendt Date: Wed, 12 Aug 2026 20:19:57 +0000 Subject: [PATCH 8/8] refactor(sandbox): read the auth challenge from the /v2/ discovery endpoint Converges registry/container_disk.py on the Windows branch's shape so the file is byte-identical across both branches and rebases away to nothing when one merges. _detect_auth_backend() pings https://{host}/v2/ rather than fetching a manifest: /v2/ is the documented OCI discovery endpoint for reading a WWW-Authenticate challenge, and it is cheaper than pulling a manifest just to be refused. The _registry_host() helper returns None for refs with an implicit registry ("ubuntu:24.04"), which short-circuit to the token backend instead of having a manifest URL built for a host that was never named. Behavior is unchanged: public.ecr.aws still resolves to token, private ECR to basic, and a cold-cache pull plus KVM boot of the pinned public image still succeeds. Also collapses DEFAULT_LINUX_REGISTRY_IMAGE into a single literal now that it no longer needs the private registry's longer host. Co-Authored-By: Claude Opus 5 (1M context) --- libs/python/cua-sandbox/cua_sandbox/image.py | 2 +- .../cua_sandbox/registry/container_disk.py | 56 +++++++++++++------ .../cua-sandbox/tests/test_container_disk.py | 40 ++++++++++--- .../tests/test_fleet_cloud_transport.py | 2 +- 4 files changed, 73 insertions(+), 27 deletions(-) diff --git a/libs/python/cua-sandbox/cua_sandbox/image.py b/libs/python/cua-sandbox/cua_sandbox/image.py index 1fbd0d90d6..e66520bca7 100644 --- a/libs/python/cua-sandbox/cua_sandbox/image.py +++ b/libs/python/cua-sandbox/cua_sandbox/image.py @@ -28,7 +28,7 @@ logger = logging.getLogger(__name__) -DEFAULT_LINUX_REGISTRY_IMAGE = "public.ecr.aws/k5j5w0x5/" "cua-ubuntu-24.04:main-38352d34" +DEFAULT_LINUX_REGISTRY_IMAGE = "public.ecr.aws/k5j5w0x5/cua-ubuntu-24.04:main-38352d34" _IMAGE_CACHE = Path.home() / ".cua" / "cua-sandbox" / "image-cache" diff --git a/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py b/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py index d8fd56db32..8757deed61 100644 --- a/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py +++ b/libs/python/cua-sandbox/cua_sandbox/registry/container_disk.py @@ -14,7 +14,6 @@ from pathlib import Path from typing import Any, Optional -import oras.container import oras.provider import requests from cua_sandbox.registry.cache import CACHE_ROOT @@ -25,15 +24,15 @@ _CONTAINER_DISK_PATHS = {"disk/disk.img", "./disk/disk.img"} _LOCK_POLL_INTERVAL_SECONDS = 0.1 -# oras fixes its auth backend at construction time and cannot negotiate one, but -# registries disagree about the scheme. Private ECR challenges with Basic — the token -# backend fails there with "Cannot respond to request for authentication" — while -# public.ecr.aws, ghcr.io and Docker Hub challenge with Bearer, where basic auth has -# no credential to send. Asking the registry which it wants beats maintaining a -# host list. -_BASIC_AUTH_BACKEND = "basic" +# 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" -_CHALLENGE_TIMEOUT_SECONDS = 30 +_BASIC_AUTH_BACKEND = "basic" +_PING_TIMEOUT_SECONDS = 10 _INDEX_MEDIA_TYPES = frozenset( { @@ -74,7 +73,7 @@ def pull_container_disk( if destination.exists(): return destination - registry = registry_factory(auth_backend=auth_backend or _auth_backend_for(ref)) + 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()) @@ -121,17 +120,38 @@ def pull_container_disk( raise FileNotFoundError(f"OCI image {ref!r} does not contain /disk/disk.img") -def _auth_backend_for(ref: str) -> str: - """Name the oras auth backend that matches the registry's authentication challenge.""" - url = f"https://{oras.container.Container(ref).manifest_url()}" +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(url, timeout=_CHALLENGE_TIMEOUT_SECONDS) - except requests.RequestException: + 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 - challenge = response.headers.get("Www-Authenticate", "").strip().lower() - backend = _BASIC_AUTH_BACKEND if challenge.startswith("basic") else _TOKEN_AUTH_BACKEND - logger.debug("Registry challenged %s with %r; using %s auth", ref, challenge, 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 diff --git a/libs/python/cua-sandbox/tests/test_container_disk.py b/libs/python/cua-sandbox/tests/test_container_disk.py index 86bb446e17..fd2f2f9082 100644 --- a/libs/python/cua-sandbox/tests/test_container_disk.py +++ b/libs/python/cua-sandbox/tests/test_container_disk.py @@ -11,9 +11,10 @@ from cua_sandbox.registry import container_disk from cua_sandbox.registry.container_disk import ( _LOCK_POLL_INTERVAL_SECONDS, - _auth_backend_for, + _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" @@ -58,16 +59,25 @@ class TestAuthBackendSelection: """Registries disagree about the challenge scheme; oras cannot negotiate one.""" def _challenge(self, monkeypatch, header): + probed = [] + def get(url, **kwargs): - return SimpleNamespace(headers={"Www-Authenticate": header} if header else {}) + 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.""" - self._challenge(monkeypatch, 'Basic realm="https://12345.dkr.ecr.us-west-2.amazonaws.com/"') + 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 _auth_backend_for("12345.dkr.ecr.us-west-2.amazonaws.com/workspace:tag") == "basic" + assert backend == "basic" + assert probed == ["https://12345.dkr.ecr.us-west-2.amazonaws.com/v2/"] @pytest.mark.parametrize( "header", @@ -80,7 +90,23 @@ def test_a_basic_challenge_selects_basic_auth(self, monkeypatch): def test_a_bearer_or_absent_challenge_selects_token_auth(self, monkeypatch, header): self._challenge(monkeypatch, header) - assert _auth_backend_for("public.ecr.aws/example/workspace:tag") == "token" + 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): @@ -88,7 +114,7 @@ def get(url, **kwargs): monkeypatch.setattr(container_disk.requests, "get", get) - assert _auth_backend_for("registry.example/workspace:tag") == "token" + assert _detect_auth_backend("registry.example/workspace:tag") == "token" def test_explicit_auth_backend_skips_the_challenge_probe(tmp_path, monkeypatch): @@ -187,7 +213,7 @@ def get_blob(self, container, digest, *, stream): 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, "_auth_backend_for", lambda ref: "basic") + monkeypatch.setattr(container_disk, "_detect_auth_backend", lambda ref: "basic") class Registry: def __init__(self, *, auth_backend): 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 c3a7b75904..0b090f00e3 100644 --- a/libs/python/cua-sandbox/tests/test_fleet_cloud_transport.py +++ b/libs/python/cua-sandbox/tests/test_fleet_cloud_transport.py @@ -76,7 +76,7 @@ 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" + "public.ecr.aws/k5j5w0x5/cua-ubuntu-24.04:main-38352d34" )