From af4491991029ef872f5c16498ccb53732093f690 Mon Sep 17 00:00:00 2001 From: morluto <76467478+morluto@users.noreply.github.com> Date: Wed, 7 Oct 2026 09:29:51 +0800 Subject: [PATCH 1/5] fix(artifact): make publication retries reliable --- .../artifact/artifacts/docker_image.py | 82 ++++++-- src/agent_env/artifact/artifacts/file.py | 49 ++--- .../artifact/test_publication_reliability.py | 183 ++++++++++++++++++ 3 files changed, 260 insertions(+), 54 deletions(-) create mode 100644 tst/unit/artifact/test_publication_reliability.py diff --git a/src/agent_env/artifact/artifacts/docker_image.py b/src/agent_env/artifact/artifacts/docker_image.py index f2bb5e0d..e864d9c9 100644 --- a/src/agent_env/artifact/artifacts/docker_image.py +++ b/src/agent_env/artifact/artifacts/docker_image.py @@ -9,9 +9,12 @@ import posixpath import re import shlex +import signal import subprocess import tarfile import tempfile +import threading +import time import uuid from dataclasses import dataclass from importlib.metadata import version as pkg_version @@ -29,6 +32,7 @@ ProgressCallback = Callable[[str, str, int], None] _SAFE_GIT_NAME = re.compile(r"^[a-zA-Z0-9._-]+$") +_DOCKER_SAVE_CHUNK_SIZE = 1024 * 1024 def _validate_git_name(value: str, label: str) -> str: @@ -60,6 +64,67 @@ def _git_clone_commands(owner: str, repo: str, ref: str | None, token: str | Non f"trap 'rm -f /tmp/git-askpass.sh' EXIT; GIT_ASKPASS=/tmp/git-askpass.sh git clone --depth 1 {ref_flag} {clone_url} /tmp/repo", ] + +def _save_image_tar_gz(image_ref: str, output_path: Path, timeout_seconds: float) -> None: + """Stream docker save through gzip under one deadline without buffering the image.""" + deadline = time.monotonic() + timeout_seconds + with tempfile.TemporaryFile() as stderr_file: + process_options = {"start_new_session": True} if os.name == "posix" else {} + process = subprocess.Popen( + ["docker", "save", image_ref], + stdout=subprocess.PIPE, + stderr=stderr_file, + **process_options, + ) + timed_out = threading.Event() + + def kill_process() -> None: + timed_out.set() + if os.name == "posix": + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + elif process.poll() is None: + process.kill() + + watchdog = threading.Timer(max(deadline - time.monotonic(), 0), kill_process) + watchdog.daemon = True + watchdog.start() + succeeded = False + try: + assert process.stdout is not None + with gzip.open(output_path, "wb") as compressed: + while chunk := process.stdout.read(_DOCKER_SAVE_CHUNK_SIZE): + if timed_out.is_set(): + raise subprocess.TimeoutExpired(process.args, timeout_seconds) + compressed.write(chunk) + if time.monotonic() >= deadline: + raise subprocess.TimeoutExpired(process.args, timeout_seconds) + remaining = deadline - time.monotonic() + process.wait(timeout=max(remaining, 0)) + if timed_out.is_set(): + raise subprocess.TimeoutExpired(process.args, timeout_seconds) + if process.returncode != 0: + stderr_file.seek(0) + stderr = stderr_file.read().decode(errors="replace") + raise RuntimeError(f"docker save {image_ref} failed: {stderr}") + succeeded = True + except subprocess.TimeoutExpired as error: + raise RuntimeError( + f"docker save {image_ref} timed out after {timeout_seconds} seconds" + ) from error + finally: + watchdog.cancel() + watchdog.join() + if process.poll() is None: + kill_process() + process.wait() + if process.stdout is not None: + process.stdout.close() + if not succeeded: + output_path.unlink(missing_ok=True) + class DockerImageArtifact(Artifact): """A Docker image artifact stored as tar.gz in S3.""" @@ -102,22 +167,7 @@ def put( tmp_path = Path(tmp.name) try: - save_proc = subprocess.Popen( - ["docker", "save", image_ref], - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - with gzip.open(tmp_path, "wb") as gz_file: - while chunk := save_proc.stdout.read(8192): - gz_file.write(chunk) - try: - save_proc.wait(timeout=cls.DOCKER_SAVE_TIMEOUT_SECONDS) - except subprocess.TimeoutExpired: - save_proc.kill() - raise RuntimeError(f"docker save {image_ref} timed out after {cls.DOCKER_SAVE_TIMEOUT_SECONDS} seconds") - if save_proc.returncode != 0: - stderr = save_proc.stderr.read().decode() if save_proc.stderr else "" - raise RuntimeError(f"docker save {image_ref} failed: {stderr}") + _save_image_tar_gz(image_ref, tmp_path, cls.DOCKER_SAVE_TIMEOUT_SECONDS) tar_gz_object_url = objects.put_file_at( f"{prefix}{fs_safe(id)}-v{version}.tar.gz", str(tmp_path), "application/gzip" diff --git a/src/agent_env/artifact/artifacts/file.py b/src/agent_env/artifact/artifacts/file.py index 43698b8b..2eb91545 100644 --- a/src/agent_env/artifact/artifacts/file.py +++ b/src/agent_env/artifact/artifacts/file.py @@ -54,35 +54,12 @@ def put( store = get_artifact_store() filename = os.path.basename(file_path) - content_type, _ = mimetypes.guess_type(file_path) - if content_type is None: - content_type = "application/octet-stream" - - version = store.next_version(id) - - # Use put_object_file (boto3 managed multipart upload) instead of - # put_object (s3.put_object — hard 5 GB single-object limit). Streams - # directly from disk so large files don't get loaded into RAM either. - s3_url = store.put_object_file( - artifact_type="file", - id=id, - version=version, - object_name=filename, - file_path=file_path, - content_type=content_type, + prefix = store.attempt_prefix("file", id) + return cls.put_at( + id, description=description, file_path=file_path, + object_url=f"{prefix}{filename}", filename=filename, ) - # Store artifact document in MongoDB - instance = cls( - id=id, - version=version, - description=description, - filename=filename, - content_type=content_type, - s3_url=s3_url, - ) - return store.put_document(instance) - @classmethod def put_bytes( cls, @@ -97,24 +74,20 @@ def put_bytes( from agent_env.artifact.store import get_artifact_store store = get_artifact_store() - version = store.next_version(id) - - s3_url = store.put_object( - artifact_type="file", - id=id, - version=version, - object_name=filename, - data=content, - content_type=content_type, + prefix = store.attempt_prefix("file", id) + object_url = f"{prefix}{filename}" + object_store = get_config().get_object_store_to_write(object_url, id) + stored_url = object_store.put( + object_store.get_object_key(object_url), content, content_type ) instance = cls( id=id, - version=version, + version=store.next_version(id), description=description, filename=filename, content_type=content_type, - s3_url=s3_url, + s3_url=stored_url, ) return store.put_document(instance) diff --git a/tst/unit/artifact/test_publication_reliability.py b/tst/unit/artifact/test_publication_reliability.py new file mode 100644 index 00000000..5bec133c --- /dev/null +++ b/tst/unit/artifact/test_publication_reliability.py @@ -0,0 +1,183 @@ +import gzip +import os +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import pytest + +from agent_env.artifact.artifacts import docker_image as docker_image_module +from agent_env.artifact.artifacts.docker_image import _save_image_tar_gz +from agent_env.artifact.artifacts.file import FileArtifact + +_SAVE_TIMEOUT_SECONDS = 0.25 +_SAVE_SUCCESS_TIMEOUT_SECONDS = 30 +_SAVE_FAILURE_TIMEOUT_SECONDS = 3 +_FAKE_DOCKER_STDERR_BYTES = 1024 * 1024 +_STREAM_TEST_MIB = 16 +_MIB = 1024 * 1024 +_WATCHDOG_TEST_SECONDS = 2 +_RETRY_PAYLOAD_MULTIPLIER = 1000 + + +def _install_fake_docker(tmp_path: Path, monkeypatch, body: str) -> None: + executable = tmp_path / "docker" + executable.write_text(f"#!/usr/bin/env python3\n{body}\n") + executable.chmod(0o755) + monkeypatch.setenv("PATH", f"{tmp_path}{os.pathsep}{os.environ['PATH']}") + + +@pytest.mark.parametrize("api", ["path", "bytes"]) +def test_file_artifact_retry_after_document_failure_isolated(local_stores, tmp_path, api): + document_store = local_stores.get_document_store() + original_insert = document_store.insert + failed = False + + def insert(collection, document): + nonlocal failed + if not failed: + failed = True + raise RuntimeError("temporary document failure") + return original_insert(collection, document) + + document_store.insert = insert + payload = b"retry payload" + source = tmp_path / "payload.bin" + source.write_bytes(payload) + + def publish(): + if api == "path": + return FileArtifact.put("retry-file", description="retry", file_path=str(source)) + return FileArtifact.put_bytes( + "retry-file", description="retry", filename=source.name, content=payload + ) + + with pytest.raises(RuntimeError, match="temporary document failure"): + publish() + artifact = publish() + + assert artifact.version == 1 + assert artifact.load() == payload + objects = local_stores.get_object_store() + keys = objects.list("artifacts/file/") + assert len(keys) == 2 + assert len(set(keys)) == 2 + + +def test_concurrent_file_artifacts_keep_each_payload(local_stores): + payloads = [b"first" * _RETRY_PAYLOAD_MULTIPLIER, b"second" * _RETRY_PAYLOAD_MULTIPLIER] + + def publish(index): + return FileArtifact.put_bytes( + "concurrent-file", description=str(index), filename="data.bin", content=payloads[index] + ) + + with ThreadPoolExecutor(max_workers=2) as pool: + artifacts = list(pool.map(publish, range(2))) + + assert sorted(artifact.version for artifact in artifacts) == [1, 2] + assert {artifact.load() for artifact in artifacts} == set(payloads) + + +@pytest.mark.parametrize("write_stderr", [False, True]) +def test_docker_save_timeout_covers_stuck_stdout(local_stores, tmp_path, monkeypatch, write_stderr): + stderr_write = f"os.write(2, b'x' * {_FAKE_DOCKER_STDERR_BYTES})\n" if write_stderr else "" + _install_fake_docker( + tmp_path, + monkeypatch, + f"import os, time\nos.write(1, b'partial archive')\n{stderr_write}time.sleep(30)", + ) + archive = tmp_path / "image.tar.gz" + started = time.monotonic() + + with pytest.raises(RuntimeError, match="timed out"): + _save_image_tar_gz("example:latest", archive, _SAVE_TIMEOUT_SECONDS) + + assert time.monotonic() - started < _WATCHDOG_TEST_SECONDS + assert not archive.exists() + + +def test_docker_save_timeout_kills_descendant_holding_stdout(local_stores, tmp_path, monkeypatch): + child_script = "import time; time.sleep(30)" + _install_fake_docker( + tmp_path, + monkeypatch, + f"import subprocess, sys\nsubprocess.Popen([sys.executable, '-c', {child_script!r}])", + ) + archive = tmp_path / "image.tar.gz" + started = time.monotonic() + + with pytest.raises(RuntimeError, match="timed out"): + _save_image_tar_gz("example:latest", archive, _SAVE_TIMEOUT_SECONDS) + + assert time.monotonic() - started < _WATCHDOG_TEST_SECONDS + assert not archive.exists() + + +def test_docker_save_failure_removes_partial_archive(local_stores, tmp_path, monkeypatch): + _install_fake_docker( + tmp_path, + monkeypatch, + "import os\nos.write(1, b'partial archive')\nos.write(2, b'failure details')\nraise SystemExit(7)", + ) + archive = tmp_path / "image.tar.gz" + + with pytest.raises(RuntimeError, match="failure details"): + _save_image_tar_gz("example:latest", archive, _SAVE_FAILURE_TIMEOUT_SECONDS) + + assert not archive.exists() + + +def test_docker_save_compression_error_kills_child_and_cleans_archive(local_stores, tmp_path, monkeypatch): + _install_fake_docker(tmp_path, monkeypatch, "import time\ntime.sleep(30)") + archive = tmp_path / "image.tar.gz" + + def fail_compression(*args, **kwargs): + raise OSError("disk full") + + monkeypatch.setattr(docker_image_module.gzip, "open", fail_compression) + with pytest.raises(OSError, match="disk full"): + _save_image_tar_gz("example:latest", archive, _SAVE_FAILURE_TIMEOUT_SECONDS) + + assert not archive.exists() + + +def test_docker_save_streams_large_payload_into_valid_gzip(local_stores, tmp_path, monkeypatch): + _install_fake_docker( + tmp_path, + monkeypatch, + f"import os\nchunk = b'z' * {_MIB}\nfor _ in range({_STREAM_TEST_MIB}): os.write(1, chunk)", + ) + archive = tmp_path / "image.tar.gz" + _save_image_tar_gz("example:latest", archive, _SAVE_SUCCESS_TIMEOUT_SECONDS) + + with gzip.open(archive, "rb") as compressed: + payload = compressed.read() + assert len(payload) == _STREAM_TEST_MIB * _MIB + assert payload == b"z" * len(payload) + + +def test_docker_image_put_cleans_archive_after_save_failure(local_stores, tmp_path, monkeypatch): + _install_fake_docker(tmp_path, monkeypatch, "raise SystemExit(7)") + image_store = type("ImageStore", (), { + "image_ref": lambda self, repository, tag: f"registry/{repository}:{tag}", + "ensure_repository": lambda self, repository: None, + })() + monkeypatch.setattr(local_stores, "get_image_store_for", lambda artifact_id: image_store) + monkeypatch.setattr(docker_image_module, "_push_local_image", lambda *args: None) + temp_paths = [] + real_named_temp_file = docker_image_module.tempfile.NamedTemporaryFile + + def named_temp_file(*args, **kwargs): + file = real_named_temp_file(*args, **kwargs) + temp_paths.append(Path(file.name)) + return file + + monkeypatch.setattr(docker_image_module.tempfile, "NamedTemporaryFile", named_temp_file) + + with pytest.raises(RuntimeError, match="docker save"): + docker_image_module.DockerImageArtifact.put( + id="failed-image", description="failed", image_name="local:latest" + ) + + assert temp_paths and all(not path.exists() for path in temp_paths) From 36080ca21d9cdf0a1256f05ae1ddbbd537138bf6 Mon Sep 17 00:00:00 2001 From: morluto <76467478+morluto@users.noreply.github.com> Date: Wed, 7 Oct 2026 09:34:18 +0800 Subject: [PATCH 2/5] fix(artifact): reap failed save descendants --- .../artifact/artifacts/docker_image.py | 7 +++- .../artifact/test_publication_reliability.py | 39 ++++++++++++++++++ .../store/local_backends_composition_test.py | 41 +++++++++++++++++-- 3 files changed, 82 insertions(+), 5 deletions(-) diff --git a/src/agent_env/artifact/artifacts/docker_image.py b/src/agent_env/artifact/artifacts/docker_image.py index e864d9c9..02dd15ee 100644 --- a/src/agent_env/artifact/artifacts/docker_image.py +++ b/src/agent_env/artifact/artifacts/docker_image.py @@ -77,9 +77,13 @@ def _save_image_tar_gz(image_ref: str, output_path: Path, timeout_seconds: float **process_options, ) timed_out = threading.Event() + kill_sent = False def kill_process() -> None: + nonlocal kill_sent timed_out.set() + if kill_sent: + return if os.name == "posix": try: os.killpg(process.pid, signal.SIGKILL) @@ -87,6 +91,7 @@ def kill_process() -> None: pass elif process.poll() is None: process.kill() + kill_sent = True watchdog = threading.Timer(max(deadline - time.monotonic(), 0), kill_process) watchdog.daemon = True @@ -117,7 +122,7 @@ def kill_process() -> None: finally: watchdog.cancel() watchdog.join() - if process.poll() is None: + if not succeeded or process.poll() is None: kill_process() process.wait() if process.stdout is not None: diff --git a/tst/unit/artifact/test_publication_reliability.py b/tst/unit/artifact/test_publication_reliability.py index 5bec133c..c13ccb70 100644 --- a/tst/unit/artifact/test_publication_reliability.py +++ b/tst/unit/artifact/test_publication_reliability.py @@ -18,6 +18,8 @@ _MIB = 1024 * 1024 _WATCHDOG_TEST_SECONDS = 2 _RETRY_PAYLOAD_MULTIPLIER = 1000 +_DESCENDANT_MARKER_DELAY_SECONDS = 0.5 +_DESCENDANT_SETTLE_SECONDS = 0.6 def _install_fake_docker(tmp_path: Path, monkeypatch, body: str) -> None: @@ -142,6 +144,43 @@ def fail_compression(*args, **kwargs): assert not archive.exists() +def test_docker_save_compression_error_kills_descendant_after_leader_exit( + local_stores, tmp_path, monkeypatch +): + marker = tmp_path / "descendant-survived" + child_script = ( + f"import time; time.sleep({_DESCENDANT_MARKER_DELAY_SECONDS}); " + f"open({str(marker)!r}, 'w').close()" + ) + _install_fake_docker( + tmp_path, + monkeypatch, + f"import subprocess, sys\nsubprocess.Popen([sys.executable, '-c', {child_script!r}])", + ) + archive = tmp_path / "image.tar.gz" + launched_processes = [] + real_popen = docker_image_module.subprocess.Popen + + def recording_popen(*args, **kwargs): + process = real_popen(*args, **kwargs) + launched_processes.append(process) + return process + + monkeypatch.setattr(docker_image_module.subprocess, "Popen", recording_popen) + + def fail_after_leader_exit(*args, **kwargs): + launched_processes[0].wait(timeout=_SAVE_FAILURE_TIMEOUT_SECONDS) + raise OSError("disk full") + + monkeypatch.setattr(docker_image_module.gzip, "open", fail_after_leader_exit) + with pytest.raises(OSError, match="disk full"): + _save_image_tar_gz("example:latest", archive, _SAVE_FAILURE_TIMEOUT_SECONDS) + + time.sleep(_DESCENDANT_SETTLE_SECONDS) + assert not marker.exists() + assert not archive.exists() + + def test_docker_save_streams_large_payload_into_valid_gzip(local_stores, tmp_path, monkeypatch): _install_fake_docker( tmp_path, diff --git a/tst/unit/store/local_backends_composition_test.py b/tst/unit/store/local_backends_composition_test.py index 12966764..e5b0d109 100644 --- a/tst/unit/store/local_backends_composition_test.py +++ b/tst/unit/store/local_backends_composition_test.py @@ -20,6 +20,7 @@ from agent_env.artifact.artifacts.file import FileArtifact from agent_env.artifact.artifacts.file_artifact_universe import FileArtifactUniverse from agent_env.artifact.artifacts.skill import SkillArtifact, download_skill +from agent_env.artifact.store import get_artifact_store from agent_env.cli.artifact.file_artifact_universe import file_artifact_universe from agent_env.store.ids import fs_safe, key_segment from agent_env.config import get_config, set_image_store, set_object_store @@ -179,15 +180,39 @@ def test_a_local_id_file_artifact_lands_under_its_encoded_segment(local_stores, fb = FileArtifact.put_bytes(id=HOSTILE, description="d", filename="raw.bin", content=b"raw") store = local_stores.get_object_store() - assert store.get_object_key(fa.object_url) == f"artifacts/file/{HOSTILE_SEGMENT}/1/payload.json" - assert store.get_object_key(fb.object_url) == f"artifacts/file/{HOSTILE_SEGMENT}/2/raw.bin" + assert re.fullmatch( + rf"artifacts/file/{re.escape(HOSTILE_SEGMENT)}/1-[0-9a-f]{{8}}/payload.json", + store.get_object_key(fa.object_url), + ) + assert re.fullmatch( + rf"artifacts/file/{re.escape(HOSTILE_SEGMENT)}/2-[0-9a-f]{{8}}/raw.bin", + store.get_object_key(fb.object_url), + ) assert FileArtifact.get(HOSTILE, 1).load() == b'{"hostile": true}' assert FileArtifact.get(HOSTILE).load() == b"raw" -def test_a_legacy_id_keeps_its_object_key_byte_identical(local_stores, tmp_path): +def test_a_legacy_id_keeps_its_encoded_segment_in_the_object_key(local_stores, tmp_path): fa = FileArtifact.put(id="Legacy/Id v1", description="d", file_path=_write(tmp_path, "p.txt", b"x")) - assert local_stores.get_object_store().get_object_key(fa.object_url) == "artifacts/file/Legacy/Id v1/1/p.txt" + key = local_stores.get_object_store().get_object_key(fa.object_url) + assert re.fullmatch(r"artifacts/file/Legacy/Id v1/1-[0-9a-f]{8}/p.txt", key) + assert fa.load() == b"x" + + +def test_a_published_file_with_an_existing_version_locator_still_loads(local_stores, tmp_path): + objects = local_stores.get_object_store() + old_locator = objects.put_file( + "artifacts/file/Legacy/Id v1/1/p.txt", _write(tmp_path, "old.txt", b"published before attempt prefixes") + ) + stored = get_artifact_store().put_document( + FileArtifact( + id="old-file", version=1, description="old", filename="p.txt", + content_type="text/plain", s3_url=old_locator, + ) + ) + + assert stored.object_url == old_locator + assert FileArtifact.get("old-file").load() == b"published before attempt prefixes" def _cli_dir(tmp_path, name, files): @@ -222,10 +247,18 @@ def __init__(self, *args, **kwargs): self.stdout = io.BytesIO(b"image-tar-bytes") self.stderr = io.BytesIO(b"") self.returncode = 0 + self.args = args[0] if args else [] def wait(self, timeout=None): + self.returncode = 0 return 0 + def poll(self): + return self.returncode + + def kill(self): + self.returncode = -9 + @pytest.mark.parametrize("entity_id, registry, repository, tarball", [ (HOSTILE, "localhost:5000", HOSTILE_SEGMENT, From 7956ba53e3e78d7b662a43a6df57eeee084c091e Mon Sep 17 00:00:00 2001 From: morluto <76467478+morluto@users.noreply.github.com> Date: Wed, 7 Oct 2026 09:46:03 +0800 Subject: [PATCH 3/5] fix(artifact): bound docker save stderr --- .../artifact/artifacts/docker_image.py | 146 +++++++++++------- .../artifact/test_publication_reliability.py | 74 +++++++-- 2 files changed, 146 insertions(+), 74 deletions(-) diff --git a/src/agent_env/artifact/artifacts/docker_image.py b/src/agent_env/artifact/artifacts/docker_image.py index 02dd15ee..0c792152 100644 --- a/src/agent_env/artifact/artifacts/docker_image.py +++ b/src/agent_env/artifact/artifacts/docker_image.py @@ -33,6 +33,8 @@ _SAFE_GIT_NAME = re.compile(r"^[a-zA-Z0-9._-]+$") _DOCKER_SAVE_CHUNK_SIZE = 1024 * 1024 +_DOCKER_SAVE_STDERR_CHUNK_SIZE = 8192 +_DOCKER_SAVE_STDERR_LIMIT_BYTES = 16 * 1024 def _validate_git_name(value: str, label: str) -> str: @@ -68,67 +70,93 @@ def _git_clone_commands(owner: str, repo: str, ref: str | None, token: str | Non def _save_image_tar_gz(image_ref: str, output_path: Path, timeout_seconds: float) -> None: """Stream docker save through gzip under one deadline without buffering the image.""" deadline = time.monotonic() + timeout_seconds - with tempfile.TemporaryFile() as stderr_file: - process_options = {"start_new_session": True} if os.name == "posix" else {} - process = subprocess.Popen( - ["docker", "save", image_ref], - stdout=subprocess.PIPE, - stderr=stderr_file, - **process_options, - ) - timed_out = threading.Event() - kill_sent = False - - def kill_process() -> None: - nonlocal kill_sent - timed_out.set() - if kill_sent: - return - if os.name == "posix": - try: - os.killpg(process.pid, signal.SIGKILL) - except ProcessLookupError: - pass - elif process.poll() is None: - process.kill() - kill_sent = True - - watchdog = threading.Timer(max(deadline - time.monotonic(), 0), kill_process) - watchdog.daemon = True - watchdog.start() - succeeded = False + process_options = {"start_new_session": True} if os.name == "posix" else {} + process = subprocess.Popen( + ["docker", "save", image_ref], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + **process_options, + ) + timed_out = threading.Event() + kill_sent = False + stderr_tail = bytearray() + stderr_size = 0 + stderr_errors: list[BaseException] = [] + + def kill_process_group() -> None: + nonlocal kill_sent + timed_out.set() + if kill_sent: + return + if os.name == "posix": + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + elif process.returncode is None: + process.kill() + kill_sent = True + + def drain_stderr() -> None: + nonlocal stderr_size try: - assert process.stdout is not None - with gzip.open(output_path, "wb") as compressed: - while chunk := process.stdout.read(_DOCKER_SAVE_CHUNK_SIZE): - if timed_out.is_set(): - raise subprocess.TimeoutExpired(process.args, timeout_seconds) - compressed.write(chunk) - if time.monotonic() >= deadline: - raise subprocess.TimeoutExpired(process.args, timeout_seconds) - remaining = deadline - time.monotonic() - process.wait(timeout=max(remaining, 0)) - if timed_out.is_set(): - raise subprocess.TimeoutExpired(process.args, timeout_seconds) - if process.returncode != 0: - stderr_file.seek(0) - stderr = stderr_file.read().decode(errors="replace") - raise RuntimeError(f"docker save {image_ref} failed: {stderr}") - succeeded = True - except subprocess.TimeoutExpired as error: - raise RuntimeError( - f"docker save {image_ref} timed out after {timeout_seconds} seconds" - ) from error - finally: - watchdog.cancel() - watchdog.join() - if not succeeded or process.poll() is None: - kill_process() + assert process.stderr is not None + while chunk := process.stderr.read(_DOCKER_SAVE_STDERR_CHUNK_SIZE): + stderr_size += len(chunk) + stderr_tail.extend(chunk) + if len(stderr_tail) > _DOCKER_SAVE_STDERR_LIMIT_BYTES: + del stderr_tail[:-_DOCKER_SAVE_STDERR_LIMIT_BYTES] + except Exception as error: + stderr_errors.append(error) + + stderr_thread = threading.Thread(target=drain_stderr, daemon=True) + stderr_thread.start() + watchdog = threading.Timer(max(deadline - time.monotonic(), 0), kill_process_group) + watchdog.daemon = True + watchdog.start() + succeeded = False + try: + assert process.stdout is not None + with gzip.open(output_path, "wb") as compressed: + while chunk := process.stdout.read(_DOCKER_SAVE_CHUNK_SIZE): + if timed_out.is_set(): + raise subprocess.TimeoutExpired(process.args, timeout_seconds) + compressed.write(chunk) + stderr_thread.join(timeout=max(deadline - time.monotonic(), 0)) + if stderr_thread.is_alive() or timed_out.is_set(): + raise subprocess.TimeoutExpired(process.args, timeout_seconds) + if stderr_errors: + raise RuntimeError(f"docker save {image_ref} stderr could not be read") from stderr_errors[0] + + watchdog.cancel() + watchdog.join() + if timed_out.is_set() or time.monotonic() >= deadline: + raise subprocess.TimeoutExpired(process.args, timeout_seconds) + process.wait(timeout=max(deadline - time.monotonic(), 0)) + if process.returncode != 0: + stderr = stderr_tail.decode(errors="replace") + if stderr_size > _DOCKER_SAVE_STDERR_LIMIT_BYTES: + stderr = f"[stderr truncated; showing last {_DOCKER_SAVE_STDERR_LIMIT_BYTES} bytes]\n{stderr}" + raise RuntimeError(f"docker save {image_ref} failed: {stderr}") + succeeded = True + except subprocess.TimeoutExpired as error: + raise RuntimeError( + f"docker save {image_ref} timed out after {timeout_seconds} seconds" + ) from error + finally: + watchdog.cancel() + watchdog.join() + if process.returncode is None and not succeeded: + kill_process_group() + if process.returncode is None: process.wait() - if process.stdout is not None: - process.stdout.close() - if not succeeded: - output_path.unlink(missing_ok=True) + stderr_thread.join() + if process.stdout is not None: + process.stdout.close() + if process.stderr is not None: + process.stderr.close() + if not succeeded: + output_path.unlink(missing_ok=True) class DockerImageArtifact(Artifact): """A Docker image artifact stored as tar.gz in S3.""" diff --git a/tst/unit/artifact/test_publication_reliability.py b/tst/unit/artifact/test_publication_reliability.py index c13ccb70..41106fce 100644 --- a/tst/unit/artifact/test_publication_reliability.py +++ b/tst/unit/artifact/test_publication_reliability.py @@ -14,12 +14,15 @@ _SAVE_SUCCESS_TIMEOUT_SECONDS = 30 _SAVE_FAILURE_TIMEOUT_SECONDS = 3 _FAKE_DOCKER_STDERR_BYTES = 1024 * 1024 +_LARGE_STDERR_BYTES = 2 * 1024 * 1024 _STREAM_TEST_MIB = 16 _MIB = 1024 * 1024 _WATCHDOG_TEST_SECONDS = 2 _RETRY_PAYLOAD_MULTIPLIER = 1000 _DESCENDANT_MARKER_DELAY_SECONDS = 0.5 _DESCENDANT_SETTLE_SECONDS = 0.6 +_LEADER_EXIT_POLL_SECONDS = 0.01 +_MAX_DOCKER_ERROR_CHARS = 20 * 1024 def _install_fake_docker(tmp_path: Path, monkeypatch, body: str) -> None: @@ -130,6 +133,45 @@ def test_docker_save_failure_removes_partial_archive(local_stores, tmp_path, mon assert not archive.exists() +def test_docker_save_bounds_captured_stderr_for_large_failure(local_stores, tmp_path, monkeypatch): + _install_fake_docker( + tmp_path, + monkeypatch, + "import os\n" + "os.write(2, b'begin')\n" + f"os.write(2, b'x' * {_LARGE_STDERR_BYTES})\n" + "os.write(2, b'end-marker')\n" + "raise SystemExit(7)", + ) + archive = tmp_path / "image.tar.gz" + + with pytest.raises(RuntimeError, match="stderr truncated") as raised: + _save_image_tar_gz("example:latest", archive, _SAVE_SUCCESS_TIMEOUT_SECONDS) + + assert "end-marker" in str(raised.value) + assert "begin" not in str(raised.value) + assert len(str(raised.value)) < _MAX_DOCKER_ERROR_CHARS + assert not archive.exists() + + +def test_docker_save_does_not_signal_group_after_leader_is_reaped(local_stores, tmp_path, monkeypatch): + _install_fake_docker(tmp_path, monkeypatch, "raise SystemExit(7)") + archive = tmp_path / "image.tar.gz" + signaled_groups = [] + real_killpg = docker_image_module.os.killpg + + def track_killpg(group_id, signal_number): + signaled_groups.append(group_id) + return real_killpg(group_id, signal_number) + + monkeypatch.setattr(docker_image_module.os, "killpg", track_killpg) + with pytest.raises(RuntimeError, match="docker save"): + _save_image_tar_gz("example:latest", archive, _SAVE_SUCCESS_TIMEOUT_SECONDS) + + assert signaled_groups == [] + assert not archive.exists() + + def test_docker_save_compression_error_kills_child_and_cleans_archive(local_stores, tmp_path, monkeypatch): _install_fake_docker(tmp_path, monkeypatch, "import time\ntime.sleep(30)") archive = tmp_path / "image.tar.gz" @@ -147,29 +189,30 @@ def fail_compression(*args, **kwargs): def test_docker_save_compression_error_kills_descendant_after_leader_exit( local_stores, tmp_path, monkeypatch ): - marker = tmp_path / "descendant-survived" + leader_exit_marker = tmp_path / "leader-exited" + survival_marker = tmp_path / "descendant-survived" child_script = ( - f"import time; time.sleep({_DESCENDANT_MARKER_DELAY_SECONDS}); " - f"open({str(marker)!r}, 'w').close()" + "import os, sys, time\n" + "parent = int(sys.argv[1])\n" + "while os.getppid() == parent: time.sleep(0.01)\n" + "open(sys.argv[2], 'w').close()\n" + f"time.sleep({_DESCENDANT_MARKER_DELAY_SECONDS})\n" + "open(sys.argv[3], 'w').close()\n" ) _install_fake_docker( tmp_path, monkeypatch, - f"import subprocess, sys\nsubprocess.Popen([sys.executable, '-c', {child_script!r}])", + "import os, subprocess, sys\n" + f"subprocess.Popen([sys.executable, '-c', {child_script!r}, str(os.getpid()), " + f"{str(leader_exit_marker)!r}, {str(survival_marker)!r}])", ) archive = tmp_path / "image.tar.gz" - launched_processes = [] - real_popen = docker_image_module.subprocess.Popen - - def recording_popen(*args, **kwargs): - process = real_popen(*args, **kwargs) - launched_processes.append(process) - return process - - monkeypatch.setattr(docker_image_module.subprocess, "Popen", recording_popen) def fail_after_leader_exit(*args, **kwargs): - launched_processes[0].wait(timeout=_SAVE_FAILURE_TIMEOUT_SECONDS) + marker_deadline = time.monotonic() + _SAVE_FAILURE_TIMEOUT_SECONDS + while not leader_exit_marker.exists() and time.monotonic() < marker_deadline: + time.sleep(_LEADER_EXIT_POLL_SECONDS) + assert leader_exit_marker.exists() raise OSError("disk full") monkeypatch.setattr(docker_image_module.gzip, "open", fail_after_leader_exit) @@ -177,7 +220,8 @@ def fail_after_leader_exit(*args, **kwargs): _save_image_tar_gz("example:latest", archive, _SAVE_FAILURE_TIMEOUT_SECONDS) time.sleep(_DESCENDANT_SETTLE_SECONDS) - assert not marker.exists() + assert leader_exit_marker.exists() + assert not survival_marker.exists() assert not archive.exists() From f025960e5a2e9bcccbcf9c7cf500a414564f186b Mon Sep 17 00:00:00 2001 From: morluto <76467478+morluto@users.noreply.github.com> Date: Wed, 7 Oct 2026 10:03:07 +0800 Subject: [PATCH 4/5] fix(artifact): enforce deadlines for inherited pipes --- .../artifact/artifacts/docker_image.py | 100 +++++++----------- .../artifact/test_publication_reliability.py | 53 ++++++++++ .../store/local_backends_composition_test.py | 25 +---- 3 files changed, 99 insertions(+), 79 deletions(-) diff --git a/src/agent_env/artifact/artifacts/docker_image.py b/src/agent_env/artifact/artifacts/docker_image.py index 0c792152..ae0664a8 100644 --- a/src/agent_env/artifact/artifacts/docker_image.py +++ b/src/agent_env/artifact/artifacts/docker_image.py @@ -8,12 +8,12 @@ import os import posixpath import re +import selectors import shlex import signal import subprocess import tarfile import tempfile -import threading import time import uuid from dataclasses import dataclass @@ -70,69 +70,49 @@ def _git_clone_commands(owner: str, repo: str, ref: str | None, token: str | Non def _save_image_tar_gz(image_ref: str, output_path: Path, timeout_seconds: float) -> None: """Stream docker save through gzip under one deadline without buffering the image.""" deadline = time.monotonic() + timeout_seconds - process_options = {"start_new_session": True} if os.name == "posix" else {} process = subprocess.Popen( ["docker", "save", image_ref], stdout=subprocess.PIPE, stderr=subprocess.PIPE, - **process_options, + bufsize=0, + start_new_session=True, ) - timed_out = threading.Event() - kill_sent = False stderr_tail = bytearray() stderr_size = 0 - stderr_errors: list[BaseException] = [] - - def kill_process_group() -> None: - nonlocal kill_sent - timed_out.set() - if kill_sent: - return - if os.name == "posix": - try: - os.killpg(process.pid, signal.SIGKILL) - except ProcessLookupError: - pass - elif process.returncode is None: - process.kill() - kill_sent = True - - def drain_stderr() -> None: - nonlocal stderr_size - try: - assert process.stderr is not None - while chunk := process.stderr.read(_DOCKER_SAVE_STDERR_CHUNK_SIZE): - stderr_size += len(chunk) - stderr_tail.extend(chunk) - if len(stderr_tail) > _DOCKER_SAVE_STDERR_LIMIT_BYTES: - del stderr_tail[:-_DOCKER_SAVE_STDERR_LIMIT_BYTES] - except Exception as error: - stderr_errors.append(error) - - stderr_thread = threading.Thread(target=drain_stderr, daemon=True) - stderr_thread.start() - watchdog = threading.Timer(max(deadline - time.monotonic(), 0), kill_process_group) - watchdog.daemon = True - watchdog.start() succeeded = False try: assert process.stdout is not None - with gzip.open(output_path, "wb") as compressed: - while chunk := process.stdout.read(_DOCKER_SAVE_CHUNK_SIZE): - if timed_out.is_set(): - raise subprocess.TimeoutExpired(process.args, timeout_seconds) - compressed.write(chunk) - stderr_thread.join(timeout=max(deadline - time.monotonic(), 0)) - if stderr_thread.is_alive() or timed_out.is_set(): - raise subprocess.TimeoutExpired(process.args, timeout_seconds) - if stderr_errors: - raise RuntimeError(f"docker save {image_ref} stderr could not be read") from stderr_errors[0] - - watchdog.cancel() - watchdog.join() - if timed_out.is_set() or time.monotonic() >= deadline: + assert process.stderr is not None + with selectors.DefaultSelector() as selector: + for stream, name in ((process.stdout, "stdout"), (process.stderr, "stderr")): + os.set_blocking(stream.fileno(), False) + selector.register(stream, selectors.EVENT_READ, name) + with gzip.open(output_path, "wb") as compressed: + while selector.get_map(): + remaining = deadline - time.monotonic() + if remaining <= 0: + raise subprocess.TimeoutExpired(process.args, timeout_seconds) + for key, _ in selector.select(remaining): + chunk_size = ( + _DOCKER_SAVE_CHUNK_SIZE + if key.data == "stdout" + else _DOCKER_SAVE_STDERR_CHUNK_SIZE + ) + chunk = os.read(key.fd, chunk_size) + if not chunk: + selector.unregister(key.fileobj) + elif key.data == "stdout": + compressed.write(chunk) + else: + stderr_size += len(chunk) + stderr_tail.extend(chunk) + if len(stderr_tail) > _DOCKER_SAVE_STDERR_LIMIT_BYTES: + del stderr_tail[:-_DOCKER_SAVE_STDERR_LIMIT_BYTES] + + remaining = deadline - time.monotonic() + if remaining <= 0: raise subprocess.TimeoutExpired(process.args, timeout_seconds) - process.wait(timeout=max(deadline - time.monotonic(), 0)) + process.wait(timeout=remaining) if process.returncode != 0: stderr = stderr_tail.decode(errors="replace") if stderr_size > _DOCKER_SAVE_STDERR_LIMIT_BYTES: @@ -144,17 +124,19 @@ def drain_stderr() -> None: f"docker save {image_ref} timed out after {timeout_seconds} seconds" ) from error finally: - watchdog.cancel() - watchdog.join() - if process.returncode is None and not succeeded: - kill_process_group() if process.returncode is None: - process.wait() - stderr_thread.join() + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + except PermissionError: + process.kill() if process.stdout is not None: process.stdout.close() if process.stderr is not None: process.stderr.close() + if process.returncode is None: + process.wait() if not succeeded: output_path.unlink(missing_ok=True) diff --git a/tst/unit/artifact/test_publication_reliability.py b/tst/unit/artifact/test_publication_reliability.py index 41106fce..064acd0e 100644 --- a/tst/unit/artifact/test_publication_reliability.py +++ b/tst/unit/artifact/test_publication_reliability.py @@ -1,5 +1,6 @@ import gzip import os +import threading import time from concurrent.futures import ThreadPoolExecutor from pathlib import Path @@ -23,6 +24,10 @@ _DESCENDANT_SETTLE_SECONDS = 0.6 _LEADER_EXIT_POLL_SECONDS = 0.01 _MAX_DOCKER_ERROR_CHARS = 20 * 1024 +_DESCENDANT_RELEASE_SECONDS = 3 +_DESCENDANT_AUTO_EXIT_SECONDS = 5 +_ESCAPED_DESCENDANT_TIMEOUT_SECONDS = 1.5 +_ESCAPED_DESCENDANT_TEST_TIMEOUT_SECONDS = 4 def _install_fake_docker(tmp_path: Path, monkeypatch, body: str) -> None: @@ -32,6 +37,13 @@ def _install_fake_docker(tmp_path: Path, monkeypatch, body: str) -> None: monkeypatch.setenv("PATH", f"{tmp_path}{os.pathsep}{os.environ['PATH']}") +@pytest.fixture +def detached_child_release(tmp_path, request): + release = tmp_path / "release-descendant" + request.addfinalizer(lambda: release.touch()) + return release + + @pytest.mark.parametrize("api", ["path", "bytes"]) def test_file_artifact_retry_after_document_failure_isolated(local_stores, tmp_path, api): document_store = local_stores.get_document_store() @@ -119,6 +131,47 @@ def test_docker_save_timeout_kills_descendant_holding_stdout(local_stores, tmp_p assert not archive.exists() +@pytest.mark.parametrize("held_stream", ["stdout", "stderr"]) +@pytest.mark.timeout(_ESCAPED_DESCENDANT_TEST_TIMEOUT_SECONDS) +def test_docker_save_timeout_returns_when_escaped_descendant_holds_pipe( + local_stores, tmp_path, monkeypatch, detached_child_release, held_stream +): + child_pid_file = tmp_path / "detached-child.pid" + child_exit_file = tmp_path / "detached-child-exited" + child_script = ( + "import pathlib, sys, time\n" + f"release = pathlib.Path({str(detached_child_release)!r})\n" + f"deadline = time.monotonic() + {_DESCENDANT_AUTO_EXIT_SECONDS}\n" + f"while not release.exists() and time.monotonic() < deadline: time.sleep({_LEADER_EXIT_POLL_SECONDS})\n" + f"pathlib.Path({str(child_exit_file)!r}).touch()\n" + ) + pipe_args = "stdout=None, stderr=subprocess.DEVNULL" if held_stream == "stdout" else "stdout=subprocess.DEVNULL, stderr=None" + _install_fake_docker( + tmp_path, + monkeypatch, + "import pathlib, subprocess, sys\n" + f"child = subprocess.Popen([sys.executable, '-c', {child_script!r}], " + f"start_new_session=True, {pipe_args})\n" + f"pathlib.Path({str(child_pid_file)!r}).write_text(str(child.pid))", + ) + archive = tmp_path / "image.tar.gz" + thread_ids_before = {thread.ident for thread in threading.enumerate()} + started = time.monotonic() + + with pytest.raises(RuntimeError, match="timed out"): + _save_image_tar_gz("example:latest", archive, _ESCAPED_DESCENDANT_TIMEOUT_SECONDS) + + assert time.monotonic() - started < _ESCAPED_DESCENDANT_TEST_TIMEOUT_SECONDS + assert child_pid_file.exists() + assert not archive.exists() + detached_child_release.touch() + exit_deadline = time.monotonic() + _DESCENDANT_RELEASE_SECONDS + while not child_exit_file.exists() and time.monotonic() < exit_deadline: + time.sleep(_LEADER_EXIT_POLL_SECONDS) + assert child_exit_file.exists() + assert {thread.ident for thread in threading.enumerate()} <= thread_ids_before + + def test_docker_save_failure_removes_partial_archive(local_stores, tmp_path, monkeypatch): _install_fake_docker( tmp_path, diff --git a/tst/unit/store/local_backends_composition_test.py b/tst/unit/store/local_backends_composition_test.py index e5b0d109..afeab6ac 100644 --- a/tst/unit/store/local_backends_composition_test.py +++ b/tst/unit/store/local_backends_composition_test.py @@ -8,7 +8,6 @@ """ import gzip -import io import re from pathlib import Path @@ -242,22 +241,9 @@ def test_a_local_id_bundle_never_lists_another_ids_files(local_stores, cli_routi assert sorted(p.name for p in out.rglob("*") if p.is_file()) == ["a"] -class _DockerSave: - def __init__(self, *args, **kwargs): - self.stdout = io.BytesIO(b"image-tar-bytes") - self.stderr = io.BytesIO(b"") - self.returncode = 0 - self.args = args[0] if args else [] - - def wait(self, timeout=None): - self.returncode = 0 - return 0 - - def poll(self): - return self.returncode - - def kill(self): - self.returncode = -9 +def _save_fake_docker_image(image_ref, output_path, timeout_seconds): + with gzip.open(output_path, "wb") as image_archive: + image_archive.write(b"image-tar-bytes") @pytest.mark.parametrize("entity_id, registry, repository, tarball", [ @@ -275,7 +261,7 @@ def test_a_docker_image_names_its_repository_and_tarball_from_the_encoded_id( monkeypatch.setattr(LocalRegistryImageStore, "ensure_repository", lambda self, repository: local_registry.append(repository)) pushed = [] monkeypatch.setattr(docker_image, "_push_local_image", lambda src, ref, store: pushed.append(ref)) - monkeypatch.setattr(docker_image.subprocess, "Popen", _DockerSave) + monkeypatch.setattr(docker_image, "_save_image_tar_gz", _save_fake_docker_image) art = docker_image.DockerImageArtifact.put(id=entity_id, description="d", image_name="src:latest") @@ -305,7 +291,6 @@ def test_an_local_id_outside_the_cli_is_refused_before_any_image_or_object_is_wr set_image_store(images) pushed = [] monkeypatch.setattr(docker_image, "_push_local_image", lambda src, ref, store: pushed.append(ref)) - monkeypatch.setattr(docker_image.subprocess, "Popen", _DockerSave) with pytest.raises(ValueError, match="only the @local namespace's store holds"): put(tmp_path) @@ -354,7 +339,7 @@ def interrupted(**kwargs): def test_a_docker_image_put_that_stops_before_its_document_doesnt_block_the_next(local_stores, monkeypatch): set_image_store(FakeImageStore()) monkeypatch.setattr(docker_image, "_push_local_image", lambda src, ref, store: None) - monkeypatch.setattr(docker_image.subprocess, "Popen", _DockerSave) + monkeypatch.setattr(docker_image, "_save_image_tar_gz", _save_fake_docker_image) put_tar = docker_image.DockerImageArtifact.put_tar def interrupted(*args, **kwargs): From c14168fd650007a30056d4f00f102d5fa574cc9c 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 5/5] 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 86985046..f7fc1df8 100644 --- a/src/agent_env/providers/sandbox_providers/sandbox.py +++ b/src/agent_env/providers/sandbox_providers/sandbox.py @@ -134,6 +134,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.""" @@ -150,6 +165,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. @@ -182,7 +211,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.""" @@ -341,7 +370,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",), + ]