From c8492e2adeb073ababb07fe03fce66dca95a1e8b Mon Sep 17 00:00:00 2001
From: morluto <76467478+morluto@users.noreply.github.com>
Date: Wed, 7 Oct 2026 09:29:03 +0800
Subject: [PATCH 1/4] fix(lifecycle): drain gateway work and verify VM image
 loads

---
 src/agent_env/env/gateway/triggers.py         |   7 +-
 .../providers/sandbox_providers/sandbox.py    |  50 +++--
 tst/unit/env/gateway/triggers_test.py         |  33 ++++
 .../sandbox_providers/vm_sandbox_test.py      | 172 ++++++++++++++++--
 4 files changed, 238 insertions(+), 24 deletions(-)

diff --git a/src/agent_env/env/gateway/triggers.py b/src/agent_env/env/gateway/triggers.py
index a2b07f77..823dba7f 100644
--- a/src/agent_env/env/gateway/triggers.py
+++ b/src/agent_env/env/gateway/triggers.py
@@ -538,12 +538,17 @@ async def _drive_time(self) -> None:
                 logger.exception("clock time-driver tick failed")
 
     async def stop_driver(self) -> None:
-        """Cancel the background poller on gateway teardown (idempotent)."""
+        """Cancel the poller and drain trigger work before gateway clients close."""
         task, self._driver_task = self._driver_task, None
         if task is not None:
             task.cancel()
             with contextlib.suppress(asyncio.CancelledError):
                 await task
+        while self._tasks:
+            tasks = tuple(self._tasks)
+            for pending in tasks:
+                pending.cancel()
+            await asyncio.gather(*tasks, return_exceptions=True)
 
     def _resolve_watch_roles(self, body: dict) -> set[str]:
         raw = body.get("watch_roles")
diff --git a/src/agent_env/providers/sandbox_providers/sandbox.py b/src/agent_env/providers/sandbox_providers/sandbox.py
index d89f8b09..eff555e5 100644
--- a/src/agent_env/providers/sandbox_providers/sandbox.py
+++ b/src/agent_env/providers/sandbox_providers/sandbox.py
@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import logging
 import os
 import posixpath
@@ -40,6 +41,7 @@
 # the partial bytes already consumed, corrupting the stream. Pipe consumers must
 # download to a temp file first, then read the file (see load_docker_images).
 CURL_RETRY_FLAGS = "--retry 5 --retry-all-errors --retry-delay 1"
+_DOCKER_IMAGE_INSPECT_FORMAT = "{{.Id}}"
 # Docker label on what a step starts on a sandbox's Docker host (containers, images, networks), valued with the
 # sandbox id, so a sandbox that shares its host (the local one) can remove its own when it terminates.
 SANDBOX_LABEL = "agentenv.sandbox"
@@ -281,31 +283,55 @@ async def sign(artifact) -> str | None:
 
     async def _load_docker_images(self, artifacts: list, signed_urls: list[str | None]) -> None:
         load_commands = []
+        staging_paths = [f"/tmp/_docker_image_{self.sandbox_id}_{idx}.tar.gz" for idx in range(len(artifacts))]
         for idx, (artifact, signed) in enumerate(zip(artifacts, signed_urls, strict=True)):
-            tmp_tar = f"/tmp/_docker_image_{self.sandbox_id}_{idx}.tar.gz"
+            tmp_tar = staging_paths[idx]
             if signed is not None:
                 # Download to a file first (retry-safe with -o); a `curl | ... docker load`
                 # pipe can't be retried without corrupting the stream (curl won't rewind).
                 load_commands.append(
-                    f'(curl -fsSL {CURL_RETRY_FLAGS} "{signed}" -o {shlex.quote(tmp_tar)} '
-                    f"&& gunzip -c {shlex.quote(tmp_tar)} | docker load && rm -f {shlex.quote(tmp_tar)})"
+                    f'(curl -fsSL {CURL_RETRY_FLAGS} {shlex.quote(signed)} -o {shlex.quote(tmp_tar)} '
+                    f"&& gunzip -c {shlex.quote(tmp_tar)} | docker load)"
                 )
             else:
-                await self._download_object_to_vm(artifact.tar_gz_object_url, tmp_tar)
                 load_commands.append(
-                    f"(gunzip -c {shlex.quote(tmp_tar)} | docker load && rm -f {shlex.quote(tmp_tar)})"
+                    f"(gunzip -c {shlex.quote(tmp_tar)} | docker load)"
                 )
             logger.info(f"  Queued: {artifact.image_name}")
-        await self.exec_script(" & ".join(load_commands) + " & wait", max_retries=2)
+        try:
+            for idx, (artifact, signed) in enumerate(zip(artifacts, signed_urls, strict=True)):
+                if signed is None:
+                    await self._download_object_to_vm(artifact.tar_gz_object_url, staging_paths[idx])
+            # exec_script explicitly invokes `bash -c`; pipefail and wait-status
+            # collection therefore use the shell the provider actually guarantees.
+            workers = "\n".join(
+                f"{command} & pid_{idx}=$!" for idx, command in enumerate(load_commands)
+            )
+            waits = "\n".join(f"wait $pid_{idx} || status=1" for idx in range(len(load_commands)))
+            cleanup_command = "rm -f " + " ".join(shlex.quote(path) for path in staging_paths)
+            script = (
+                f"trap {shlex.quote(cleanup_command)} EXIT\n"
+                "set -o pipefail\n"
+                f"{workers}\n"
+                "status=0\n"
+                f"{waits}\n"
+                "exit $status"
+            )
+            await self.exec_script(script)
+        except BaseException:
+            # Download failures happen before the shell trap can be installed.
+            with contextlib.suppress(Exception):
+                await self._remove_vm_temp_file(*staging_paths)
+            raise
 
         logger.info("Verifying Docker images...")
