fix(artifact): bound docker save and isolate file upload attempts

src/agent_env/artifact/artifacts/docker_image.py

@@ -10,6 +10,7 @@
 import shlex
 import subprocess
 import tempfile
+import threading
 import uuid
 from dataclasses import dataclass
 from pathlib import Path
@@ -145,16 +146,7 @@ def put(
                 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}")
+                _stream_save_output(save_proc, gz_file, image_ref, 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"
@@ -615,6 +607,53 @@ def docker_login(auth) -> None:
         raise RuntimeError(f"docker login failed: {login.stderr}")
 
 
+# How much of a failed save's stderr its error message carries.
+_SAVE_STDERR_TAIL_BYTES = 65536
+
+
+def _stream_save_output(save_proc: subprocess.Popen, gz_file, image_ref: str, timeout_seconds: int) -> None:
+    """Stream a ``docker save`` process's stdout into ``gz_file`` until it exits, bounded by ``timeout_seconds``.
+
+    stderr drains on a thread into a bounded tail, and a watchdog kills the process at the
+    deadline: reading stdout to EOF first deadlocks once the process fills its stderr pipe, and
+    a stalled process blocks the read itself, so timing out only the wait after EOF never fires.
+    The process is reaped on every path.
+    """
+    stderr_tail = bytearray()
+
+    def drain_stderr() -> None:
+        while chunk := save_proc.stderr.read(8192):
+            stderr_tail.extend(chunk)
+            del stderr_tail[:-_SAVE_STDERR_TAIL_BYTES]
+
+    drainer = threading.Thread(target=drain_stderr, daemon=True)
+    drainer.start()
+    timed_out = threading.Event()
+
+    def watchdog() -> None:
+        timed_out.set()
+        save_proc.kill()
+
+    timer = threading.Timer(timeout_seconds, watchdog)
+    timer.start()
+    try:
+        while chunk := save_proc.stdout.read(8192):
+            gz_file.write(chunk)
+    except BaseException:
+        save_proc.kill()
+        raise
+    finally:
+        # The process is reaped on every path: it either exited, or the watchdog killed it
+        # (which is also what unblocks a read stuck on a stalled process's stdout).
+        save_proc.wait()
+        timer.cancel()
+        drainer.join(timeout=5)
+    if timed_out.is_set():
+        raise RuntimeError(f"docker save {image_ref} timed out after {timeout_seconds} seconds")
+    if save_proc.returncode != 0:
+        raise RuntimeError(f"docker save {image_ref} failed: {stderr_tail.decode(errors='replace')}")
+
+
 def _parse_copy_sources(dockerfile_text: str, dockerfile_rel_path: str | None = None) -> list[str]:
     """Parse COPY/ADD source paths from Dockerfile text. Returns top-level directories to include in context tar."""
     paths: list[str] = []

src/agent_env/artifact/artifacts/file.py

@@ -10,12 +10,18 @@
 
 from agent_env.artifact.artifact import Artifact, _write_twin
 from agent_env.config import get_config
+from agent_env.store.base import NotFoundError, ObjectAlreadyExistsError
 
 if TYPE_CHECKING:
+    from agent_env.artifact.store import ArtifactStore
     from agent_env.bundle.authoring import AuthoringContext
     from agent_env.entity.export import ExportContext
 
 
+# How many reserved versions one upload tries before giving up on the version path.
+_UPLOAD_VERSION_ATTEMPTS = 3
+
+
 class FileArtifact(Artifact):
     """A generic file artifact that stores any file in the object store.
 
@@ -69,19 +75,10 @@ def put(
         if content_type is None:
             content_type = "application/octet-stream"
 
-        version = store.next_version(id)
-
         # put_object_file, not put_object: a hosted store uploads a file in parts, where a
         # single put is capped (5 GB on S3). It streams 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,
-        )
+        stored_url = cls._upload(store, id, filename, content_type, file_path=file_path)
 
         # The object is keyed by the reserved version; the document takes the next free one, the same unless a
         # concurrent write took it first.
@@ -108,16 +105,7 @@ 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,
-        )
+        stored_url = cls._upload(store, id, filename, content_type, content=content)
 
         # The object is keyed by the reserved version; the document takes the next free one, the same unless a
         # concurrent write took it first.
@@ -130,6 +118,61 @@ def put_bytes(
         )
         return store.put_document(instance)
 
+    @classmethod
+    def _upload(
+        cls,
+        store: "ArtifactStore",
+        id: str,
+        filename: str,
+        content_type: str,
+        *,
+        file_path: str | None = None,
+        content: bytes | None = None,
+    ) -> str:
+        """Upload to the reserved version's write-once path, working around what is already there.
+
+        A collision with a published version means a concurrent writer won the reservation:
+        reserve the next version and retry. A collision with no document behind it is an
+        orphan an earlier attempt left when its document insert failed; write-once objects
+        are never deleted or overwritten, so this attempt uploads under a prefix of its own
+        (``ArtifactStore.attempt_prefix``) instead of staying blocked by the orphan.
+        """
+        version = store.next_version(id)
+        for _ in range(_UPLOAD_VERSION_ATTEMPTS):
+            try:
+                if file_path is not None:
+                    return store.put_object_file(
+                        artifact_type="file",
+                        id=id,
+                        version=version,
+                        object_name=filename,
+                        file_path=file_path,
+                        content_type=content_type,
+                    )
+                return store.put_object(
+                    artifact_type="file",
+                    id=id,
+                    version=version,
+                    object_name=filename,
+                    data=content,
+                    content_type=content_type,
+                )
+            except ObjectAlreadyExistsError:
+                try:
+                    store.get(id, version)
+                except NotFoundError:
+                    break
+                version = store.next_version(id)
+        else:
+            raise ObjectAlreadyExistsError(
+                f"every reserved version of artifact {id!r} was published by a concurrent writer"
+            )
+        prefix = store.attempt_prefix("file", id)
+        objects = get_config().get_object_store_to_write(prefix, id)
+        if file_path is not None:
+            return objects.put_file_at(prefix + filename, file_path, content_type)
+        return objects.put(objects.get_object_key(prefix + filename), content, content_type)
+
     @classmethod
     def from_toml(cls, data: dict, ctx: AuthoringContext) -> Self:
         """Write the folder's one file (``put_attempt``); ``description`` defaults to the file's name."""

tst/unit/artifact/test_docker_image_save.py

@@ -0,0 +1,80 @@
+"""DockerImageArtifact.put's ``docker save`` streaming: bounded, drained, and reaped.
+
+Regression guard: put used to read the save process's stdout to EOF before waiting
+on it, with stderr never drained. A save that filled its stderr pipe deadlocked, and
+a stalled save blocked the read itself, so the timeout in the wait never fired.
+The docker CLI is a PATH stub; the image store and the push are stubbed at the same
+boundaries the other docker image tests use. No Docker daemon is involved.
+"""
+
+import gzip
+import os
+import time
+from types import SimpleNamespace
+
+import pytest
+
+import agent_env.artifact.artifacts.docker_image as docker_image
+from agent_env.artifact.artifacts.docker_image import DockerImageArtifact
+from agent_env.config import get_config
+
+_STUB = """#!/bin/sh
+if [ "$1" = "save" ]; then
+  case "$FAKE_DOCKER_MODE" in
+    flood) python3 -c 'import sys; sys.stderr.write("x" * 300000); sys.stdout.write("FAKETAR")' ;;
+    stall) exec sleep 60 ;;
+    fail) echo "boom: no such image" >&2; exit 3 ;;
+    *) printf 'FAKETAR' ;;
+  esac
+  exit 0
+fi
+exit 0
+"""
+
+
+@pytest.fixture
+def stubbed_put(local_stores, tmp_path, monkeypatch):
+    bindir = tmp_path / "bin"
+    bindir.mkdir()
+    stub = bindir / "docker"
+    stub.write_text(_STUB)
+    stub.chmod(0o755)
+    monkeypatch.setenv("PATH", f"{bindir}:{os.environ['PATH']}")
+    monkeypatch.setattr(docker_image, "push_local_image", lambda *a, **k: None)
+    image_store = SimpleNamespace(
+        image_ref=lambda repo, tag: f"localhost:5000/{repo}:{tag}",
+        ensure_repository=lambda repo: None,
+    )
+    monkeypatch.setattr(get_config(), "get_image_store_for", lambda id: image_store)
+    monkeypatch.setattr(DockerImageArtifact, "DOCKER_SAVE_TIMEOUT_SECONDS", 2)
+    return monkeypatch
+
+
+def _put(monkeypatch, mode, artifact_id):
+    monkeypatch.setenv("FAKE_DOCKER_MODE", mode)
+    return DockerImageArtifact.put(artifact_id, description="d", image_name="img:latest")
+
+
+def test_save_with_saturated_stderr_completes(stubbed_put):
+    art = _put(stubbed_put, "flood", "img-flood")
+
+    assert gzip.decompress(art.load()) == b"FAKETAR"
+
+
+def test_stalled_save_is_killed_at_the_timeout(stubbed_put):
+    start = time.monotonic()
+    with pytest.raises(RuntimeError, match="timed out after 2 seconds"):
+        _put(stubbed_put, "stall", "img-stall")
+    assert time.monotonic() - start < 20
+
+
+def test_failed_save_reports_its_stderr(stubbed_put):
+    with pytest.raises(RuntimeError, match="boom: no such image"):
+        _put(stubbed_put, "fail", "img-fail")
+
+
+def test_save_streams_the_archive_to_the_object_store(stubbed_put):
+    art = _put(stubbed_put, "ok", "img-ok")
+
+    assert art.version == 1
+    assert gzip.decompress(art.load()) == b"FAKETAR"