feat(store)!: a version named on a put must be the next one

src/agent_env/artifact/artifact.py

@@ -34,7 +34,7 @@ class Artifact(BaseModel):
     """Base class for all environment artifacts. Immutable - each modification creates a new version."""
 
     id: str = Field(description="Unique artifact identifier")
-    version: int = Field(default=0, description="Artifact version (auto-calculated on put)")
+    version: int = Field(default=0, description="Artifact version: 0 takes the next free one on put; another must be the next")
     type: str = Field(description="The type of Artifact")
 
     toml_refs: ClassVar[tuple[EntityRef, ...]] = ()

src/agent_env/artifact/artifacts/cli.py

@@ -80,15 +80,13 @@ def put(
             raise ValueError(f"entrypoint {entrypoint!r} isn't among the CLI's files")
 
         store = get_artifact_store()
-        version = store.next_version(id)
         cli_object_url = store.attempt_prefix(cls.model_fields["type"].default, id)
 
         universe = FileArtifactUniverse.put_bundled(id=derive_id(id, "files"), files=dict(files),
                                                     prefix_url=cli_object_url)
 
         instance = cls(
             id=id,
-            version=version,
             cli_files_id=universe.id,
             cli_object_url=cli_object_url,
             entrypoint=entrypoint,

src/agent_env/artifact/artifacts/docker_image.py

@@ -197,10 +197,8 @@ def put_tar(
             if url:
                 get_config().check_object_url(id, url)
         store = get_artifact_store()
-        version = store.next_version(id)
         instance = cls(
             id=id,
-            version=version,
             description=description,
             image_name=image_name,
             tar_gz_object_url=tar_gz_object_url,
@@ -236,11 +234,9 @@ def put_ref(cls, id: str, *, description: str, image_name: str, earlier_pin: str
         if context_path is not None:
             context = _context_of(Path(context_path), dockerfile_path)
             store = get_artifact_store()
-            version = store.next_version(id)
             prefix = store.attempt_prefix("docker_image", id)
             return store.put_document(cls(
                 id=id,
-                version=version,
                 description=description,
                 image_name=image_name,
                 image_digest=image_digest,
@@ -253,8 +249,7 @@ def put_ref(cls, id: str, *, description: str, image_name: str, earlier_pin: str
         else:
             pinned = pin_digest(image_name, config.get_image_store().auth(image_name))
         store = get_artifact_store()
-        version = store.next_version(id)
-        return store.put_document(cls(id=id, version=version, description=description, image_name=pinned))
+        return store.put_document(cls(id=id, description=description, image_name=pinned))
 
     @classmethod
     def put_context(cls, id: str, *, description: str, context_path: str, dockerfile_path: str,
@@ -270,13 +265,11 @@ def put_context(cls, id: str, *, description: str, context_path: str, dockerfile
         platform = normalize_platform(platform) if platform else None
         context = _context_of(Path(context_path), dockerfile_path)
         store = get_artifact_store()
-        version = store.next_version(id)
         prefix = store.attempt_prefix("docker_image", id)
         build_context_object_url = _upload_context(get_config().get_object_store_to_write(prefix, id), prefix, context)
         source_digest = context.source_digest(platform)
         return store.put_document(cls(
             id=id,
-            version=version,
             description=description,
             image_name=_context_image_name(id, source_digest),
             build_context_object_url=build_context_object_url,

src/agent_env/artifact/artifacts/file.py

@@ -83,10 +83,10 @@ def put(
             content_type=content_type,
         )
 
-        # Store artifact document in MongoDB
+        # The object is keyed by the reserved version; the document takes the next free one, the same unless a
+        # concurrent write took it first.
         instance = cls(
             id=id,
-            version=version,
             description=description,
             filename=filename,
             content_type=content_type,
@@ -119,9 +119,10 @@ def put_bytes(
             content_type=content_type,
         )
 
+        # The object is keyed by the reserved version; the document takes the next free one, the same unless a
+        # concurrent write took it first.
         instance = cls(
             id=id,
-            version=version,
             description=description,
             filename=filename,
             content_type=content_type,
@@ -199,11 +200,10 @@ def put_at(
             content_type = "application/octet-stream"
 
         artifact_store = get_artifact_store()
-        version = artifact_store.next_version(id)
+        artifact_store.check_writable(id)
         stored_url = store.put_file_at(object_url, file_path, content_type)
         instance = cls(
             id=id,
-            version=version,
             description=description,
             filename=filename,
             content_type=content_type,
@@ -240,10 +240,8 @@ def put_existing(
         content_type = metadata.content_type or mimetypes.guess_type(filename)[0] or "application/octet-stream"
 
         artifact_store = get_artifact_store()
-        version = artifact_store.next_version(id)
         instance = cls(
             id=id,
-            version=version,
             description=description,
             filename=filename,
             content_type=content_type,

src/agent_env/artifact/artifacts/file_artifact_universe.py

@@ -91,10 +91,8 @@ def put(
             get_config().check_object_url(id, bundle_object_url)
 
         store = get_artifact_store()
-        version = store.next_version(id)
         instance = cls(
             id=id,
-            version=version,
             file_artifact_refs={
                 fname: fa.as_ref() for fname, fa in file_artifacts.items()
             },

src/agent_env/artifact/artifacts/skill.py

@@ -96,7 +96,6 @@ def put(cls, id: str, *, skill_dir: Optional[Path] = None,
         frontmatter, _body = parse_skill_md(skill_md_bytes)
 
         store = get_artifact_store()
-        version = store.next_version(id)
         skill_object_url = store.attempt_prefix(cls.model_fields["type"].default, id)
 
         universe = FileArtifactUniverse.put_bundled(
@@ -107,7 +106,6 @@ def put(cls, id: str, *, skill_dir: Optional[Path] = None,
 
         instance = cls(
             id=id,
-            version=version,
             skill_files_id=universe.id,
             skill_object_url=skill_object_url,
             agent_skills_spec_version=AGENT_SKILLS_SPEC_VERSION,

src/agent_env/artifact/store.py

@@ -118,12 +118,15 @@ def get(self, id: str, version: Optional[int] = None) -> "Artifact":
             raise NotFoundError(f"Artifact {id}{version_str} not found")
         return artifact
 
-    def next_version(self, id: str) -> int:
-        # Every put asks for its version before writing an object or pushing an image, so an id
-        # the store refuses, or a bare one while an @local task runs, is refused here rather than
-        # after those side effects, in put_document.
+    def check_writable(self, id: str) -> None:
+        """Refuse an id the store won't take, or a bare one while an @local task runs. A put calls this, or
+        ``next_version`` or ``attempt_prefix``, which do, before it writes an object or pushes an image, so the
+        refusal comes before those side effects rather than after them in put_document."""
         self._doc_store.check_id(id)
         get_config().check_local_run_write(id)
+
+    def next_version(self, id: str) -> int:
+        self.check_writable(id)
         return self._versioned.next_version(id)
 
     def attempt_prefix(self, artifact_type: str, id: str) -> str:
@@ -173,11 +176,10 @@ def get_object(self, object_url: str) -> bytes:
         return get_config().get_object_store_at(object_url).get(object_url)
 
     def put_document(self, artifact: _A) -> _A:
-        """Insert an artifact document, retrying the version-allocation race.
-
-        Concurrent writers to one id can allocate the same next_version; the
-        unique index rejects the losers, so we re-allocate and retry rather
-        than fail the caller (which drops a pass@k rollout).
+        """Insert an artifact document. One that names a version is stored as exactly that one, which must be the
+        id's next (ConcurrentModificationError otherwise). One at version 0 takes the next free version, and
+        re-allocates when a concurrent writer takes it first rather than fail the caller (which drops a pass@k
+        rollout).
         """
 
         self._doc_store.check_id(artifact.id)
@@ -186,9 +188,12 @@ def put_document(self, artifact: _A) -> _A:
 
         store = self._doc_store
         versioned = self._versioned_cache.for_store(store)
-        version = artifact.version if artifact.version > 0 else versioned.next_version(artifact.id)
         doc = artifact.model_dump(by_alias=True)
         doc["created_at_utc"] = datetime.now(timezone.utc)
+        if artifact.version > 0:
+            versioned.insert_next(doc, artifact.version)
+            return artifact.model_copy()
+        version = versioned.next_version(artifact.id)
         last_exc: Optional[DuplicateKeyError] = None
         for attempt in range(_MAX_VERSION_RETRIES):
             doc["version"] = version

src/agent_env/store/document_store/document_store.py

@@ -8,6 +8,7 @@
 from dataclasses import dataclass, field
 from typing import Any, Callable, Generic, Optional, TypeVar, Union
 
+from agent_env.store.base import ConcurrentModificationError
 from agent_env.store.ids import is_local_id
 
 logger = logging.getLogger(__name__)
@@ -446,8 +447,8 @@ def get(self, id: str, version: Optional[int] = None) -> Optional[T]:
     def next_version(self, id: str) -> int:
         """Allocate the version a write to ``id`` would land on.
 
-        The id is checked here and not only at ``put`` because every artifact
-        helper calls this first and then writes remote data — an object to the object store, an image to a
+        The id is checked here and not only at ``put`` because an artifact helper that reserves its version
+        calls this first and then writes remote data — an object to the object store, an image to a
         registry — before it has a document to store. Refusing the id at the end would leave
         that data orphaned with no artifact record pointing at it.
         """
@@ -456,13 +457,17 @@ def next_version(self, id: str) -> int:
         return (doc["version"] + 1) if doc is not None else 1
 
     def put(self, entity: T, max_retries: int = 5) -> int:
-        """Serialize and insert the entity under the next free version; return that version."""
+        """Serialize and insert the entity; return its version. An entity that names a version is stored as exactly
+        that one, which ``insert_next`` requires to be the id's next; one that doesn't takes the next free version."""
         if self._serialize is None:
             raise RuntimeError("VersionedEntityStore was built without a serializer; put() is unsupported")
         doc = self._serialize(entity)
         if not isinstance(doc.get("id"), str):
             raise ValueError(f"an entity in {self._collection!r} needs a string id, got {doc.get('id')!r}")
         self._doc_store.check_id(doc["id"])
+        if doc.get("version"):
+            self.insert_next(doc, doc["version"])
+            return doc["version"]
         last_error: Optional[DuplicateKeyError] = None
         for _ in range(max_retries):
             version = self.next_version(doc["id"])
@@ -476,6 +481,23 @@ def put(self, entity: T, max_retries: int = 5) -> int:
             f"Could not assign a version for id={doc.get('id')!r} after {max_retries} retries"
         ) from last_error
 
+    def insert_next(self, doc: dict, version: int) -> None:
+        """Insert ``doc`` as ``version``, which must be its id's next version: ConcurrentModificationError when it
+        isn't, or when another write takes it first, so a write never lands on or skips past one it didn't see."""
+        latest = self._doc_store.latest_version(self._collection, doc["id"])
+        current = latest["version"] if latest is not None else 0
+        if version != current + 1:
+            raise ConcurrentModificationError(
+                f"{doc['id']!r} in {self._collection!r} is at version {current}, so its next version is "
+                f"{current + 1}, not {version}"
+            )
+        try:
+            self._doc_store.insert(self._collection, {**doc, "version": version})
+        except DuplicateKeyError as e:
+            raise ConcurrentModificationError(
+                f"another write took version {version} of {doc['id']!r} in {self._collection!r} first"
+            ) from e
+
     def query(
         self,
         filter: Filter,

src/agent_env/task_step/task_steps/collect_artifacts.py

@@ -966,18 +966,15 @@ def _register_file_artifact(self, store, file_artifacts, artifact_id, key, objec
         from agent_env.artifact import FileArtifact
 
         fa_id = derive_id(artifact_id, hashlib.sha256(object_name.encode("utf-8")).hexdigest()[:16])
-        fa_version = store.next_version(fa_id)
         fa = FileArtifact(
             id=fa_id,
-            version=fa_version,
             description=f"Collected artifact '{object_name}' from task instance {context.instance_id or artifact_id}",
             filename=object_name.split("/")[-1],
             content_type=content_type,
             object_url=object_url,
         )
         try:
-            store.put_document(fa)
-            file_artifacts[key] = fa
+            file_artifacts[key] = store.put_document(fa)
         except Exception as doc_err:
             # Don't fail the whole collection if the document write fails.
             logger.warning(f"Failed to register FileArtifact for {key}: {doc_err}")

src/agent_env/task_step/task_steps/snapshot_env.py

@@ -652,7 +652,6 @@ def _register_uploaded_bundle():
                             get_config().check_object_url(fa_id, object_url)
                             fa = FileArtifact(
                                 id=fa_id,
-                                version=store.next_version(fa_id),
                                 description=f"State snapshot of {environment_name} from env {env.id}",
                                 filename=filename,
                                 content_type=content_type,

tst/store/conformance.py

@@ -14,6 +14,7 @@
 
 from agent_env.store import (
     AbsentOrNull,
+    ConcurrentModificationError,
     DuplicateKeyError,
     Eq,
     Filter,
@@ -266,6 +267,20 @@ def versioned_entity_store_roundtrip(store, coll):
     assert ves.count(Filter.of(id="e")) == 2
 
 
+def versioned_entity_store_put_at_a_version(store, coll):
+    """A put naming a version lands on exactly it, which must be the id's next; any other writes nothing."""
+    ves: VersionedEntityStore[dict] = VersionedEntityStore(
+        store, coll, serialize=dict, deserialize=lambda d: d
+    )
+    assert ves.put({"id": "e", "version": 1, "n": 1}) == 1
+    assert ves.put({"id": "e", "version": 2, "n": 2}) == 2
+    for taken_latest_or_past_next in (1, 2, 4):
+        with pytest.raises(ConcurrentModificationError):
+            ves.put({"id": "e", "version": taken_latest_or_past_next, "n": 0})
+    assert ves.put({"id": "e", "n": 3}) == 3
+    assert ves.get("e")["n"] == 3 and ves.count(Filter.of(id="e")) == 3
+
+
 def latest_per_id_reduces_to_newest_version(store, coll):
     """One row per id, the highest-version one; filter applies before grouping."""
     for doc in [
@@ -440,6 +455,24 @@ def racing_pre_image_updates_have_exactly_one_winner(store, coll):
     assert store.find_one(coll, Filter.of(id="x")) == {"id": "x", "owner": winners[0]}
 
 
+def racing_puts_at_one_version_have_exactly_one_winner(store, coll):
+    """Writers that all name the id's next version race for it; one lands and the rest are refused."""
+    ves: VersionedEntityStore[dict] = VersionedEntityStore(
+        store, coll, serialize=dict, deserialize=lambda d: d
+    )
+    ves.put({"id": "e"})
+
+    def put(racer):
+        try:
+            return ves.put({"id": "e", "version": 2, "owner": racer})
+        except ConcurrentModificationError:
+            return None
+
+    landed = _race(put)
+    assert landed.count(2) == 1 and landed.count(None) == _RACERS - 1
+    assert ves.get("e")["owner"] == landed.index(2) and ves.count(Filter.of(id="e")) == 2
+
+
 def concurrent_compare_and_swap_loses_no_update(store, coll):
     """Each writer's mutation depends on what it read, so a write applied under a stale guard
     drops another writer's step."""
@@ -483,6 +516,7 @@ def update_one_and_get_upsert_inserts_on_a_miss(store, coll):
     dotted_array_traversal,
     positional_array_index_filter_and_set,
     versioned_entity_store_roundtrip,
+    versioned_entity_store_put_at_a_version,
     latest_per_id_reduces_to_newest_version,
     latest_per_id_sort_offset_limit_apply_after_grouping,
     latest_per_id_page_returns_entity_total_and_window,
@@ -496,5 +530,6 @@ def update_one_and_get_upsert_inserts_on_a_miss(store, coll):
     racing_guarded_replaces_have_exactly_one_winner,
     racing_update_one_and_get_returns_only_its_own_write,
     racing_pre_image_updates_have_exactly_one_winner,
+    racing_puts_at_one_version_have_exactly_one_winner,
     concurrent_compare_and_swap_loses_no_update,
 ]

tst/unit/a2a_agent/put_version_test.py

@@ -0,0 +1,31 @@
+"""A2AAgent.put with a version: it lands on exactly that version, which must be the agent's next."""
+
+import pytest
+
+from agent_env.a2a_agent import A2AAgent
+from agent_env.artifact.artifacts.docker_image import DockerImageArtifact
+from agent_env.artifact.store import get_artifact_store
+from agent_env.store.base import ConcurrentModificationError
+
+
+@pytest.fixture
+def image(local_stores):
+    return get_artifact_store().put_document(DockerImageArtifact(
+        id="agent__agent_image", description="d", image_name="ghcr.io/acme/agent:1", tar_gz_s3_url="s3://b/x"))
+
+
+def test_a_put_naming_the_next_version_lands_on_it(image):
+    assert A2AAgent.put(id="agent", docker_image_artifact=image, version=1).version == 1
+    assert A2AAgent.put(id="agent", docker_image_artifact=image, version=2).version == 2
+
+
+@pytest.mark.parametrize("version", [2, 5])
+def test_a_put_naming_another_version_writes_nothing(image, version):
+    A2AAgent.put(id="agent", docker_image_artifact=image)
+    A2AAgent.put(id="agent", docker_image_artifact=image, metadata={"by": "other"})
+
+    with pytest.raises(ConcurrentModificationError, match=f"at version 2, so its next version is 3, not {version}"):
+        A2AAgent.put(id="agent", docker_image_artifact=image, version=version, metadata={"by": "stale"})
+
+    latest = A2AAgent.get("agent")
+    assert latest.version == 2 and latest.metadata["by"] == "other"