feat(firecracker): pull pinned agent images

bot_bottle/backend/firecracker/image_builder.py

@@ -72,6 +72,46 @@ def cached_agent_rootfs_dir(dockerfile: Path) -> Path | None:
     return base if (base / ".bb-ready").is_file() else None
 
 
+def _image_rootfs_digest(image: str) -> str:
+    h = hashlib.sha256()
+    h.update(image.encode())
+    h.update(b"\0")
+    h.update(util._GUEST_INIT.encode())
+    return h.hexdigest()[:16]
+
+
+def cached_agent_image_rootfs_dir(image: str) -> Path | None:
+    """Return a ready rootfs exported from an immutable OCI image."""
+    base = util.cache_dir() / "rootfs" / f"agent-image-{_image_rootfs_digest(image)}"
+    return base if (base / ".bb-ready").is_file() else None
+
+
+def acquire_agent_image_rootfs_dir(
+    image: str, *, smoke_test: tuple[str, ...] = (),
+) -> Path:
+    """Pull a digest-pinned image in the infra VM and export its rootfs."""
+    if "@sha256:" not in image:
+        die(f"prebuilt Firecracker agent image is not digest-pinned: {image}")
+    digest = _image_rootfs_digest(image)
+    base = util.cache_dir() / "rootfs" / f"agent-image-{digest}"
+    cached = cached_agent_image_rootfs_dir(image)
+    if cached is not None:
+        info(f"using cached agent rootfs {cached.name}")
+        return cached
+    with _build_lock():
+        if (base / ".bb-ready").is_file():
+            return base
+        staging = util.cache_dir() / "rootfs" / f".pulling-{digest}"
+        shutil.rmtree(staging, ignore_errors=True)
+        staging.mkdir(parents=True)
+        _pull_in_infra(image, staging, smoke_test, digest)
+        util.inject_guest_boot(staging)
+        (staging / ".bb-ready").write_text("ok\n")
+        shutil.rmtree(base, ignore_errors=True)
+        os.rename(staging, base)
+    return base
+
+
 def build_agent_rootfs_dir(
     dockerfile: Path, *, image_tag: str, smoke_test: tuple[str, ...] = (),
 ) -> Path:
@@ -167,6 +207,43 @@ def _build_in_infra(
         _cleanup()
 
 
+def _pull_in_infra(
+    image: str, base: Path, smoke_test: tuple[str, ...], digest: str,
+) -> None:
+    """Pull and export a published agent image inside the orchestrator VM."""
+    service = FirecrackerInfraService()
+    service.ensure_running()
+    key, ip = service.orchestrator().ssh_target()
+    tag = f"bot-bottle-agent-pull-{digest}"
+    smoke_ctr, export_ctr = f"{tag}-smoke", f"{tag}-export"
+    quoted_image = shlex.quote(image)
+
+    def cleanup() -> None:
+        _ssh(
+            key,
+            ip,
+            f"buildah rm {smoke_ctr} {export_ctr} >/dev/null 2>&1; "
+            f"buildah rmi {_STORE_FLAG} {tag} >/dev/null 2>&1",
+            timeout=60,
+        )
+
+    cleanup()
+    try:
+        result = _ssh(
+            key,
+            ip,
+            f"buildah pull {_STORE_FLAG} {quoted_image} && "
+            f"buildah tag {_STORE_FLAG} {quoted_image} {tag}",
+            timeout=_BUILD_TIMEOUT_SECONDS,
+        )
+        if result.returncode != 0:
+            die(f"pulling pinned agent image failed: {result.stderr.strip()}")
+        _smoke_test(key, ip, tag, smoke_ctr, smoke_test)
+        _stream_rootfs(key, ip, tag, export_ctr, base)
+    finally:
+        cleanup()
+
+
 def _ssh(private_key: Path, guest_ip: str, script: str,
          *, timeout: float = 60.0) -> subprocess.CompletedProcess[str]:
     return subprocess.run(

bot_bottle/backend/firecracker/launch.py

@@ -228,8 +228,13 @@ def build_or_load_agent_base(plan: FirecrackerBottlePlan) -> Path:
         info(f"resuming from committed rootfs {committed_tar}")
         return util.build_committed_rootfs_dir(committed_tar)
     dockerfile = Path(plan.dockerfile_path)
+    prebuilt = "@sha256:" in plan.image
     if plan.spec.image_policy == "cached":
-        cached = image_builder.cached_agent_rootfs_dir(dockerfile)
+        cached = (
+            image_builder.cached_agent_image_rootfs_dir(plan.image)
+            if prebuilt
+            else image_builder.cached_agent_rootfs_dir(dockerfile)
+        )
         if cached is None:
             die(
                 f"cached agent rootfs for {plan.image!r} not found; "
@@ -237,6 +242,11 @@ def build_or_load_agent_base(plan: FirecrackerBottlePlan) -> Path:
             )
         info(f"using cached agent rootfs {cached.name}")
         return cached
+    if prebuilt:
+        return image_builder.acquire_agent_image_rootfs_dir(
+            plan.image,
+            smoke_test=runtime_for(plan.agent_provider_template).smoke_test,
+        )
     return image_builder.build_agent_rootfs_dir(
         dockerfile,
         image_tag=plan.image,
@@ -253,7 +263,11 @@ def stale_checks(plan: FirecrackerBottlePlan) -> None:
     if committed and committed_tar.is_file():
         check_stale_path(f"agent rootfs {committed_tar}", committed_tar)
         return
-    cached = image_builder.cached_agent_rootfs_dir(Path(plan.dockerfile_path))
+    cached = (
+        image_builder.cached_agent_image_rootfs_dir(plan.image)
+        if "@sha256:" in plan.image
+        else image_builder.cached_agent_rootfs_dir(Path(plan.dockerfile_path))
+    )
     if cached is not None:
         check_stale_path(f"agent rootfs {cached}", cached / ".bb-ready")
 

tests/unit/test_firecracker_image_builder.py

@@ -90,6 +90,20 @@ class TestBuildAgentRootfsDir(unittest.TestCase):
             second = image_builder._rootfs_digest(self.dockerfile)
         self.assertNotEqual(first, second)
 
+    def test_pinned_image_is_pulled_and_cached(self):
+        image = f"registry.example/agent@sha256:{'a' * 64}"
+        with patch.object(image_builder.util, "cache_dir", return_value=self.cache), \
+             patch.object(image_builder, "_pull_in_infra") as pull, \
+             patch.object(image_builder.util, "inject_guest_boot"):
+            first = image_builder.acquire_agent_image_rootfs_dir(image)
+            second = image_builder.acquire_agent_image_rootfs_dir(image)
+        pull.assert_called_once()
+        self.assertEqual(first, second)
+
+    def test_mutable_image_cannot_use_prebuilt_path(self):
+        with self.assertRaises(SystemExit):
+            image_builder.acquire_agent_image_rootfs_dir("agent:latest")
+
 
 class TestSmokeTest(unittest.TestCase):
     def test_buildah_receives_centralized_image_build_args(self):