diff --git a/src/agent_env/artifact/artifacts/docker_image.py b/src/agent_env/artifact/artifacts/docker_image.py index 8a77cec..d6be073 100644 --- a/src/agent_env/artifact/artifacts/docker_image.py +++ b/src/agent_env/artifact/artifacts/docker_image.py @@ -8,10 +8,13 @@ import os import posixpath import re +import selectors import shlex +import signal import subprocess import tarfile import tempfile +import time import uuid from dataclasses import dataclass from importlib.metadata import version as pkg_version @@ -31,6 +34,9 @@ ProgressCallback = Callable[[str, str, int], None] _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: @@ -62,6 +68,80 @@ 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 + process = subprocess.Popen( + ["docker", "save", image_ref], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + bufsize=0, + start_new_session=True, + ) + stderr_tail = bytearray() + stderr_size = 0 + succeeded = False + try: + assert process.stdout is not None + 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=remaining) + 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: + if process.returncode is None: + 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) + class DockerImageArtifact(Artifact): """A Docker image artifact stored as tar.gz in the object store.""" @@ -115,22 +195,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 900339d..b4b4fb0 100644 --- a/src/agent_env/artifact/artifacts/file.py +++ b/src/agent_env/artifact/artifacts/file.py @@ -64,35 +64,11 @@ 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. - stored_url = store.put_object_file( - artifact_type="file", - id=id, - version=version, - object_name=filename, - file_path=file_path, - content_type=content_type, - ) - - # Store artifact document in MongoDB - instance = cls( - id=id, - version=version, - description=description, - filename=filename, - content_type=content_type, - object_url=stored_url, + prefix = store.attempt_prefix("file", id) + return cls.put_at( + id, description=description, file_path=file_path, + object_url=f"{prefix}{filename}", filename=filename, ) - return store.put_document(instance) - @classmethod def put_bytes( cls, @@ -107,20 +83,16 @@ def put_bytes( from agent_env.artifact.store import get_artifact_store store = get_artifact_store() - version = store.next_version(id) - - stored_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, diff --git a/src/agent_env/providers/sandbox_providers/sandbox.py b/src/agent_env/providers/sandbox_providers/sandbox.py index 8698504..f7fc1df 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/artifact/test_publication_reliability.py b/tst/unit/artifact/test_publication_reliability.py new file mode 100644 index 0000000..064acd0 --- /dev/null +++ b/tst/unit/artifact/test_publication_reliability.py @@ -0,0 +1,319 @@ +import gzip +import os +import threading +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 +_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 +_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: + 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.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() + 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() + + +@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, + 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_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" + + 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_compression_error_kills_descendant_after_leader_exit( + local_stores, tmp_path, monkeypatch +): + leader_exit_marker = tmp_path / "leader-exited" + survival_marker = tmp_path / "descendant-survived" + child_script = ( + "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, + "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" + + def fail_after_leader_exit(*args, **kwargs): + 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) + with pytest.raises(OSError, match="disk full"): + _save_image_tar_gz("example:latest", archive, _SAVE_FAILURE_TIMEOUT_SECONDS) + + time.sleep(_DESCENDANT_SETTLE_SECONDS) + assert leader_exit_marker.exists() + assert not survival_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, + 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) 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 0000000..daffb84 --- /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",), + ] diff --git a/tst/unit/store/local_backends_composition_test.py b/tst/unit/store/local_backends_composition_test.py index 0ee370c..202f442 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 @@ -20,6 +19,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 @@ -184,15 +184,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,14 +246,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 - - def wait(self, timeout=None): - return 0 +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", [ @@ -247,7 +266,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") @@ -277,7 +296,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) @@ -326,7 +344,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):