-        exit_code, stdout, stderr = await self.exec_with_output("sudo", "docker", "images")
+        image_refs = [artifact.image_name for artifact in artifacts]
+        exit_code, _, stderr = await self.exec_with_output(
+            "sudo", "docker", "image", "inspect", "--format", _DOCKER_IMAGE_INSPECT_FORMAT,
+            *image_refs,
+        )
         if exit_code != 0:
-            raise RuntimeError(f"docker images failed: {stderr}")
-        for artifact in artifacts:
-            base_name = artifact.image_name.split(":")[0]
-            if base_name not in stdout:
-                raise RuntimeError(f"{artifact.image_name} image not found. stdout: {stdout}")
+            raise RuntimeError(f"Docker image reference verification failed for {image_refs}: {stderr}")
         logger.info("  All images loaded successfully")
 
     async def load_s3_file(self, s3_url: str, destination_path: str) -> None:
diff --git a/tst/unit/env/gateway/triggers_test.py b/tst/unit/env/gateway/triggers_test.py
index 7f16d956..41ebe35f 100644
--- a/tst/unit/env/gateway/triggers_test.py
+++ b/tst/unit/env/gateway/triggers_test.py
@@ -96,6 +96,39 @@ def _registration(**overrides):
     return base
 
 
+@pytest.mark.asyncio
+async def test_stop_driver_cancels_and_awaits_pending_action_before_client_teardown(engine):
+    started = asyncio.Event()
+    cancelled = asyncio.Event()
+
+    async def slow_call(_tool_name, _arguments):
+        started.set()
+        try:
+            await asyncio.Event().wait()
+        except asyncio.CancelledError:
+            cancelled.set()
+            raise
+
+    engine._internal_call = slow_call
+    engine.register({"watch_roles": ["default"], "triggers": [
+        {"id": "shutdown", "when": {"type": "action", "tool": "provoking_tool"},
+         "actions": [{"type": "tool", "tool": "slack_send_message", "args": {}}]}
+    ]})
+    engine.start_driver()
+    engine.on_tool_call("default", "provoking_tool", {}, _result())
+    await asyncio.wait_for(started.wait(), 1)
+    tracked = tuple(engine._tasks)
+    assert tracked
+
+    await engine.stop_driver()
+
+    assert cancelled.is_set()
+    assert all(task.done() for task in tracked)
+    assert not engine._tasks
+    await engine.stop_driver()
+    assert not engine._tasks
+
+
 @pytest.mark.parametrize("bad,fragment", [
     ({"triggers": "nope"}, "triggers must be a list"),
     ({"triggers": [{"id": "x", "when": {"type": "step"}, "actions": []}]}, "when.type"),
diff --git a/tst/unit/providers/sandbox_providers/vm_sandbox_test.py b/tst/unit/providers/sandbox_providers/vm_sandbox_test.py
index 75ae6d81..7b3d4e9f 100644
--- a/tst/unit/providers/sandbox_providers/vm_sandbox_test.py
+++ b/tst/unit/providers/sandbox_providers/vm_sandbox_test.py
@@ -1,4 +1,7 @@
 import asyncio
+import os
+import subprocess
+from pathlib import Path
 import threading
 import time
 from concurrent.futures import ThreadPoolExecutor
@@ -33,10 +36,9 @@ def signing_store():
 class _RecordingVmSandbox(VmSandbox):
     """Concrete VmSandbox that records executed scripts instead of touching a VM."""
 
-    def __init__(self, images_stdout: str = ""):
+    def __init__(self):
         self.sandbox_id = "vm-test"
         self.scripts: list[str] = []
-        self._images_stdout = images_stdout
 
     async def terminate(self) -> None:  # pragma: no cover - not exercised
         pass
@@ -45,10 +47,9 @@ async def exec(self, *command):  # pragma: no cover - not exercised
         return None
 
     async def exec_with_output(self, *args):
-        # load_docker_images runs `sudo docker images` to verify; everything
-        # else flows through exec_script as `sudo bash -c <script>`.
-        if args[:2] == ("sudo", "docker") and "images" in args:
-            return 0, self._images_stdout, ""
+        # Image loading verifies exact refs through `docker image inspect`.
+        if args[:4] == ("sudo", "docker", "image", "inspect"):
+            return 0, "sha256:image", ""
         if args[:2] == ("sudo", "bash"):
             self.scripts.append(script_run(args))
         return 0, "", ""
@@ -65,7 +66,7 @@ async def test_load_docker_images_downloads_to_file_before_load(signing_store):
     because curl can't rewind bytes already piped to stdout, so the retry flags
     must only apply to a `-o file` download.
     """
-    sandbox = _RecordingVmSandbox(images_stdout="myimage\n")
+    sandbox = _RecordingVmSandbox()
     artifact = SimpleNamespace(tar_gz_object_url="s3://bucket/img.tar.gz", image_name="myimage:latest")
 
     await sandbox.load_docker_images([artifact])
@@ -79,11 +80,14 @@ async def test_load_docker_images_downloads_to_file_before_load(signing_store):
     assert "| gunzip" not in load_script.split("-o")[0]
     # gunzip reads from the downloaded file, not curl's stdout.
     assert "gunzip -c" in load_script
+    assert "set -o pipefail" in load_script
+    assert "wait $pid_0 || status=1" in load_script
+    assert "trap 'rm -f" in load_script
 
 
 @pytest.mark.asyncio
 async def test_load_docker_images_uses_unique_tmp_per_artifact(signing_store):
-    sandbox = _RecordingVmSandbox(images_stdout="a\nb\n")
+    sandbox = _RecordingVmSandbox()
     artifacts = [
         SimpleNamespace(tar_gz_object_url="s3://bucket/a.tar.gz", image_name="a:1"),
         SimpleNamespace(tar_gz_object_url="s3://bucket/b.tar.gz", image_name="b:1"),
@@ -96,6 +100,152 @@ async def test_load_docker_images_uses_unique_tmp_per_artifact(signing_store):
     assert "/tmp/_docker_image_vm-test_1.tar.gz" in load_script
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("failure_stage", "expected_loads"),
+    [("docker", ["fail", "ok"]), ("gunzip", ["ok", "ok"])],
+)
+async def test_load_image_worker_failure_propagates_and_cleans_staging(
+    tmp_path, failure_stage, expected_loads,
+):
+    bin_dir = tmp_path / "bin"
+    bin_dir.mkdir()
+    docker_log = tmp_path / "docker-load.log"
+    _install_docker_load_stubs(
+        bin_dir,
+        fail_first=failure_stage == "docker",
+        gunzip_fail_first=failure_stage == "gunzip",
+    )
+    sandbox = _ShellVmSandbox(bin_dir, docker_log, f"lifecycle-{tmp_path.name}")
+    images = [
+        SimpleNamespace(tar_gz_object_url=f"s3://bucket/{i}", image_name=f"repo:{i}")
+        for i in range(2)
+    ]
+    with pytest.raises(RuntimeError, match="load script exited"):
+        await sandbox._load_docker_images(images, ["https://signed/a", "https://signed/b"])
+    assert sorted(docker_log.read_text().splitlines()) == expected_loads
+    assert all(
+        not Path(f"/tmp/_docker_image_{sandbox.sandbox_id}_{i}.tar.gz").exists()
+        for i in range(2)
+    )
+    assert not any("image" in call and "inspect" in call for call in sandbox.verify_calls)
+
+
+@pytest.mark.asyncio
+async def test_load_image_success_runs_every_worker_and_verifies_refs(tmp_path):
+    bin_dir = tmp_path / "bin"
+    bin_dir.mkdir()
+    docker_log = tmp_path / "docker-load.log"
+    _install_docker_load_stubs(bin_dir)
+
+    sandbox = _ShellVmSandbox(bin_dir, docker_log, f"lifecycle-{tmp_path.name}")
+    images = [
+        SimpleNamespace(tar_gz_object_url=f"s3://bucket/{i}", image_name=f"repo:{i}")
+        for i in range(2)
+    ]
+    await sandbox._load_docker_images(images, ["https://signed/a", "https://signed/b"])
+    assert sorted(docker_log.read_text().splitlines()) == ["ok", "ok"]
+    assert sandbox.verify_calls == [(
+        "sudo", "docker", "image", "inspect", "--format", "{{.Id}}", "repo:0", "repo:1",
+    )]
+
+
+class _ShellVmSandbox(_RecordingVmSandbox):
+    def __init__(self, bin_dir, docker_log, sandbox_id):
+        super().__init__()
+        self.bin_dir = bin_dir
+        self.docker_log = docker_log
+        self.sandbox_id = sandbox_id
+        self.verify_calls = []
+
+    async def exec_script(self, script, *, max_retries=0):
+        env = dict(
+            os.environ,
+            PATH=f"{self.bin_dir}:{os.environ['PATH']}",
+            DOCKER_LOG=str(self.docker_log),
+        )
+        result = subprocess.run(["bash", "-c", script], env=env, capture_output=True, text=True)
+        if result.returncode:
+            raise RuntimeError(f"load script exited {result.returncode}")
+        self.scripts.append(script)
+
+    async def exec_with_output(self, *args):
+        self.verify_calls.append(args)
+        return await super().exec_with_output(*args)
+
+
+def _install_docker_load_stubs(bin_dir, *, fail_first=False, gunzip_fail_first=False):
+    curl_result = (
+        'case "$url" in */a) printf fail > "$1" ;; *) printf ok > "$1" ;; esac\n'
+        if fail_first else 'printf ok > "$1"\n'
+    )
+    gunzip_script = '#!/bin/sh\ncat "$2"\n'
+    if gunzip_fail_first:
+        gunzip_script = '#!/bin/sh\ncat "$2"\ncase "$2" in *_0.tar.gz) exit 24 ;; esac\n'
+    for name, body in {
+        "curl": (
+            '#!/bin/sh\nurl=\nwhile [ "$1" != "-o" ]; do url=$1; shift; done\n'
+            "shift\n" + curl_result
+        ),
+        "gunzip": gunzip_script,
+        "docker": '#!/bin/sh\npayload=$(cat)\nprintf "%s\\n" "$payload" >> "$DOCKER_LOG"\n'
+        '[ "$payload" = fail ] && exit 23\nexit 0\n',
+    }.items():
+        path = bin_dir / name
+        path.write_text(body)
+        path.chmod(0o755)
+
+
+@pytest.mark.asyncio
+async def test_load_image_requires_exact_reference_even_when_old_tag_is_listed(signing_store):
+    class StaleTagVm(_RecordingVmSandbox):
+        async def exec_with_output(self, *args):
+            if args[:4] == ("sudo", "docker", "image", "inspect"):
+                existing_refs = {"registry.example:5000/team/repo:old-tag"}
+                requested_refs = set(args[6:])
+                if requested_refs <= existing_refs:
+                    return 0, "sha256:old", ""
+                return 1, "", "No such image"
+            return await super().exec_with_output(*args)
+
+    sandbox = StaleTagVm()
+    image = SimpleNamespace(
+        tar_gz_object_url="s3://bucket/img",
+        image_name="registry.example:5000/team/repo:new-tag",
+    )
+    with pytest.raises(RuntimeError, match="registry.example:5000/team/repo:new-tag"):
+        await sandbox.load_docker_images([image])
+
+
+@pytest.mark.asyncio
+async def test_unsigned_download_failure_cleans_staging_file(signing_store, tmp_path):
+    cleaned = []
+
+    class UnsignedStore(_SigningStore):
+        def signed_get_url(self, object_url, expires_in=3600):
+            return None
+
+    set_object_store(UnsignedStore())
+
+    class FailedDownloadVm(_RecordingVmSandbox):
+        async def _download_object_to_vm(self, _object_url, vm_path):
+            Path(vm_path).write_text("partial")
+            raise RuntimeError("download failed")
+
+        async def _remove_vm_temp_file(self, *vm_paths):
+            cleaned.extend(vm_paths)
+            for vm_path in vm_paths:
+                Path(vm_path).unlink(missing_ok=True)
+
+    sandbox = FailedDownloadVm()
+    sandbox.sandbox_id = f"lifecycle-{tmp_path.name}"
+    artifact = SimpleNamespace(tar_gz_object_url="obj://image", image_name="repo:tag")
+    with pytest.raises(RuntimeError, match="download failed"):
+        await sandbox._load_docker_images([artifact], [None])
+    assert cleaned == [f"/tmp/_docker_image_{sandbox.sandbox_id}_0.tar.gz"]
+    assert not Path(cleaned[0]).exists()
+
+
 class _ScriptRecorder(VmSandbox):
     """Records exec_script invocations, running the real write_file_from_text."""
 
@@ -238,7 +388,7 @@ def get(self, object_url):
 
     set_object_store(_LocalStore())
     try:
-        sandbox = _RecordingVmSandbox(images_stdout="myimage\n")
+        sandbox = _RecordingVmSandbox()
         artifact = SimpleNamespace(tar_gz_object_url="file:///store/img.tar.gz", image_name="myimage:latest")
         await sandbox.load_docker_images([artifact])
     finally:
@@ -279,7 +429,7 @@ async def test_the_store_is_called_off_the_event_loop(signs):
     store = _LoopCheckingStore(signs)
     set_object_store(store)
     try:
-        sandbox = _RecordingVmSandbox(images_stdout="a\nb\n")
+        sandbox = _RecordingVmSandbox()
         artifacts = [SimpleNamespace(tar_gz_object_url=f"s3://bucket/{n}.tar.gz", image_name=f"{n}:1") for n in "ab"]
         await sandbox.load_docker_images(artifacts)
         await sandbox.load_s3_file("s3://bucket/data.json", "/tmp/data.json")
@@ -365,7 +515,7 @@ async def test_a_load_bounds_how_many_signs_run_at_once():
     asyncio.get_running_loop().set_default_executor(ThreadPoolExecutor(max_workers=32))
     try:
         names = [f"i{n}" for n in range(20)]
-        sandbox = _RecordingVmSandbox(images_stdout="\n".join(names))
+        sandbox = _RecordingVmSandbox()
         await sandbox.load_docker_images(
             [SimpleNamespace(tar_gz_object_url=f"s3://bucket/{n}.tar.gz", image_name=f"{n}:1") for n in names]
         )

From c7acaa69e23734358d941437dbc577f4b1d18c9f Mon Sep 17 00:00:00 2001
From: morluto <76467478+morluto@users.noreply.github.com>
Date: Wed, 7 Oct 2026 09:44:42 +0800
Subject: [PATCH 2/4] fix(sandbox): retry image loads and clean staging files

---
 .../providers/sandbox_providers/sandbox.py    |  30 ++--
 .../sandbox_providers/vm_sandbox_test.py      | 129 ++++++++++++++++--
 2 files changed, 128 insertions(+), 31 deletions(-)

diff --git a/src/agent_env/providers/sandbox_providers/sandbox.py b/src/agent_env/providers/sandbox_providers/sandbox.py
index eff555e5..2d10c682 100644
--- a/src/agent_env/providers/sandbox_providers/sandbox.py
+++ b/src/agent_env/providers/sandbox_providers/sandbox.py
@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import asyncio
-import contextlib
+import base64
 import logging
 import os
 import posixpath
@@ -42,6 +42,7 @@
 # download to a temp file first, then read the file (see load_docker_images).
 CURL_RETRY_FLAGS = "--retry 5 --retry-all-errors --retry-delay 1"
 _DOCKER_IMAGE_INSPECT_FORMAT = "{{.Id}}"
+_DOCKER_IMAGE_LOAD_EXEC_RETRIES = 2
 # Docker label on what a step starts on a sandbox's Docker host (containers, images, networks), valued with the
 # sandbox id, so a sandbox that shares its host (the local one) can remove its own when it terminates.
 SANDBOX_LABEL = "agentenv.sandbox"
@@ -308,21 +309,17 @@ async def _load_docker_images(self, artifacts: list, signed_urls: list[str | Non
                 f"{command} & pid_{idx}=$!" for idx, command in enumerate(load_commands)
             )
             waits = "\n".join(f"wait $pid_{idx} || status=1" for idx in range(len(load_commands)))
-            cleanup_command = "rm -f " + " ".join(shlex.quote(path) for path in staging_paths)
             script = (
-                f"trap {shlex.quote(cleanup_command)} EXIT\n"
                 "set -o pipefail\n"
                 f"{workers}\n"
                 "status=0\n"
                 f"{waits}\n"
                 "exit $status"
             )
-            await self.exec_script(script)
-        except BaseException:
-            # Download failures happen before the shell trap can be installed.
-            with contextlib.suppress(Exception):
-                await self._remove_vm_temp_file(*staging_paths)
-            raise
+            # Signed archives are downloaded per attempt; keep unsigned archives staged until retries finish.
+            await self.exec_script(script, max_retries=_DOCKER_IMAGE_LOAD_EXEC_RETRIES)
+        finally:
+            await self._remove_vm_temp_file(*staging_paths)
 
         logger.info("Verifying Docker images...")
         image_refs = [artifact.image_name for artifact in artifacts]
@@ -402,18 +399,19 @@ async def write_file_from_url(self, url: str, destination_path: str) -> None:
 
     async def _write_bytes_to_vm_path(self, data: bytes, vm_path: str) -> None:
         """Stream bytes from agent-env onto the VM host at vm_path (base64 over exec)."""
-        import base64
-
         encoded = base64.b64encode(data).decode()
         if len(encoded) <= self._WFT_CHUNK_BYTES:
             await self.exec_script(f"base64 -d <<'ENDB64' > {shlex.quote(vm_path)}\n{encoded}\nENDB64")
             return
         # Too big for one heredoc arg: append the (shell-safe) base64 in bounded chunks.
         vm_b64 = f"{vm_path}.b64"
-        await self.exec_script(f": > {shlex.quote(vm_b64)}")
-        for i in range(0, len(encoded), self._WFT_CHUNK_BYTES):
-            await self.exec_script(f"printf '%s' {shlex.quote(encoded[i:i + self._WFT_CHUNK_BYTES])} >> {shlex.quote(vm_b64)}")
-        await self.exec_script(f"base64 -d {shlex.quote(vm_b64)} > {shlex.quote(vm_path)} && rm -f {shlex.quote(vm_b64)}")
+        try:
+            await self.exec_script(f": > {shlex.quote(vm_b64)}")
+            for i in range(0, len(encoded), self._WFT_CHUNK_BYTES):
+                await self.exec_script(f"printf '%s' {shlex.quote(encoded[i:i + self._WFT_CHUNK_BYTES])} >> {shlex.quote(vm_b64)}")
+            await self.exec_script(f"base64 -d {shlex.quote(vm_b64)} > {shlex.quote(vm_path)}")
+        finally:
+            await self._remove_vm_temp_file(vm_b64)
 
     async def write_host_file(self, data: bytes, vm_path: str) -> None:
         """Write bytes to vm_path on the VM host itself, not into the agent container."""
@@ -428,7 +426,7 @@ async def write_file_from_text(self, content: str, destination_path: str) -> Non
             await self._write_bytes_to_vm_path(content.encode(), vm_path)
             await self._copy_into_container(vm_path, destination_path)
         finally:
-            await self._remove_vm_temp_file(vm_path, f"{vm_path}.b64")
+            await self._remove_vm_temp_file(vm_path)
 
 
 def port_bindings(host_ips: Iterable[str], host_port: int, container_port: int) -> list[str]:
diff --git a/tst/unit/providers/sandbox_providers/vm_sandbox_test.py b/tst/unit/providers/sandbox_providers/vm_sandbox_test.py
index 7b3d4e9f..6ab07956 100644
--- a/tst/unit/providers/sandbox_providers/vm_sandbox_test.py
+++ b/tst/unit/providers/sandbox_providers/vm_sandbox_test.py
@@ -82,7 +82,7 @@ async def test_load_docker_images_downloads_to_file_before_load(signing_store):
     assert "gunzip -c" in load_script
     assert "set -o pipefail" in load_script
     assert "wait $pid_0 || status=1" in load_script
-    assert "trap 'rm -f" in load_script
+    assert "trap " not in load_script
 
 
 @pytest.mark.asyncio
@@ -121,13 +121,14 @@ async def test_load_image_worker_failure_propagates_and_cleans_staging(
         SimpleNamespace(tar_gz_object_url=f"s3://bucket/{i}", image_name=f"repo:{i}")
         for i in range(2)
     ]
-    with pytest.raises(RuntimeError, match="load script exited"):
+    with pytest.raises(RuntimeError, match=r"Script failed \(exit 1\)"):
         await sandbox._load_docker_images(images, ["https://signed/a", "https://signed/b"])
     assert sorted(docker_log.read_text().splitlines()) == expected_loads
     assert all(
         not Path(f"/tmp/_docker_image_{sandbox.sandbox_id}_{i}.tar.gz").exists()
         for i in range(2)
     )
+    assert len([script for script in sandbox.scripts if "docker load" in script]) == 1
     assert not any("image" in call and "inspect" in call for call in sandbox.verify_calls)
 
 
@@ -150,6 +151,48 @@ async def test_load_image_success_runs_every_worker_and_verifies_refs(tmp_path):
     )]
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("signed", [True, False], ids=["signed", "unsigned"])
+async def test_load_image_retries_after_transient_exec_result_without_losing_archive(
+    tmp_path, monkeypatch, signed,
+):
+    bin_dir = tmp_path / "bin"
+    bin_dir.mkdir()
+    docker_log = tmp_path / "docker-load.log"
+    _install_docker_load_stubs(bin_dir)
+    sandbox = _ShellVmSandbox(bin_dir, docker_log, f"lifecycle-{tmp_path.name}")
+    sandbox.transient_results = 1
+    artifacts = [
+        SimpleNamespace(tar_gz_object_url=f"s3://bucket/{idx}", image_name=f"repo:{idx}")
+        for idx in range(2)
+    ]
+
+    if not signed:
+        downloaded = []
+
+        async def save_unsigned(_object_url, vm_path):
+            downloaded.append(vm_path)
+            Path(vm_path).write_text("ok")
+
+        sandbox._download_object_to_vm = save_unsigned
+    async def no_sleep(_delay):
+        return None
+
+    monkeypatch.setattr(sandbox_module.asyncio, "sleep", no_sleep)
+    await sandbox._load_docker_images(
+        artifacts,
+        ["https://signed/a", "https://signed/b"] if signed else [None, None],
+    )
+
+    assert sorted(docker_log.read_text().splitlines()) == ["ok"] * 4
+    assert len([script for script in sandbox.scripts if "docker load" in script]) == 2
+    assert not any(Path(path).exists() for path in (
+        f"/tmp/_docker_image_{sandbox.sandbox_id}_{idx}.tar.gz" for idx in range(2)
+    ))
+    if not signed:
+        assert len(downloaded) == 2
+
+
 class _ShellVmSandbox(_RecordingVmSandbox):
     def __init__(self, bin_dir, docker_log, sandbox_id):
         super().__init__()
@@ -157,20 +200,24 @@ def __init__(self, bin_dir, docker_log, sandbox_id):
         self.docker_log = docker_log
         self.sandbox_id = sandbox_id
         self.verify_calls = []
-
-    async def exec_script(self, script, *, max_retries=0):
-        env = dict(
-            os.environ,
-            PATH=f"{self.bin_dir}:{os.environ['PATH']}",
-            DOCKER_LOG=str(self.docker_log),
-        )
-        result = subprocess.run(["bash", "-c", script], env=env, capture_output=True, text=True)
-        if result.returncode:
-            raise RuntimeError(f"load script exited {result.returncode}")
-        self.scripts.append(script)
+        self.transient_results = 0
 
     async def exec_with_output(self, *args):
-        self.verify_calls.append(args)
+        if args[:4] == ("sudo", "docker", "image", "inspect"):
+            self.verify_calls.append(args)
+        if args[:3] == ("sudo", "bash", "-c"):
+            script = args[3]
+            self.scripts.append(script)
+            env = dict(
+                os.environ,
+                PATH=f"{self.bin_dir}:{os.environ['PATH']}",
+                DOCKER_LOG=str(self.docker_log),
+            )
+            result = subprocess.run(["bash", "-c", script], env=env, capture_output=True, text=True)
+            if self.transient_results:
+                self.transient_results -= 1
+                return -1, result.stdout, result.stderr
+            return result.returncode, result.stdout, result.stderr
         return await super().exec_with_output(*args)
 
 
@@ -246,6 +293,57 @@ async def _remove_vm_temp_file(self, *vm_paths):
     assert not Path(cleaned[0]).exists()
 
 
+@pytest.mark.asyncio
+async def test_failed_chunk_write_removes_base64_staging_file(tmp_path):
+    sandbox = _ChunkWriteSandbox()
+    destination = tmp_path / "payload.bin"
+    staging = Path(f"{destination}.b64")
+    with pytest.raises(RuntimeError, match=r"Script failed \(exit 17\)"):
+        await sandbox._write_bytes_to_vm_path(b"x" * (VmSandbox._WFT_CHUNK_BYTES + 1), str(destination))
+    assert not staging.exists()
+    assert sandbox.scripts[-1] == f"rm -f {staging}"
+
+
+@pytest.mark.asyncio
+async def test_cancelled_chunk_write_removes_base64_staging_file(tmp_path):
+    sandbox = _ChunkWriteSandbox(cancel_chunk=True)
+    destination = tmp_path / "payload.bin"
+    staging = Path(f"{destination}.b64")
+    write = asyncio.create_task(
+        sandbox._write_bytes_to_vm_path(b"x" * (VmSandbox._WFT_CHUNK_BYTES + 1), str(destination))
+    )
+    await asyncio.wait_for(sandbox.chunk_started.wait(), 1)
+    write.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await write
+    assert not staging.exists()
+    assert sandbox.scripts[-1] == f"rm -f {staging}"
+
+
+class _ChunkWriteSandbox(VmSandbox):
+    def __init__(self, *, cancel_chunk=False):
+        self.scripts = []
+        self.cancel_chunk = cancel_chunk
+        self.chunk_started = asyncio.Event()
+
+    async def terminate(self):  # pragma: no cover
+        pass
+
+    async def exec(self, *command):  # pragma: no cover
+        return None
+
+    async def exec_with_output(self, *args):
+        script = script_run(args)
+        self.scripts.append(script)
+        if script.startswith("printf '%s'"):
+            if self.cancel_chunk:
+                self.chunk_started.set()
+                await asyncio.Future()
+            return 17, "", "chunk write failed"
+        result = subprocess.run(["bash", "-c", script], capture_output=True, text=True)
+        return result.returncode, result.stdout, result.stderr
+
+
 class _ScriptRecorder(VmSandbox):
     """Records exec_script invocations, running the real write_file_from_text."""
 
@@ -536,7 +634,8 @@ async def test_write_host_file_writes_on_the_host_in_bounded_chunks():
     assert not any("docker" in s for s in sandbox.scripts)
     b64_path = "/tmp/agentenv_run_code/input.json.b64"
     assert _b64_from_chunk_scripts(sandbox.scripts, b64_path) == base64.b64encode(data).decode()
-    assert sandbox.scripts[-1] == f"base64 -d {b64_path} > /tmp/agentenv_run_code/input.json && rm -f {b64_path}"
+    assert f"base64 -d {b64_path} > /tmp/agentenv_run_code/input.json" in sandbox.scripts
+    assert sandbox.scripts[-1] == f"rm -f {b64_path}"
 
 
 class _ArgsRecorder(VmSandbox):

From 8ab23b5dc9b6fe45cd44afaa4a5e09973f44489e Mon Sep 17 00:00:00 2001
From: morluto <76467478+morluto@users.noreply.github.com>
Date: Thu, 8 Oct 2026 02:24:34 +0800
Subject: [PATCH 3/4] fix(sandbox): preserve deprecated provider overrides

---
 .../providers/sandbox_providers/sandbox.py    |  33 +++++-
 .../sandbox_method_compatibility_test.py      | 109 ++++++++++++++++++
 2 files changed, 140 insertions(+), 2 deletions(-)
 create mode 100644 tst/unit/providers/sandbox_providers/sandbox_method_compatibility_test.py

diff --git a/src/agent_env/providers/sandbox_providers/sandbox.py b/src/agent_env/providers/sandbox_providers/sandbox.py
index 779d5445..e143efad 100644
--- a/src/agent_env/providers/sandbox_providers/sandbox.py
+++ b/src/agent_env/providers/sandbox_providers/sandbox.py
@@ -136,6 +136,21 @@ def from_dict(cls, data: dict) -> "NetworkPolicy":
         )
 
 
+def _object_file_method(legacy: Callable) -> Callable:
+    async def method(self, object_url: str, destination_path: str) -> None:
+        await legacy(self, object_url, destination_path)
+
+    return method
+
+
+def _legacy_file_method(canonical: Callable, old_symbol: str, new_name: str) -> Callable:
+    async def method(self, s3_url: str, destination_path: str) -> None:
+        warn_deprecated(old_symbol, new_name, kind="method")
+        await canonical(self, s3_url, destination_path)
+
+    return method
+
+
 class Sandbox(ABC):
     """Universal sandbox contract — anything that can host a process and expose ports."""
 
@@ -152,6 +167,20 @@ class Sandbox(ABC):
     _VM_READY_TIMEOUT = 1200      # wait_for_vm wall-clock budget (s)
     _VM_READY_POLL_INTERVAL = 30  # sparse polling (s)
 
+    def __init_subclass__(cls, **kwargs) -> None:
+        """Keep overrides of either file-method spelling in the dispatch path, including super() calls."""
+        super().__init_subclass__(**kwargs)
+        for legacy, neutral, owner in (
+            ("write_file_from_s3", "write_file_from_object", "Sandbox"),
+            ("load_s3_file", "load_object_file", "VmSandbox"),
+        ):
+            legacy_impl = cls.__dict__.get(legacy)
+            neutral_impl = cls.__dict__.get(neutral)
+            if legacy_impl is not None and neutral_impl is None:
+                setattr(cls, neutral, _object_file_method(legacy_impl))
+            elif neutral_impl is not None and legacy_impl is None:
+                setattr(cls, legacy, _legacy_file_method(neutral_impl, f"{owner}.{legacy}", neutral))
+
     def host_port(self, port: int) -> int:
         """The host-side port a published container port is reachable on.
 
@@ -184,7 +213,7 @@ async def write_file_from_object(self, object_url: str, destination_path: str) -
     async def write_file_from_s3(self, s3_url: str, destination_path: str) -> None:
         """Deprecated: ``write_file_from_object``."""
         warn_deprecated("Sandbox.write_file_from_s3", "write_file_from_object", kind="method")
-        await self.write_file_from_object(s3_url, destination_path)
+        await Sandbox.write_file_from_object(self, s3_url, destination_path)
 
     async def write_file_from_url(self, url: str, destination_path: str) -> None:
         """Download an HTTP(S) URL into the agent process's filesystem at destination_path."""
@@ -363,7 +392,7 @@ async def load_object_file(self, object_url: str, destination_path: str) -> None
     async def load_s3_file(self, s3_url: str, destination_path: str) -> None:
         """Deprecated: ``load_object_file``."""
         warn_deprecated("VmSandbox.load_s3_file", "load_object_file", kind="method")
-        await self.load_object_file(s3_url, destination_path)
+        await VmSandbox.load_object_file(self, s3_url, destination_path)
 
     async def _download_object_to_vm(self, object_url: str, vm_path: str) -> None:
         """Place object_url onto the VM host at vm_path, backend-agnostically."""
diff --git a/tst/unit/providers/sandbox_providers/sandbox_method_compatibility_test.py b/tst/unit/providers/sandbox_providers/sandbox_method_compatibility_test.py
new file mode 100644
index 00000000..daffb84d
--- /dev/null
+++ b/tst/unit/providers/sandbox_providers/sandbox_method_compatibility_test.py
@@ -0,0 +1,109 @@
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+from agent_env.providers.sandbox_providers.sandbox import Sandbox, VmSandbox, stage_files_into_container
+
+
+_METHODS = (
+    (Sandbox, "write_file_from_s3", "write_file_from_object"),
+    (VmSandbox, "write_file_from_s3", "write_file_from_object"),
+    (VmSandbox, "load_s3_file", "load_object_file"),
+)
+
+
+async def _terminate(self):
+    pass
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("base, legacy, neutral", _METHODS)
+async def test_new_calls_reach_inherited_legacy_provider_overrides(base, legacy, neutral):
+    calls = []
+
+    async def implementation(self, s3_url, destination_path):
+        calls.append((s3_url, destination_path))
+
+    provider = type("LegacySandbox", (base,), {legacy: implementation, "terminate": _terminate})
+    child = type("InheritedSandbox", (provider,), {})
+
+    await getattr(child(), neutral)(object_url="file:///objects/input", destination_path="/app/input")
+
+    assert calls == [("file:///objects/input", "/app/input")]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("base, legacy, neutral", _METHODS)
+async def test_old_calls_reach_modern_overrides_of_legacy_providers(base, legacy, neutral):
+    calls = []
+
+    async def old(self, s3_url, destination_path):
+        raise AssertionError("the parent provider was bypassed by the modern override")
+
+    async def new(self, object_url, destination_path):
+        calls.append((object_url, destination_path))
+
+    parent = type("LegacySandbox", (base,), {legacy: old, "terminate": _terminate})
+    child = type("ModernSandbox", (parent,), {neutral: new})
+
+    with pytest.warns(DeprecationWarning, match=f"use {neutral}"):
+        await getattr(child(), legacy)(s3_url="file:///objects/input", destination_path="/app/input")
+
+    assert calls == [("file:///objects/input", "/app/input")]
+
+
+@pytest.mark.asyncio
+async def test_artifact_staging_uses_an_older_container_provider(tmp_path):
+    class LegacySandbox(Sandbox):
+        async def terminate(self):
+            pass
+
+        async def exec(self, *command):
+            Path(command[-1]).mkdir(parents=True, exist_ok=True)
+
+        async def write_file_from_s3(self, s3_url, destination_path):
+            Path(destination_path).write_bytes(b"artifact payload")
+
+    artifacts = {"input.json": SimpleNamespace(object_url="file:///objects/input")}
+
+    loaded = await stage_files_into_container(LegacySandbox(), artifacts, str(tmp_path))
+
+    assert loaded == {"input.json": str(tmp_path / "input.json")}
+    assert (tmp_path / "input.json").read_bytes() == b"artifact payload"
+
+
+@pytest.mark.asyncio
+async def test_legacy_vm_overrides_can_delegate_to_super(monkeypatch):
+    calls = []
+
+    class LegacyVm(VmSandbox):
+        async def terminate(self):
+            pass
+
+        async def load_s3_file(self, s3_url, destination_path):
+            calls.append("load override")
+            await super().load_s3_file(s3_url, destination_path)
+
+        async def write_file_from_s3(self, s3_url, destination_path):
+            calls.append("write override")
+            await super().write_file_from_s3(s3_url, destination_path)
+
+        async def _download_object_to_vm(self, object_url, vm_path):
+            calls.append((object_url, vm_path))
+
+        async def _copy_into_container(self, vm_path, destination_path):
+            calls.append((vm_path, destination_path))
+
+        async def _remove_vm_temp_file(self, *vm_paths):
+            calls.append(vm_paths)
+
+    monkeypatch.setattr(LegacyVm, "_staging_path", staticmethod(lambda kind, dest: "/tmp/staged"))
+
+    with pytest.warns(DeprecationWarning):
+        await LegacyVm().write_file_from_object("file:///objects/input", "/app/input")
+
+    assert calls == [
+        "write override", "load override", ("file:///objects/input", "/tmp/staged"),
+        ("/tmp/staged", "/app/input"), ("/tmp/staged",),
+    ]

From ade2141530cdbfa4095969a9170d8d0446a30836 Mon Sep 17 00:00:00 2001
From: morluto <76467478+morluto@users.noreply.github.com>
Date: Thu, 8 Oct 2026 02:24:34 +0800
Subject: [PATCH 4/4] test(lifecycle): update local object routing fixture

---
 tst/unit/providers/sandbox_providers/vm_sandbox_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tst/unit/providers/sandbox_providers/vm_sandbox_test.py b/tst/unit/providers/sandbox_providers/vm_sandbox_test.py
index 319ed3af..04017b4c 100644
--- a/tst/unit/providers/sandbox_providers/vm_sandbox_test.py
+++ b/tst/unit/providers/sandbox_providers/vm_sandbox_test.py
@@ -719,7 +719,7 @@ async def push(sandbox, store, object_url, vm_path):
         pushed.append((store, object_url))
 
     monkeypatch.setattr(sandbox_module, "push_object_over_exec", push)
-    sandbox = _RecordingVmSandbox(images_stdout="myimage\n")
+    sandbox = _RecordingVmSandbox()
 
     await sandbox.load_object_file(data, "/tmp/data.json")
     await sandbox.load_docker_images([SimpleNamespace(tar_gz_object_url=image, image_name="myimage:latest")])